This is an automated email from the ASF dual-hosted git repository.
siddteotia pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new e6d3dcdbda [multistage] Fix leaf stage return (#9892)
e6d3dcdbda is described below
commit e6d3dcdbdaf9e08f0ffd939e55361bb946b03f93
Author: Rong Rong <[email protected]>
AuthorDate: Mon Dec 5 15:51:07 2022 -0800
[multistage] Fix leaf stage return (#9892)
* initial commit to move schema logic into LeafStageTransfer
only work with selection result block. others to be add later
* refactor
* adding distinct handling
* also support agg and group by
* add tests
* adding more tests
Co-authored-by: Rong Rong <[email protected]>
---
.../apache/pinot/query/runtime/QueryRunner.java | 46 +---
.../LeafStageTransferableBlockOperator.java | 234 +++++++++++++++++++--
.../LeafStageTransferableBlockOperatorTest.java | 78 +++++--
.../src/test/resources/queries/Aggregates.json | 93 +++++---
.../src/test/resources/queries/Distincts.json | 47 +++++
.../test/resources/queries/LexicalStructure.json | 21 +-
.../test/resources/queries/SelectExpressions.json | 15 +-
7 files changed, 414 insertions(+), 120 deletions(-)
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
index 7b69aa40ce..5cf1cfa999 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
@@ -20,10 +20,8 @@ package org.apache.pinot.query.runtime;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
-import java.util.LinkedList;
import java.util.List;
import java.util.Map;
-import java.util.Objects;
import java.util.concurrent.ExecutorService;
import org.apache.helix.HelixManager;
import org.apache.helix.store.zk.ZkHelixPropertyStore;
@@ -31,13 +29,10 @@ import org.apache.helix.zookeeper.datamodel.ZNRecord;
import org.apache.pinot.common.exception.QueryException;
import org.apache.pinot.common.metadata.ZKMetadataProvider;
import org.apache.pinot.common.metrics.ServerMetrics;
-import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.data.manager.InstanceDataManager;
import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
-import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
import org.apache.pinot.core.query.executor.ServerQueryExecutorV1Impl;
import org.apache.pinot.core.query.request.ServerQueryRequest;
-import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
import org.apache.pinot.core.transport.ServerInstance;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.mailbox.MultiplexingMailboxService;
@@ -189,47 +184,10 @@ public class QueryRunner {
private InstanceResponseBlock processServerQuery(ServerQueryRequest
serverQueryRequest,
ExecutorService executorService) {
try {
- InstanceResponseBlock result =
_serverExecutor.execute(serverQueryRequest, executorService);
-
- if (result.getRows() != null &&
serverQueryRequest.getQueryContext().getOrderByExpressions() != null) {
- // we only re-arrange columns to match the projection in the case of
order by - this is to ensure
- // that V1 results match what the expected projection schema in the
calcite logical operator; if
- // we realize that there are other situations where we need to
post-process v1 results to adhere to
- // the expected results we should factor this out and also apply the
canonicalization of the data
- // types during this post-process step (also see
LeafStageTransferableBlockOperator#canonicalizeRow)
- DataSchema dataSchema = result.getDataSchema();
- List<String> selectionColumns =
-
SelectionOperatorUtils.getSelectionColumns(serverQueryRequest.getQueryContext(),
dataSchema);
-
- int[] columnIndices =
SelectionOperatorUtils.getColumnIndices(selectionColumns, dataSchema);
- int numColumns = columnIndices.length;
-
- DataSchema resultDataSchema =
SelectionOperatorUtils.getSchemaForProjection(dataSchema, columnIndices);
-
- // Extract the result rows
- LinkedList<Object[]> rowsInSelectionResults = new LinkedList<>();
- for (Object[] row : result.getRows()) {
- assert row != null;
- Object[] extractedRow = new Object[numColumns];
- for (int i = 0; i < numColumns; i++) {
- Object value = row[columnIndices[i]];
- if (value != null) {
- extractedRow[i] = value;
- }
- }
-
- rowsInSelectionResults.addFirst(extractedRow);
- }
-
- return new InstanceResponseBlock(
- new SelectionResultsBlock(resultDataSchema,
rowsInSelectionResults),
- serverQueryRequest.getQueryContext());
- } else {
- return result;
- }
+ return _serverExecutor.execute(serverQueryRequest, executorService);
} catch (Exception e) {
InstanceResponseBlock errorResponse = new InstanceResponseBlock();
-
errorResponse.getExceptions().put(QueryException.QUERY_EXECUTION_ERROR_CODE,
Objects.toString(e.getMessage()));
+
errorResponse.getExceptions().put(QueryException.QUERY_EXECUTION_ERROR_CODE,
e.getMessage());
return errorResponse;
}
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java
index 551955c388..8def59dd32 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java
@@ -18,11 +18,14 @@
*/
package org.apache.pinot.query.runtime.operator;
+import com.clearspring.analytics.util.Preconditions;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.PriorityQueue;
+import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.datablock.DataBlockUtils;
@@ -32,7 +35,13 @@ import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.operator.BaseOperator;
import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
+import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock;
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.DistinctResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import
org.apache.pinot.core.query.aggregation.function.DistinctAggregationFunction;
+import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
@@ -54,10 +63,12 @@ public class LeafStageTransferableBlockOperator extends
BaseOperator<Transferabl
private final InstanceResponseBlock _errorBlock;
private final List<InstanceResponseBlock> _baseResultBlock;
+ private final DataSchema _desiredDataSchema;
private int _currentIndex;
public LeafStageTransferableBlockOperator(List<InstanceResponseBlock>
baseResultBlock, DataSchema dataSchema) {
_baseResultBlock = baseResultBlock;
+ _desiredDataSchema = dataSchema;
_errorBlock = baseResultBlock.stream().filter(e ->
!e.getExceptions().isEmpty()).findFirst().orElse(null);
_currentIndex = 0;
}
@@ -84,13 +95,10 @@ public class LeafStageTransferableBlockOperator extends
BaseOperator<Transferabl
} else {
if (_currentIndex < _baseResultBlock.size()) {
InstanceResponseBlock responseBlock =
_baseResultBlock.get(_currentIndex++);
- BaseResultsBlock resultsBlock = responseBlock.getResultsBlock();
- if (resultsBlock != null) {
- List<Object[]> rows =
- toList(resultsBlock.getRows(responseBlock.getQueryContext()),
responseBlock.getDataSchema());
- return new TransferableBlock(rows, responseBlock.getDataSchema(),
DataBlock.Type.ROW);
+ if (responseBlock.getResultsBlock() != null) {
+ return composeTransferableBlock(responseBlock, _desiredDataSchema);
} else {
- return new TransferableBlock(Collections.emptyList(),
responseBlock.getDataSchema(), DataBlock.Type.ROW);
+ return new TransferableBlock(Collections.emptyList(),
_desiredDataSchema, DataBlock.Type.ROW);
}
} else {
_currentIndex = -1;
@@ -99,6 +107,181 @@ public class LeafStageTransferableBlockOperator extends
BaseOperator<Transferabl
}
}
+ /**
+ * this is data transfer block compose method is here to ensure that V1
results match what the expected projection
+ * schema in the calcite logical operator.
+ * <p> it applies different clean up mechanism based on different type of
{@link BaseResultsBlock} and the
+ * {@link org.apache.pinot.core.query.request.context.QueryContext}.</p>
+ * <p> this also applies to the canonicalization of the data types during
post post-process step.</p>
+ *
+ * @param responseBlock result block from leaf stage
+ * @param desiredDataSchema the desired schema for send operator
+ * @return the converted {@link TransferableBlock} that conform with the
desiredDataSchema
+ */
+ private static TransferableBlock
composeTransferableBlock(InstanceResponseBlock responseBlock,
+ DataSchema desiredDataSchema) {
+ BaseResultsBlock resultsBlock = responseBlock.getResultsBlock();
+ if (resultsBlock instanceof SelectionResultsBlock) {
+ return composeSelectTransferableBlock(responseBlock, desiredDataSchema);
+ } else if (resultsBlock instanceof AggregationResultsBlock) {
+ return composeAggregationTransferableBlock(responseBlock,
desiredDataSchema);
+ } else if (resultsBlock instanceof GroupByResultsBlock) {
+ return composeGroupByTransferableBlock(responseBlock, desiredDataSchema);
+ } else if (resultsBlock instanceof DistinctResultsBlock) {
+ return composeDistinctTransferableBlock(responseBlock,
desiredDataSchema);
+ } else {
+ throw new IllegalArgumentException("Unsupported result block type: " +
resultsBlock);
+ }
+ }
+
+ /**
+ * we only need to rearrange columns when distinct is not conforming with
selection columns, specifically:
+ * <ul>
+ * <li> when distinct is not returning final result:
+ * it should never happen as non-final result contains Object opaque
columns v2 engine can't process.</li>
+ * <li> when distinct columns are not all being selected:
+ * it should never happen as leaf stage MUST return the entire
list.</li>
+ * </ul>
+ *
+ * @see
LeafStageTransferableBlockOperator#composeTransferableBlock(InstanceResponseBlock,
DataSchema).
+ */
+ @SuppressWarnings("ConstantConditions")
+ private static TransferableBlock
composeDistinctTransferableBlock(InstanceResponseBlock responseBlock,
+ DataSchema desiredDataSchema) {
+ List<String> selectionColumns = Arrays.asList(
+ ((DistinctAggregationFunction)
responseBlock.getQueryContext().getAggregationFunctions()[0]).getColumns());
+ int[] columnIndices =
SelectionOperatorUtils.getColumnIndices(selectionColumns, desiredDataSchema);
+ Preconditions.checkState(inOrder(columnIndices), "Incompatible distinct
table schema for leaf stage."
+ + "Expected: " + desiredDataSchema + ". Actual: " + desiredDataSchema);
+ return composeDirectTransferableBlock(responseBlock, desiredDataSchema);
+ }
+
+ /**
+ * Calcite generated {@link DataSchema} should conform with Pinot's group by
result schema thus we only need to check
+ * for correctness similar to distinct case.
+ *
+ * @see
LeafStageTransferableBlockOperator#composeDirectTransferableBlock(InstanceResponseBlock,
DataSchema).
+ * @see
LeafStageTransferableBlockOperator#composeTransferableBlock(InstanceResponseBlock,
DataSchema).
+ */
+ @SuppressWarnings("ConstantConditions")
+ private static TransferableBlock
composeGroupByTransferableBlock(InstanceResponseBlock responseBlock,
+ DataSchema desiredDataSchema) {
+ DataSchema resultSchema = responseBlock.getDataSchema();
+ // GROUP-BY column names conforms with selection expression
+ List<String> selectionColumns =
responseBlock.getQueryContext().getSelectExpressions().stream()
+ .map(e -> e.toString()).collect(Collectors.toList());
+ int[] columnIndices =
SelectionOperatorUtils.getColumnIndices(selectionColumns, resultSchema);
+ Preconditions.checkState(inOrder(columnIndices), "Incompatible group by
result schema for leaf stage."
+ + "Expected: " + desiredDataSchema + ". Actual: " + resultSchema);
+ return composeDirectTransferableBlock(responseBlock, desiredDataSchema);
+ }
+
+
+ /**
+ * Calcite generated {@link DataSchema} should conform with Pinot's agg
result schema thus we only need to check
+ * for correctness similar to distinct case.
+ *
+ * @see
LeafStageTransferableBlockOperator#composeDirectTransferableBlock(InstanceResponseBlock,
DataSchema).
+ * @see
LeafStageTransferableBlockOperator#composeTransferableBlock(InstanceResponseBlock,
DataSchema).
+ */
+ @SuppressWarnings("ConstantConditions")
+ private static TransferableBlock
composeAggregationTransferableBlock(InstanceResponseBlock responseBlock,
+ DataSchema desiredDataSchema) {
+ DataSchema resultSchema = responseBlock.getDataSchema();
+ // AGG-ONLY column names are derived from AggFunction.getColumnName()
+ List<String> selectionColumns =
Arrays.stream(responseBlock.getQueryContext().getAggregationFunctions()).map(
+ a -> a.getColumnName()).collect(Collectors.toList());
+ int[] columnIndices =
SelectionOperatorUtils.getColumnIndices(selectionColumns, resultSchema);
+ Preconditions.checkState(inOrder(columnIndices), "Incompatible aggregate
result schema for leaf stage."
+ + "Expected: " + desiredDataSchema + ". Actual: " + resultSchema);
+ return composeDirectTransferableBlock(responseBlock, desiredDataSchema);
+ }
+
+ /**
+ * Only re-arrange columns to match the projection in the case of select /
order-by, when the desiredDataSchema
+ * doesn't conform with the result block schema exactly.
+ *
+ * @see
LeafStageTransferableBlockOperator#composeTransferableBlock(InstanceResponseBlock,
DataSchema).
+ */
+ @SuppressWarnings("ConstantConditions")
+ private static TransferableBlock
composeSelectTransferableBlock(InstanceResponseBlock responseBlock,
+ DataSchema desiredDataSchema) {
+ DataSchema resultSchema = responseBlock.getDataSchema();
+ List<String> selectionColumns =
SelectionOperatorUtils.getSelectionColumns(responseBlock.getQueryContext(),
+ resultSchema);
+ int[] columnIndices =
SelectionOperatorUtils.getColumnIndices(selectionColumns, resultSchema);
+ if (!inOrder(columnIndices)) {
+ DataSchema adjustedResultSchema =
SelectionOperatorUtils.getSchemaForProjection(resultSchema, columnIndices);
+
Preconditions.checkState(isDataSchemaColumnTypesCompatible(desiredDataSchema.getColumnDataTypes(),
+ adjustedResultSchema.getColumnDataTypes()), "Incompatible result
data schema: "
+ + "Expecting: " + desiredDataSchema + " Actual: " +
adjustedResultSchema);
+ // Extract the result rows
+ Collection<Object[]> resultRows = responseBlock.getRows();
+ List<Object[]> extractedRows = new ArrayList<>(resultRows.size());
+ return composeColumnIndexedTransferableBlock(responseBlock,
adjustedResultSchema, columnIndices);
+ } else {
+ return composeDirectTransferableBlock(responseBlock, desiredDataSchema);
+ }
+ }
+
+ /**
+ * Created {@link TransferableBlock} using column indices.
+ *
+ * @see
LeafStageTransferableBlockOperator#composeTransferableBlock(InstanceResponseBlock,
DataSchema).
+ */
+ private static TransferableBlock
composeColumnIndexedTransferableBlock(InstanceResponseBlock responseBlock,
+ DataSchema desiredDataSchema, int[] columnIndices) {
+ Collection<Object[]> resultRows = responseBlock.getRows();
+ List<Object[]> extractedRows = new ArrayList<>(resultRows.size());
+ if (responseBlock.getQueryContext().getOrderByExpressions() != null) {
+ // extract result row in ordered fashion
+ PriorityQueue<Object[]> priorityQueue = (PriorityQueue<Object[]>)
resultRows;
+ while (!priorityQueue.isEmpty()) {
+ extractedRows.add(canonicalizeRow(priorityQueue.poll(),
desiredDataSchema, columnIndices));
+ }
+ } else {
+ // extract result row in non-ordered fashion
+ for (Object[] row : resultRows) {
+ extractedRows.add(canonicalizeRow(row, desiredDataSchema,
columnIndices));
+ }
+ }
+ return new TransferableBlock(extractedRows, desiredDataSchema,
DataBlock.Type.ROW);
+ }
+
+ /**
+ * Fallback mechanism for {@link TransferableBlock}, used when no special
handling is necessary. This method only
+ * performs {@link DataSchema.ColumnDataType} canonicalization.
+ *
+ * @see
LeafStageTransferableBlockOperator#composeTransferableBlock(InstanceResponseBlock,
DataSchema).
+ */
+ private static TransferableBlock
composeDirectTransferableBlock(InstanceResponseBlock responseBlock,
+ DataSchema desiredDataSchema) {
+ Collection<Object[]> resultRows = responseBlock.getRows();
+ List<Object[]> extractedRows = new ArrayList<>(resultRows.size());
+ if (resultRows instanceof List) {
+ for (Object[] orgRow : resultRows) {
+ extractedRows.add(canonicalizeRow(orgRow, desiredDataSchema));
+ }
+ } else if (resultRows instanceof PriorityQueue) {
+ PriorityQueue<Object[]> priorityQueue = (PriorityQueue<Object[]>)
resultRows;
+ while (!priorityQueue.isEmpty()) {
+ extractedRows.add(canonicalizeRow(priorityQueue.poll(),
desiredDataSchema));
+ }
+ } else {
+ throw new UnsupportedOperationException("Unsupported collection type: "
+ resultRows.getClass());
+ }
+ return new TransferableBlock(extractedRows, desiredDataSchema,
DataBlock.Type.ROW);
+ }
+
+ private static boolean inOrder(int[] columnIndices) {
+ for (int i = 0; i < columnIndices.length; i++) {
+ if (columnIndices[i] != i) {
+ return false;
+ }
+ }
+ return true;
+ }
+
/**
* This util is used to canonicalize row generated from V1 engine, which is
stored using
* {@link DataSchema#getStoredColumnDataTypes()} format. However, the
transferable block ser/de stores data in the
@@ -111,28 +294,35 @@ public class LeafStageTransferableBlockOperator extends
BaseOperator<Transferabl
private static Object[] canonicalizeRow(Object[] row, DataSchema dataSchema)
{
Object[] resultRow = new Object[row.length];
for (int colId = 0; colId < row.length; colId++) {
- resultRow[colId] =
dataSchema.getColumnDataType(colId).convert(row[colId]);
+ Object value = row[colId];
+ if (value != null) {
+ resultRow[colId] = dataSchema.getColumnDataType(colId).convert(value);
+ }
}
return resultRow;
}
- private static List<Object[]> toList(Collection<Object[]> collection,
DataSchema dataSchema) {
- if (collection == null || collection.isEmpty()) {
- return new ArrayList<>();
- }
- List<Object[]> resultRows = new ArrayList<>(collection.size());
- if (collection instanceof List) {
- for (Object[] orgRow : collection) {
- resultRows.add(canonicalizeRow(orgRow, dataSchema));
+ private static Object[] canonicalizeRow(Object[] row, DataSchema dataSchema,
int[] columnIndices) {
+ Object[] resultRow = new Object[columnIndices.length];
+ for (int colId = 0; colId < columnIndices.length; colId++) {
+ Object value = row[columnIndices[colId]];
+ if (value != null) {
+ resultRow[colId] = dataSchema.getColumnDataType(colId).convert(value);
}
- } else if (collection instanceof PriorityQueue) {
- PriorityQueue<Object[]> priorityQueue = (PriorityQueue<Object[]>)
collection;
- while (!priorityQueue.isEmpty()) {
- resultRows.add(canonicalizeRow(priorityQueue.poll(), dataSchema));
+ }
+ return resultRow;
+ }
+
+ private static boolean
isDataSchemaColumnTypesCompatible(DataSchema.ColumnDataType[] desiredTypes,
+ DataSchema.ColumnDataType[] givenTypes) {
+ if (desiredTypes.length != givenTypes.length) {
+ return false;
+ }
+ for (int i = 0; i < desiredTypes.length; i++) {
+ if (desiredTypes[i] != givenTypes[i] &&
!givenTypes[i].isSuperTypeOf(desiredTypes[i])) {
+ return false;
}
- } else {
- throw new UnsupportedOperationException("Unsupported collection type: "
+ collection.getClass());
}
- return resultRows;
+ return true;
}
}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java
index 496a09e4d5..89df15c62b 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java
@@ -26,28 +26,57 @@ import org.apache.pinot.common.exception.QueryException;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.testng.Assert;
import org.testng.annotations.Test;
+// TODO: add tests for Agg / GroupBy / Distinct result blocks
public class LeafStageTransferableBlockOperatorTest {
@Test
public void shouldReturnDataBlockThenMetadataBlock() {
// Given:
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext("SELECT strCol, intCol FROM tbl");
DataSchema schema = new DataSchema(new String[]{"strCol", "intCol"},
new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING,
DataSchema.ColumnDataType.INT});
List<InstanceResponseBlock> resultsBlockList =
Collections.singletonList(new InstanceResponseBlock(
- new SelectionResultsBlock(schema, Arrays.asList(new Object[]{"foo",
1}, new Object[]{"", 2})), null));
+ new SelectionResultsBlock(schema, Arrays.asList(new Object[]{"foo",
1}, new Object[]{"", 2})), queryContext));
LeafStageTransferableBlockOperator operator = new
LeafStageTransferableBlockOperator(resultsBlockList, schema);
// When:
- TransferableBlock transferableBlock = operator.nextBlock();
+ TransferableBlock resultBlock = operator.nextBlock();
// Then:
- Assert.assertEquals(transferableBlock.getContainer().get(0), new
Object[]{"foo", 1});
- Assert.assertEquals(transferableBlock.getContainer().get(1), new
Object[]{"", 2});
+ Assert.assertEquals(resultBlock.getContainer().get(0), new Object[]{"foo",
1});
+ Assert.assertEquals(resultBlock.getContainer().get(1), new Object[]{"",
2});
+ Assert.assertTrue(operator.nextBlock().isEndOfStreamBlock(), "Expected EOS
after reading two rows");
+ }
+
+ @Test
+ public void shouldHandleDesiredDataSchemaConversionCorrectly() {
+ // Given:
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT boolCol, tsCol, boolCol AS newNamedBoolCol FROM tbl");
+ DataSchema resultSchema = new DataSchema(new String[]{"boolCol", "tsCol"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BOOLEAN,
DataSchema.ColumnDataType.TIMESTAMP});
+ DataSchema desiredSchema = new DataSchema(new String[]{"boolCol", "tsCol",
"newNamedBoolCol"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BOOLEAN,
DataSchema.ColumnDataType.TIMESTAMP,
+ DataSchema.ColumnDataType.BOOLEAN});
+ List<InstanceResponseBlock> resultsBlockList =
Collections.singletonList(new InstanceResponseBlock(
+ new SelectionResultsBlock(resultSchema, Arrays.asList(new Object[]{1,
1660000000000L},
+ new Object[]{0, 1600000000000L})), queryContext));
+ LeafStageTransferableBlockOperator operator = new
LeafStageTransferableBlockOperator(resultsBlockList,
+ desiredSchema);
+
+ // When:
+ TransferableBlock resultBlock = operator.nextBlock();
+
+ // Then:
+ Assert.assertEquals(resultBlock.getContainer().get(0), new Object[]{true,
new Timestamp(1660000000000L), true});
+ Assert.assertEquals(resultBlock.getContainer().get(1), new Object[]{false,
new Timestamp(1600000000000L), false});
Assert.assertTrue(operator.nextBlock().isEndOfStreamBlock(), "Expected EOS
after reading two rows");
}
@@ -55,67 +84,70 @@ public class LeafStageTransferableBlockOperatorTest {
public void shouldHandleCanonicalizationCorrectly() {
// TODO: not all stored types are supported, add additional datatype when
they are supported.
// Given:
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext("SELECT boolCol, tsCol FROM tbl");
DataSchema schema = new DataSchema(new String[]{"boolCol", "tsCol"},
new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BOOLEAN,
DataSchema.ColumnDataType.TIMESTAMP});
List<InstanceResponseBlock> resultsBlockList =
Collections.singletonList(new InstanceResponseBlock(
new SelectionResultsBlock(schema, Arrays.asList(new Object[]{1,
1660000000000L},
- new Object[]{0, 1600000000000L})), null));
+ new Object[]{0, 1600000000000L})), queryContext));
LeafStageTransferableBlockOperator operator = new
LeafStageTransferableBlockOperator(resultsBlockList, schema);
// When:
- TransferableBlock transferableBlock = operator.nextBlock();
+ TransferableBlock resultBlock = operator.nextBlock();
// Then:
- Assert.assertEquals(transferableBlock.getContainer().get(0), new
Object[]{true, new Timestamp(1660000000000L)});
- Assert.assertEquals(transferableBlock.getContainer().get(1), new
Object[]{false, new Timestamp(1600000000000L)});
+ Assert.assertEquals(resultBlock.getContainer().get(0), new Object[]{true,
new Timestamp(1660000000000L)});
+ Assert.assertEquals(resultBlock.getContainer().get(1), new Object[]{false,
new Timestamp(1600000000000L)});
Assert.assertTrue(operator.nextBlock().isEndOfStreamBlock(), "Expected EOS
after reading two rows");
}
@Test
public void shouldReturnMultipleDataBlockThenMetadataBlock() {
// Given:
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext("SELECT strCol, intCol FROM tbl");
DataSchema schema = new DataSchema(new String[]{"strCol", "intCol"},
new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING,
DataSchema.ColumnDataType.INT});
List<InstanceResponseBlock> resultsBlockList = Arrays.asList(
new InstanceResponseBlock(new SelectionResultsBlock(schema,
- Arrays.asList(new Object[]{"foo", 1}, new Object[]{"", 2})), null),
+ Arrays.asList(new Object[]{"foo", 1}, new Object[]{"", 2})),
queryContext),
new InstanceResponseBlock(new SelectionResultsBlock(schema,
- Arrays.asList(new Object[]{"bar", 3}, new Object[]{"foo", 4})),
null),
- new InstanceResponseBlock(new SelectionResultsBlock(schema,
Collections.emptyList()), null));
+ Arrays.asList(new Object[]{"bar", 3}, new Object[]{"foo", 4})),
queryContext),
+ new InstanceResponseBlock(new SelectionResultsBlock(schema,
Collections.emptyList()), queryContext));
LeafStageTransferableBlockOperator operator = new
LeafStageTransferableBlockOperator(resultsBlockList, schema);
// When:
- TransferableBlock transferableBlock1 = operator.nextBlock();
- TransferableBlock transferableBlock2 = operator.nextBlock();
- TransferableBlock transferableBlock3 = operator.nextBlock();
+ TransferableBlock resultBlock1 = operator.nextBlock();
+ TransferableBlock resultBlock2 = operator.nextBlock();
+ TransferableBlock resultBlock3 = operator.nextBlock();
// Then:
- Assert.assertEquals(transferableBlock1.getContainer().get(0), new
Object[]{"foo", 1});
- Assert.assertEquals(transferableBlock1.getContainer().get(1), new
Object[]{"", 2});
- Assert.assertEquals(transferableBlock2.getContainer().get(0), new
Object[]{"bar", 3});
- Assert.assertEquals(transferableBlock2.getContainer().get(1), new
Object[]{"foo", 4});
- Assert.assertEquals(transferableBlock3.getContainer().size(), 0);
+ Assert.assertEquals(resultBlock1.getContainer().get(0), new
Object[]{"foo", 1});
+ Assert.assertEquals(resultBlock1.getContainer().get(1), new Object[]{"",
2});
+ Assert.assertEquals(resultBlock2.getContainer().get(0), new
Object[]{"bar", 3});
+ Assert.assertEquals(resultBlock2.getContainer().get(1), new
Object[]{"foo", 4});
+ Assert.assertEquals(resultBlock3.getContainer().size(), 0);
Assert.assertTrue(operator.nextBlock().isEndOfStreamBlock(), "Expected EOS
after reading two rows");
}
@Test
public void shouldGetErrorBlockWhenInstanceResponseContainsError() {
// Given:
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext("SELECT strCol, intCol FROM tbl");
DataSchema schema = new DataSchema(new String[]{"strCol", "intCol"},
new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING,
DataSchema.ColumnDataType.INT});
InstanceResponseBlock errorBlock = new InstanceResponseBlock();
errorBlock.addException(QueryException.QUERY_EXECUTION_ERROR.getErrorCode(),
"foobar");
List<InstanceResponseBlock> resultsBlockList = Arrays.asList(
new InstanceResponseBlock(new SelectionResultsBlock(schema,
- Arrays.asList(new Object[]{"foo", 1}, new Object[]{"", 2})), null),
+ Arrays.asList(new Object[]{"foo", 1}, new Object[]{"", 2})),
queryContext),
errorBlock,
- new InstanceResponseBlock(new SelectionResultsBlock(schema,
Collections.emptyList()), null));
+ new InstanceResponseBlock(new SelectionResultsBlock(schema,
Collections.emptyList()), queryContext));
LeafStageTransferableBlockOperator operator = new
LeafStageTransferableBlockOperator(resultsBlockList, schema);
// When:
- TransferableBlock transferableBlock = operator.nextBlock();
+ TransferableBlock resultBlock = operator.nextBlock();
// Then:
- Assert.assertTrue(transferableBlock.isErrorBlock());
+ Assert.assertTrue(resultBlock.isErrorBlock());
}
}
diff --git a/pinot-query-runtime/src/test/resources/queries/Aggregates.json
b/pinot-query-runtime/src/test/resources/queries/Aggregates.json
index b82be2cb5e..cba6bab972 100644
--- a/pinot-query-runtime/src/test/resources/queries/Aggregates.json
+++ b/pinot-query-runtime/src/test/resources/queries/Aggregates.json
@@ -3,19 +3,19 @@
"tables": {
"tbl": {
"schema": [
- { "name": "int_col", "type": "INT" },
- { "name": "double_col", "type": "DOUBLE" },
- { "name": "string_col", "type": "STRING" },
- { "name": "bool_col", "type": "BOOLEAN" }
+ {"name": "int_col", "type": "INT"},
+ {"name": "double_col", "type": "DOUBLE"},
+ {"name": "string_col", "type": "STRING"},
+ {"name": "bool_col", "type": "BOOLEAN"}
],
"inputs": [
- [ 2, 300, "a", true ],
- [ 2, 400, "a", true ],
- [ 3, 100, "b", false ],
- [ 100, 1, "b", false ],
- [ 101, 1.01, "c", false ],
- [ 150, 1.5, "c", false ],
- [ 175, 1.75, "c", true ]
+ [2, 300, "a", true],
+ [2, 400, "a", true],
+ [3, 100, "b", false],
+ [100, 1, "b", false],
+ [101, 1.01, "c", false],
+ [150, 1.5, "c", false],
+ [175, 1.75, "c", true]
]
}
},
@@ -147,21 +147,21 @@
"tables": {
"tbl": {
"schema": [
- { "name": "int_col", "type": "INT" },
- { "name": "double_col", "type": "DOUBLE" },
- { "name": "string_col", "type": "STRING" },
- { "name": "bool_col", "type": "BOOLEAN" }
+ {"name": "int_col", "type": "INT"},
+ {"name": "double_col", "type": "DOUBLE"},
+ {"name": "string_col", "type": "STRING"},
+ {"name": "bool_col", "type": "BOOLEAN"}
],
"inputs": [
- [ 2, 300, "a", false ],
- [ 2, 400, "a", true ],
- [ 3, 100, "b", true ],
- [ 0.001, 1, "b", false ],
- [ 101, 1.01, "c", false ],
- [ 150, 1.5, "c", true ],
- [ 175, 1.75, "c", true ],
- [ -10000, 1.75, "c", false ],
- [ -2, 0.5, "c", false ]
+ [2, 300, "a", false],
+ [2, 400, "a", true],
+ [3, 100, "b", true],
+ [0.001, 1, "b", false],
+ [101, 1.01, "c", false],
+ [150, 1.5, "c", true],
+ [175, 1.75, "c", true],
+ [-10000, 1.75, "c", false],
+ [-2, 0.5, "c", false]
]
}
},
@@ -211,5 +211,48 @@
"sql": "select string_col, count(int_col) from {tbl} a group by
string_col order by string_col having (select count(*) from {tbl} b where
count(int_col) = b.int_col) > 0;"
}
]
+ },
+ "additional_aggregate": {
+ "tables": {
+ "tbl": {
+ "schema": [
+ {"name": "int_col", "type": "INT"},
+ {"name": "double_col", "type": "DOUBLE"},
+ {"name": "string_col", "type": "STRING"},
+ {"name": "bool_col", "type": "BOOLEAN"}
+ ],
+ "inputs": [
+ [2, 300, "a", false],
+ [2, 400, "a", true],
+ [3, 100, "b", true],
+ [0.001, 1, "b", false],
+ [101, 1.01, "c", false],
+ [150, 1.5, "c", true],
+ [175, 1.75, "c", true],
+ [-10000, 1.75, "c", false],
+ [-2, 0.5, "c", false]
+ ]
+ }
+ },
+ "queries": [
+ { "sql": "SELECT min(int_col), min(int_col) FROM {tbl} WHERE int_col <
100" },
+ { "sql": "SELECT min(int_col), count(int_col), count(*) FROM {tbl}" },
+ { "sql": "SELECT string_col, min(int_col), count(int_col), count(*) FROM
{tbl} GROUP BY string_col" },
+ { "sql": "SELECT string_col, bool_col, min(int_col), count(int_col),
count(*) FROM {tbl} GROUP BY string_col, bool_col" },
+ { "sql": "SELECT string_col, bool_col, min(int_col), count(int_col),
count(*) FROM {tbl} GROUP BY bool_col, string_col" },
+ { "sql": "SELECT string_col, string_col AS alias, count(int_col),
count(*) FROM {tbl} GROUP BY bool_col, string_col" },
+ { "sql": "SELECT string_col, bool_col, bool_col AS alias, sum(int_col),
count(int_col), count(*) FROM {tbl} GROUP BY bool_col, string_col" },
+ { "sql": "SELECT sum(int_col), count(int_col) FROM {tbl} GROUP BY
string_col" },
+ { "sql": "SELECT string_col, min(int_col), count(int_col), count(*) FROM
{tbl} GROUP BY string_col ORDER BY string_col" },
+ { "sql": "SELECT string_col, bool_col, min(int_col), count(int_col),
count(*) FROM {tbl} GROUP BY string_col, bool_col ORDER BY string_col" },
+ { "sql": "SELECT string_col, bool_col, min(int_col), count(int_col),
count(*) FROM {tbl} GROUP BY string_col, bool_col ORDER BY bool_col,
string_col" },
+ { "sql": "SELECT string_col, bool_col, min(int_col), count(int_col),
count(*) FROM {tbl} GROUP BY bool_col, string_col ORDER BY bool_col,
string_col" },
+ { "sql": "SELECT string_col, string_col AS alias, count(int_col),
count(*) FROM {tbl} GROUP BY bool_col, string_col ORDER BY string_col" },
+ { "sql": "SELECT string_col, string_col AS alias, count(int_col),
count(*) FROM {tbl} GROUP BY bool_col, string_col ORDER BY bool_col, alias"},
+ { "sql": "SELECT string_col, bool_col, bool_col AS alias, sum(int_col),
count(int_col), count(*) FROM {tbl} GROUP BY bool_col, string_col ORDER BY
alias" },
+ { "sql": "SELECT sum(int_col), count(int_col) FROM {tbl} GROUP BY
string_col ORDER BY string_col" },
+ { "sql": "SELECT * FROM (SELECT string_col, bool_col, min(int_col) AS m,
count(int_col), count(*) AS c FROM {tbl} GROUP BY string_col, bool_col ORDER BY
string_col) WHERE c < m" },
+ { "sql": "SELECT * FROM (SELECT string_col, bool_col, min(int_col) AS m,
count(int_col), count(*) AS c FROM {tbl} GROUP BY string_col, bool_col ORDER BY
bool_col, string_col) WHERE c < m" }
+ ]
}
-}
+}
\ No newline at end of file
diff --git a/pinot-query-runtime/src/test/resources/queries/Distincts.json
b/pinot-query-runtime/src/test/resources/queries/Distincts.json
new file mode 100644
index 0000000000..22ee05725c
--- /dev/null
+++ b/pinot-query-runtime/src/test/resources/queries/Distincts.json
@@ -0,0 +1,47 @@
+{
+ "general_distinct": {
+ "tables": {
+ "tbl": {
+ "schema": [
+ { "name": "intCol", "type": "INT" },
+ { "name": "doubleCol", "type": "DOUBLE" },
+ { "name": "strCol", "type": "STRING" },
+ { "name": "boolCol", "type": "BOOLEAN" }
+ ],
+ "inputs": [
+ [ 2, 300, "a", true ],
+ [ 2, 400, "a", true ],
+ [ 3, 100, "b", false ],
+ [ 100, 1, "b", false ],
+ [ 101, 1.01, "c", false ],
+ [ 150, 1.5, "c", false ],
+ [ 175, 1.75, "c", true ]
+ ]
+ }
+ },
+ "queries": [
+ { "sql": "SELECT DISTINCT intCol FROM {tbl}" },
+ { "sql": "SELECT intCol FROM {tbl} GROUP BY intCol" },
+ { "sql": "SELECT DISTINCT strCol, boolCol FROM {tbl}" },
+ { "sql": "SELECT strCol, boolCol FROM {tbl} GROUP BY strCol, boolCol" },
+ { "sql": "SELECT strCol, boolCol FROM {tbl} GROUP BY boolCol, strCol" },
+ { "sql": "SELECT strCol FROM {tbl} GROUP BY strCol, boolCol" },
+ { "sql": "SELECT boolCol FROM {tbl} GROUP BY strCol, boolCol" },
+ { "sql": "SELECT DISTINCT intCol, doubleCol FROM {tbl} ORDER BY
doubleCol" },
+ { "sql": "SELECT intCol FROM {tbl} GROUP BY intCol ORDER BY intCol" },
+ { "sql": "SELECT DISTINCT strCol, boolCol, intCol FROM {tbl} ORDER BY
intCol" },
+ { "sql": "SELECT strCol, boolCol FROM {tbl} GROUP BY strCol, boolCol
ORDER BY boolCol" },
+ { "sql": "SELECT strCol, boolCol FROM {tbl} GROUP BY boolCol, strCol
ORDER BY strCol" },
+ { "sql": "SELECT strCol FROM {tbl} GROUP BY strCol, boolCol ORDER BY
boolCol" },
+ { "sql": "SELECT strCol FROM (SELECT strCol FROM {tbl} GROUP BY
boolCol, strCol ORDER BY strCol) WHERE strCol <> 'foo'" },
+ { "sql": "SELECT intCol FROM (SELECT DISTINCT intCol FROM {tbl}) WHERE
intCol < 100" },
+ { "sql": "SELECT intCol FROM (SELECT intCol FROM {tbl} GROUP BY intCol)
WHERE intCol < 100" },
+ { "sql": "SELECT strCol FROM (SELECT DISTINCT strCol, boolCol FROM
{tbl}) WHERE boolCol = false" },
+ { "sql": "SELECT strCol FROM (SELECT strCol, boolCol FROM {tbl} GROUP
BY strCol, boolCol) WHERE strCol <> 'foo'" },
+ { "sql": "SELECT strCol, boolCol FROM (SELECT strCol, boolCol FROM
{tbl} GROUP BY boolCol, strCol) WHERE strCol <> 'foo'" },
+ { "sql": "SELECT strCol FROM (SELECT strCol FROM {tbl} GROUP BY strCol,
boolCol) WHERE strCol != 'lol'" },
+ { "sql": "SELECT intCol FROM (SELECT DISTINCT intCol, doubleCol FROM
{tbl} ORDER BY doubleCol) WHERE doubleCol < 100" },
+ { "sql": "SELECT strCol FROM (SELECT DISTINCT strCol, boolCol, intCol
FROM {tbl} ORDER BY intCol) WHERE intCol > 10" }
+ ]
+ }
+}
diff --git
a/pinot-query-runtime/src/test/resources/queries/LexicalStructure.json
b/pinot-query-runtime/src/test/resources/queries/LexicalStructure.json
index 400ed57fac..ea6e139ddc 100644
--- a/pinot-query-runtime/src/test/resources/queries/LexicalStructure.json
+++ b/pinot-query-runtime/src/test/resources/queries/LexicalStructure.json
@@ -135,11 +135,22 @@
"inputs": [["1"]]
}
},
- "queries": [{
- "psql": "4.1.2.6",
- "description": "numeric constant supports",
- "sql": "SELECT data, 42, 3.5, 4., .001, 5e2, 1.925e-3 FROM {tbl}"
- }]
+ "queries": [
+ {
+ "ignored": true,
+ "comments": "numeric constant without specific type will cast all
floating points into BIG DECIMAL which is not supported by type-hoisting on
leaf stage",
+ "psql": "4.1.2.6",
+ "description": "numeric constant supports",
+ "sql": "SELECT data, 42, 3.5, 4., .001, 5e2, 1.925e-3 FROM {tbl}"
+ },
+ {
+ "ignored": true,
+ "comments": "floating point comparison missing precision checker",
+ "psql": "4.1.2.6",
+ "description": "numeric constant supports",
+ "sql": "SELECT data, 42, CAST(3.5 AS FLOAT), CAST(4. AS FLOAT),
CAST(.001 AS FLOAT), CAST(5e2 AS FLOAT), CAST(1.925e-3 AS FLOAT) FROM {tbl}"
+ }
+ ]
},
"constant_casting": {
"tables": {
diff --git
a/pinot-query-runtime/src/test/resources/queries/SelectExpressions.json
b/pinot-query-runtime/src/test/resources/queries/SelectExpressions.json
index f214307b2a..af925382ef 100644
--- a/pinot-query-runtime/src/test/resources/queries/SelectExpressions.json
+++ b/pinot-query-runtime/src/test/resources/queries/SelectExpressions.json
@@ -37,7 +37,20 @@
{ "sql": "SELECT intCol as \"from\" FROM {tbl1}"},
{ "sql": "SELECT intCol as key, SUM(doubleCol + floatCol) AS aggSum FROM
{tbl1} GROUP BY intCol"},
{ "sql": "SELECT intCol, SUM(avgVal) FROM (SELECT strCol, intCol,
AVG(doubleCol) AS avgVal FROM {tbl1} GROUP BY intCol, strCol) GROUP BY intCol"},
- { "sql": "SELECT strCol, MAX(sumVal), MAX(sumVal + avgVal) AS transVal
FROM (SELECT strCol, intCol, SUM(floatCol + 2 * intCol) AS sumVal,
AVG(doubleCol) AS avgVal FROM {tbl1} GROUP BY strCol, intCol) GROUP BY strCol
ORDER BY MAX(sumVal)" }
+ { "sql": "SELECT strCol, MAX(sumVal), MAX(sumVal + avgVal) AS transVal
FROM (SELECT strCol, intCol, SUM(floatCol + 2 * intCol) AS sumVal,
AVG(doubleCol) AS avgVal FROM {tbl1} GROUP BY strCol, intCol) GROUP BY strCol
ORDER BY MAX(sumVal)" },
+ { "sql": "SELECT intCol, intCol FROM {tbl1} WHERE intCol < 100" },
+ { "sql": "SELECT intCol, intCol, doubleCol, strCol, strCol FROM {tbl1}"
},
+ { "sql": "SELECT intCol, intCol FROM {tbl1} WHERE intCol < 100 ORDER BY
doubleCol" },
+ { "sql": "SELECT intCol, intCol FROM {tbl1} WHERE intCol < 100 ORDER BY
intCol" },
+ { "sql": "SELECT intCol, intCol, doubleCol, strCol, strCol FROM {tbl1}
ORDER BY intCol" },
+ { "sql": "SELECT intCol, intCol, doubleCol, strCol, strCol FROM {tbl1}
ORDER BY strCol, intCol" },
+ { "sql": "SELECT intCol, intCol, doubleCol, strCol, strCol FROM {tbl1}
ORDER BY floatCol, intCol" },
+ { "sql": "SELECT intCol, intCol, doubleCol, strCol, strCol FROM {tbl1}
ORDER BY floatCol" },
+
+ { "sql": "SELECT intCol + alias FROM (SELECT intCol, intCol as alias
FROM {tbl1}) WHERE intCol < 100" },
+ { "sql": "SELECT intCol, intCol, doubleCol, strCol, strCol FROM {tbl1}"
},
+ { "sql": "SELECT {tbl1}.intCol, {tbl1}.intCol, {tbl1}.doubleCol,
{tbl2}.strCol, {tbl2}.strCol FROM {tbl1}, {tbl2} WHERE {tbl1}.intCol =
{tbl2}.intCol" },
+ { "sql": "SELECT {tbl2}.intCol, {tbl2}.intCol FROM {tbl1}, {tbl2} WHERE
{tbl1}.intCol = {tbl2}.intCol AND {tbl1}.intCol < 100 ORDER BY
{tbl1}.doubleCol" }
]
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
