This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch xiangfu0/agg-rewrite/02-array-agg in repository https://gitbox.apache.org/repos/asf/pinot.git
commit b80d110536e5c5b4a797d0f1a209d6fca788350d Author: Xiang Fu <[email protected]> AuthorDate: Tue Oct 20 18:10:00 2026 -0700 Support ARRAY_AGG without a data type Both query engines rewrite ARRAY_AGG(expression) and ARRAY_AGG(expression, isDistinct) to the explicit form ARRAY_AGG(expression, 'dataType'[, isDistinct]) with the type of the expression, or its element type for a multi-value input: AggregateFunctionRewriteOptimizer for a column in the single-stage engine and PinotAggregateTypeArgumentRewriteRule in the multi-stage engine. --- .../AggregateFunctionRewriteOptimizer.java | 43 ++++++++++++++++ .../AggregateFunctionRewriteOptimizerTest.java | 23 +++++++++ .../pinot/queries/ArrayAggMvQueriesTest.java | 19 +++++++ .../apache/pinot/queries/ArrayAggQueriesTest.java | 33 ++++++++++++ .../pinot/queries/ArrayAggStarTreeQueriesTest.java | 42 ++++++++++++---- .../PinotAggregateTypeArgumentRewriteRule.java | 58 +++++++++++++++++----- .../src/test/resources/queries/AggregatePlans.json | 44 ++++++++++++++++ .../src/test/resources/queries/UDFAggregates.json | 41 +++++++++++++++ .../pinot/segment/spi/AggregationFunctionType.java | 22 ++++++-- .../apache/pinot/spi/utils/CommonConstants.java | 2 +- 10 files changed, 299 insertions(+), 28 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java index 5bbc3274e05..b0227ee2dda 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java @@ -35,6 +35,7 @@ import org.apache.pinot.spi.utils.CommonConstants; /// Always applied: /// - FIRST_WITH_TIME(dataColumn, timeColumn) -> FIRST_WITH_TIME(dataColumn, timeColumn, 'dataType') /// - LAST_WITH_TIME(dataColumn, timeColumn) -> LAST_WITH_TIME(dataColumn, timeColumn, 'dataType') +/// - ARRAY_AGG(column[, isDistinct]) -> ARRAY_AGG(column, 'dataType'[, isDistinct]) /// /// Applied only when enabled through the `autoRewriteAggregationType` query option: /// - MIN(stringType) -> MINSTRING @@ -97,6 +98,10 @@ public class AggregateFunctionRewriteOptimizer implements StatementOptimizer { addDataTypeArgument(function, fieldSpec); return; } + if (normalizedName.equals(AggregationFunctionType.ARRAYAGG.name()) && isArrayAggWithoutDataType(function)) { + addArrayAggDataTypeArgument(function, fieldSpec); + return; + } if (rewriteToTypedVariants) { rewriteToTypedVariant(function, functionName, fieldSpec.getDataType().getStoredType()); } @@ -126,6 +131,44 @@ public class AggregateFunctionRewriteOptimizer implements StatementOptimizer { } } + /// Matches ARRAY_AGG(column) and ARRAY_AGG(column, isDistinct); a string second argument is the explicit data type. + private static boolean isArrayAggWithoutDataType(Function function) { + int numOperands = function.getOperandsSize(); + if (numOperands == 1) { + return true; + } + if (numOperands == 2) { + Expression secondOperand = function.getOperands().get(1); + return secondOperand.isSetLiteral() && secondOperand.getLiteral().isSetBoolValue(); + } + return false; + } + + /// Inserts the column's logical type, the element type for a multi-value column, as the second argument when the + /// explicit form supports it. + private static void addArrayAggDataTypeArgument(Function function, FieldSpec fieldSpec) { + switch (fieldSpec.getDataType()) { + case BOOLEAN: + case INT: + case LONG: + case TIMESTAMP: + case FLOAT: + case DOUBLE: + case BIG_DECIMAL: + case STRING: + case BYTES: + case UUID: + function.getOperands().add(1, RequestUtils.getLiteralExpression(fieldSpec.getDataType().name())); + break; + case JSON: + // Same as the multi-stage engine, which types a JSON column as VARCHAR. + function.getOperands().add(1, RequestUtils.getLiteralExpression(FieldSpec.DataType.STRING.name())); + break; + default: + break; + } + } + private static void rewriteToTypedVariant(Function function, String functionName, FieldSpec.DataType operandType) { // Rewrite MIN(stringCol) and MAX(stringCol) to MINSTRING / MAXSTRING // Rewrite MIN(longCol) and MAX(longCol) to MINLONG / MAXLONG diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java index d44bc13a13a..c7693b581ca 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java @@ -36,6 +36,7 @@ public class AggregateFunctionRewriteOptimizerTest { .addSingleValueDimension("longCol", DataType.LONG) .addSingleValueDimension("floatCol", DataType.FLOAT) .addSingleValueDimension("doubleCol", DataType.DOUBLE) + .addSingleValueDimension("bigDecimalCol", DataType.BIG_DECIMAL) .addSingleValueDimension("stringCol", DataType.STRING) .addSingleValueDimension("boolCol", DataType.BOOLEAN) .addSingleValueDimension("timestampCol", DataType.TIMESTAMP) @@ -73,6 +74,28 @@ public class AggregateFunctionRewriteOptimizerTest { + "FROM testTable"); } + @DataProvider + public Object[][] arrayAggDataTypes() { + return new Object[][]{ + {"intCol", "INT"}, {"longCol", "LONG"}, {"floatCol", "FLOAT"}, {"doubleCol", "DOUBLE"}, + {"bigDecimalCol", "BIG_DECIMAL"}, {"stringCol", "STRING"}, {"jsonCol", "STRING"}, {"boolCol", "BOOLEAN"}, + {"timestampCol", "TIMESTAMP"}, {"bytesCol", "BYTES"}, {"uuidCol", "UUID"}, {"mvCol", "INT"} + }; + } + + @Test(dataProvider = "arrayAggDataTypes") + public void testArrayAggInfersDataType(String column, String dataType) { + assertRewrite("SELECT ARRAY_AGG(" + column + "), arrayAgg(" + column + ", true) FROM testTable", + "SELECT ARRAY_AGG(" + column + ", '" + dataType + "'), arrayAgg(" + column + ", '" + dataType + "', true) " + + "FROM testTable"); + } + + @Test + public void testArrayAggKeepsExplicitAndUnsupportedForms() { + assertUnchanged("SELECT ARRAY_AGG(intCol, 'LONG'), ARRAY_AGG(intCol, 'LONG', true), ARRAY_AGG(intCol, 1), " + + "ARRAY_AGG(missingCol), ARRAY_AGG(trim(stringCol), true) FROM testTable"); + } + @Test public void testRewritesNestedAggregationsConsistently() { // Aliases, transforms, FILTER, HAVING and ORDER BY wrap the aggregation in another function. diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggMvQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggMvQueriesTest.java index 118d27d143c..d21125a17e7 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggMvQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggMvQueriesTest.java @@ -26,6 +26,7 @@ import java.util.List; import java.util.Set; import org.apache.commons.io.FileUtils; import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock; import org.apache.pinot.core.operator.query.AggregationOperator; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; @@ -162,6 +163,24 @@ public class ArrayAggMvQueriesTest extends BaseQueriesTest { assertEquals(((String[]) resultTable.getRows().get(0)[4]).length, 2 * NUM_RECORDS); } + /// The data type argument can be omitted for a multi-value column; the optimizer supplies its element type. + @Test + public void testArrayAggMvWithInferredDataType() { + String query = "SELECT ArrayAgg(intMV), ArrayAgg(longMV, true), ArrayAgg(floatMV), ArrayAgg(doubleMV, true), " + + "ArrayAgg(stringMV) FROM testTableMv"; + ResultTable resultTable = getBrokerResponseForOptimizedQuery(query, SCHEMA).getResultTable(); + assertEquals(resultTable.getDataSchema().getColumnDataTypes(), new ColumnDataType[]{ + ColumnDataType.INT_ARRAY, ColumnDataType.LONG_ARRAY, ColumnDataType.FLOAT_ARRAY, ColumnDataType.DOUBLE_ARRAY, + ColumnDataType.STRING_ARRAY + }); + Object[] row = resultTable.getRows().get(0); + assertEquals(((int[]) row[0]).length, 8 * NUM_RECORDS); + assertEquals(((long[]) row[1]).length, 2 * NUM_RECORDS); + assertEquals(((float[]) row[2]).length, 8 * NUM_RECORDS); + assertEquals(((double[]) row[3]).length, 2 * NUM_RECORDS); + assertEquals(((String[]) row[4]).length, 8 * NUM_RECORDS); + } + @AfterClass public void tearDown() throws IOException { diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggQueriesTest.java index a1749adeeea..cba00735016 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggQueriesTest.java @@ -20,6 +20,7 @@ package org.apache.pinot.queries; import java.io.File; import java.io.IOException; +import java.lang.reflect.Array; import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; @@ -29,6 +30,7 @@ import java.util.Set; import java.util.UUID; import org.apache.commons.io.FileUtils; import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock; import org.apache.pinot.core.operator.query.AggregationOperator; @@ -83,6 +85,11 @@ public class ArrayAggQueriesTest extends BaseQueriesTest { .addSingleValueDimension(BYTES_COLUMN, DataType.BYTES) .addSingleValueDimension(UUID_COLUMN, DataType.UUID) .build(); + // The JSON column aggregates to strings + private static final ColumnDataType[] EXPECTED_COLUMN_DATA_TYPES = { + ColumnDataType.INT_ARRAY, ColumnDataType.LONG_ARRAY, ColumnDataType.FLOAT_ARRAY, ColumnDataType.DOUBLE_ARRAY, + ColumnDataType.STRING_ARRAY, ColumnDataType.STRING_ARRAY, ColumnDataType.BYTES_ARRAY, ColumnDataType.UUID_ARRAY + }; private Set<Integer> _values; private int[] _expectedCardinalityResults; @@ -226,6 +233,32 @@ public class ArrayAggQueriesTest extends BaseQueriesTest { assertEquals(((String[]) resultTable.getRows().get(0)[7]).length, _expectedCardinalityResults[7]); } + /// The data type argument can be omitted for a column; the optimizer supplies it from the schema. + @Test + public void testArrayAggWithInferredDataType() { + String query = "SELECT ArrayAgg(intColumn), ArrayAgg(longColumn), ArrayAgg(floatColumn), ArrayAgg(doubleColumn), " + + "ArrayAgg(stringColumn), ArrayAgg(jsonColumn), ArrayAgg(bytesColumn), ArrayAgg(uuidColumn) FROM testTable"; + ResultTable resultTable = getBrokerResponseForOptimizedQuery(query, SCHEMA).getResultTable(); + assertEquals(resultTable.getDataSchema().getColumnDataTypes(), EXPECTED_COLUMN_DATA_TYPES); + Object[] row = resultTable.getRows().get(0); + for (int i = 0; i < 8; i++) { + assertEquals(Array.getLength(row[i]), 4 * NUM_RECORDS); + } + } + + @Test + public void testArrayAggDistinctWithInferredDataType() { + String query = "SELECT ArrayAgg(intColumn, true), ArrayAgg(longColumn, true), ArrayAgg(floatColumn, true), " + + "ArrayAgg(doubleColumn, true), ArrayAgg(stringColumn, true), ArrayAgg(jsonColumn, true), " + + "ArrayAgg(bytesColumn, true), ArrayAgg(uuidColumn, true) FROM testTable"; + ResultTable resultTable = getBrokerResponseForOptimizedQuery(query, SCHEMA).getResultTable(); + assertEquals(resultTable.getDataSchema().getColumnDataTypes(), EXPECTED_COLUMN_DATA_TYPES); + Object[] row = resultTable.getRows().get(0); + for (int i = 0; i < 8; i++) { + assertEquals(Array.getLength(row[i]), _expectedCardinalityResults[i]); + } + } + @AfterClass public void tearDown() throws IOException { diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggStarTreeQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggStarTreeQueriesTest.java index ac3b9c4565a..6d64c2d1d10 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggStarTreeQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ArrayAggStarTreeQueriesTest.java @@ -64,6 +64,13 @@ public class ArrayAggStarTreeQueriesTest extends BaseQueriesTest { private static final String LONG_METRIC = "lm"; private static final String STRING_METRIC = "sm"; private static final String BYTES_METRIC = "bm"; + private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension(DIMENSION, DataType.INT) + .addSingleValueDimension(INT_METRIC, DataType.INT) + .addSingleValueDimension(LONG_METRIC, DataType.LONG) + .addSingleValueDimension(STRING_METRIC, DataType.STRING) + .addSingleValueDimension(BYTES_METRIC, DataType.BYTES) + .build(); private static final Map<String, String> STAR_TREE_ON = Map.of("useStarTree", "true"); private static final Map<String, String> STAR_TREE_OFF = Map.of("useStarTree", "false"); @@ -95,13 +102,6 @@ public class ArrayAggStarTreeQueriesTest extends BaseQueriesTest { throws Exception { FileUtils.deleteDirectory(INDEX_DIR); - Schema schema = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) - .addSingleValueDimension(DIMENSION, DataType.INT) - .addSingleValueDimension(INT_METRIC, DataType.INT) - .addSingleValueDimension(LONG_METRIC, DataType.LONG) - .addSingleValueDimension(STRING_METRIC, DataType.STRING) - .addSingleValueDimension(BYTES_METRIC, DataType.BYTES) - .build(); // The LONG metric is raw-encoded: arrayAgg star-tree aggregation supports both dictionary-encoded and raw // single-value source columns, so the LONG tests double as raw-metric build/query coverage. TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) @@ -119,7 +119,7 @@ public class ArrayAggStarTreeQueriesTest extends BaseQueriesTest { rows.add(row); } - SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, SCHEMA); segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath()); SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); @@ -271,11 +271,33 @@ public class ArrayAggStarTreeQueriesTest extends BaseQueriesTest { "Non-distinct arrayAgg must not use the star-tree, so it scans the same docs with the option on or off"); } + /// Without the data type argument, distinct arrayAgg is rewritten to the explicit form with each column's type, which + /// the star-tree still serves. + @Test + public void testDistinctArrayAggWithInferredDataTypeUsesStarTree() { + String query = String.format( + "SELECT %s, arrayAgg(%s, true), arrayAgg(%s, true), arrayAgg(%s, true), arrayAgg(%s, true) FROM %s " + + "GROUP BY %s ORDER BY %s", DIMENSION, INT_METRIC, LONG_METRIC, STRING_METRIC, BYTES_METRIC, + RAW_TABLE_NAME, DIMENSION, DIMENSION); + String explicitQuery = String.format( + "SELECT %s, arrayAgg(%s, 'INT', true), arrayAgg(%s, 'LONG', true), arrayAgg(%s, 'STRING', true), " + + "arrayAgg(%s, 'BYTES', true) FROM %s GROUP BY %s ORDER BY %s", DIMENSION, INT_METRIC, LONG_METRIC, + STRING_METRIC, BYTES_METRIC, RAW_TABLE_NAME, DIMENSION, DIMENSION); + BrokerResponseNative starTree = getBrokerResponseForOptimizedQuery("SET useStarTree=true; " + query, SCHEMA); + BrokerResponseNative rawScan = getBrokerResponse(explicitQuery, STAR_TREE_OFF); + assertTrue(starTree.getNumDocsScanned() < rawScan.getNumDocsScanned(), + "Expected the star-tree to serve the distinct arrayAgg query without a data type argument"); + assertSameDistinctValues(starTree.getResultTable(), rawScan.getResultTable()); + } + /// Runs the query with the star-tree enabled and disabled and asserts each arrayAgg cell (an unordered MV of /// distinct values) is set-equal between the two, per group. private void assertStarTreeMatchesRawScan(String query) { - ResultTable starTree = getBrokerResponse(query, STAR_TREE_ON).getResultTable(); - ResultTable rawScan = getBrokerResponse(query, STAR_TREE_OFF).getResultTable(); + assertSameDistinctValues(getBrokerResponse(query, STAR_TREE_ON).getResultTable(), + getBrokerResponse(query, STAR_TREE_OFF).getResultTable()); + } + + private static void assertSameDistinctValues(ResultTable starTree, ResultTable rawScan) { List<Object[]> starTreeRows = starTree.getRows(); List<Object[]> rawScanRows = rawScan.getRows(); assertEquals(starTreeRows.size(), rawScanRows.size()); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java index dc4f8455363..4c657c7bca2 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java @@ -33,6 +33,8 @@ import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.hint.RelHint; import org.apache.calcite.rel.logical.LogicalAggregate; import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeField; import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.type.SqlTypeName; @@ -43,6 +45,7 @@ import org.apache.pinot.segment.spi.AggregationFunctionType; /// Supplies the explicit data type argument from the type of the data argument: /// - FIRST_WITH_TIME(dataColumn, timeColumn) -> FIRST_WITH_TIME(dataColumn, timeColumn, 'dataType') /// - LAST_WITH_TIME(dataColumn, timeColumn) -> LAST_WITH_TIME(dataColumn, timeColumn, 'dataType') +/// - ARRAY_AGG(expression[, isDistinct]) -> ARRAY_AGG(expression, 'dataType'[, isDistinct]) /// /// The result type of the call is unchanged. The literal is projected directly beneath the aggregate, appended to the /// existing project when there is one, because [PinotAggregateExchangeNodeInsertRule] and `AggregatePushdownRule` @@ -116,7 +119,11 @@ public class PinotAggregateTypeArgumentRewriteRule extends RelOptRule { continue; } List<Integer> argList = new ArrayList<>(aggCall.getArgList()); - argList.add(literalIndexes[i]); + if (isArrayAgg(aggCall)) { + argList.add(1, literalIndexes[i]); + } else { + argList.add(literalIndexes[i]); + } // Keeps the original result type, so the rewrite does not change the result schema. aggCalls.add(aggCall.withArgList(argList)); } @@ -124,20 +131,35 @@ public class PinotAggregateTypeArgumentRewriteRule extends RelOptRule { aggRel.copy(aggRel.getTraitSet(), newInput, aggRel.getGroupSet(), aggRel.getGroupSets(), aggCalls)); } - /// Returns the data type argument to append, or `null` when the call already has one or is not such a call. + /// Returns the data type argument to add, or `null` when the call already has one or is not such a call. @Nullable private static String getDataTypeArgument(AggregateCall aggCall, RelNode input) { - if (aggCall.getArgList().size() != 2) { - return null; - } - String name = AggregationFunctionType.getNormalizedAggregationFunctionName(aggCall.getAggregation().getName()); - if (!name.equals(AggregationFunctionType.FIRSTWITHTIME.name()) - && !name.equals(AggregationFunctionType.LASTWITHTIME.name())) { - return null; + List<Integer> argList = aggCall.getArgList(); + List<RelDataTypeField> fields = input.getRowType().getFieldList(); + RelDataType dataType; + if (isArrayAgg(aggCall)) { + // ARRAY_AGG(expression) or ARRAY_AGG(expression, isDistinct); a string second argument is the data type + if (argList.size() != 1 + && (argList.size() != 2 || fields.get(argList.get(1)).getType().getSqlTypeName() != SqlTypeName.BOOLEAN)) { + return null; + } + // A multi-value input has the type of its elements + dataType = fields.get(argList.get(0)).getType(); + if (dataType.getComponentType() != null) { + dataType = dataType.getComponentType(); + } + } else { + if (argList.size() != 2) { + return null; + } + String name = AggregationFunctionType.getNormalizedAggregationFunctionName(aggCall.getAggregation().getName()); + if (!name.equals(AggregationFunctionType.FIRSTWITHTIME.name()) + && !name.equals(AggregationFunctionType.LASTWITHTIME.name())) { + return null; + } + dataType = fields.get(argList.get(0)).getType(); } - SqlTypeName dataType = - input.getRowType().getFieldList().get(aggCall.getArgList().get(0)).getType().getSqlTypeName(); - switch (dataType) { + switch (dataType.getSqlTypeName()) { case TINYINT: case SMALLINT: case INTEGER: @@ -149,16 +171,28 @@ public class PinotAggregateTypeArgumentRewriteRule extends RelOptRule { return "FLOAT"; case DOUBLE: return "DOUBLE"; + case DECIMAL: + return "BIG_DECIMAL"; case CHAR: case VARCHAR: return "STRING"; + case BINARY: + case VARBINARY: + return "BYTES"; case BOOLEAN: return "BOOLEAN"; case TIMESTAMP: return "TIMESTAMP"; + case UUID: + return "UUID"; default: // Rejected during validation; leave the call as it is rather than guess a type. return null; } } + + private static boolean isArrayAgg(AggregateCall aggCall) { + return AggregationFunctionType.getNormalizedAggregationFunctionName(aggCall.getAggregation().getName()) + .equals(AggregationFunctionType.ARRAYAGG.name()); + } } diff --git a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json index c10f7185b79..aabc1a355e0 100644 --- a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json +++ b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json @@ -189,6 +189,50 @@ "sql": "EXPLAIN PLAN FOR SELECT FIRST_WITH_TIME(a.col4, a.ts) FROM a", "expectedException": ".*Unsupported data type for .*\\(dataColumn, timeColumn\\): DECIMAL.*" }, + { + "description": "ARRAY_AGG without a data type infers it from the expression as the second argument", + "sql": "EXPLAIN PLAN FOR SELECT ARRAY_AGG(a.col1), ARRAY_AGG(a.col3, true), ARRAY_AGG(a.col4), ARRAY_AGG(a.ts_timestamp, true), ARRAY_AGG(a.col5), ARRAY_AGG(a.col7 + 1) FROM a", + "output": [ + "Execution Plan", + "\nPinotLogicalAggregate(group=[{}], agg#0=[ARRAYAGG($0, _UTF-8'STRING')], agg#1=[ARRAYAGG($1, _UTF-8'INT', true)], agg#2=[ARRAYAGG($2, _UTF-8'BIG_DECIMAL')], agg#3=[ARRAYAGG($3, _UTF-8'TIMESTAMP', true)], agg#4=[ARRAYAGG($4, _UTF-8'BOOLEAN')], agg#5=[ARRAYAGG($5, _UTF-8'LONG')], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash])", + "\n PinotLogicalAggregate(group=[{}], agg#0=[ARRAYAGG($0, _UTF-8'STRING')], agg#1=[ARRAYAGG($1, _UTF-8'INT', true)], agg#2=[ARRAYAGG($3, _UTF-8'BIG_DECIMAL')], agg#3=[ARRAYAGG($4, _UTF-8'TIMESTAMP', true)], agg#4=[ARRAYAGG($5, _UTF-8'BOOLEAN')], agg#5=[ARRAYAGG($6, _UTF-8'LONG')], aggType=[LEAF])", + "\n LogicalProject($f0=[$0], $f1=[$2], $f2=[true], $f3=[$3], $f4=[$8], $f5=[$4], $f6=[+($6, 1)], $f7=[_UTF-8'STRING'], $f8=[_UTF-8'INT'], $f9=[_UTF-8'BIG_DECIMAL'], $f10=[_UTF-8'TIMESTAMP'], $f11=[_UTF-8'BOOLEAN'], $f12=[_UTF-8'LONG'])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, + { + "description": "ARRAY_AGG without a data type next to the explicit forms in a grouped aggregate", + "sql": "EXPLAIN PLAN FOR SELECT a.col2, ARRAY_AGG(a.col3, true), ARRAY_AGG(a.col3, 'LONG', true), ARRAY_AGG(a.col1, 'STRING') FROM a GROUP BY a.col2", + "output": [ + "Execution Plan", + "\nPinotLogicalAggregate(group=[{0}], agg#0=[ARRAYAGG($1, _UTF-8'INT', true)], agg#1=[ARRAYAGG($2, _UTF-8'LONG', true)], agg#2=[ARRAYAGG($3, _UTF-8'STRING')], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash[0]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[ARRAYAGG($1, _UTF-8'INT', true)], agg#1=[ARRAYAGG($1, _UTF-8'LONG', true)], agg#2=[ARRAYAGG($4, _UTF-8'STRING')], aggType=[LEAF])", + "\n LogicalProject($f0=[$1], $f1=[$2], $f2=[true], $f3=[_UTF-8'LONG'], $f4=[$0], $f5=[_UTF-8'STRING'], $f6=[_UTF-8'INT'])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, + { + "description": "ARRAY_AGG without a data type uses the element type of a multi-value column", + "sql": "EXPLAIN PLAN FOR SELECT ARRAY_AGG(e.mcol1), ARRAY_AGG(e.mcol2, true) FROM e", + "output": [ + "Execution Plan", + "\nPinotLogicalAggregate(group=[{}], agg#0=[ARRAYAGG($0, _UTF-8'STRING')], agg#1=[ARRAYAGG($1, _UTF-8'LONG', true)], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash])", + "\n PinotLogicalAggregate(group=[{}], agg#0=[ARRAYAGG($0, _UTF-8'STRING')], agg#1=[ARRAYAGG($1, _UTF-8'LONG', true)], aggType=[LEAF])", + "\n LogicalProject($f0=[$7], $f1=[$8], $f2=[true], $f3=[_UTF-8'STRING'], $f4=[_UTF-8'LONG'])", + "\n PinotLogicalTableScan(table=[[default, e]])", + "\n" + ] + }, + { + "description": "ARRAY_AGG without a data type rejects an unsupported expression type", + "sql": "EXPLAIN PLAN FOR SELECT ARRAY_AGG(CAST(a.col1 AS DATE), true) FROM a", + "expectedException": ".*Unsupported data type for .*\\(expression\\): DATE.*" + }, { "description": "MIN on STRING column rewritten to MINSTRING", "sql": "SET autoRewriteAggregationType=true; EXPLAIN PLAN FOR SELECT MIN(a.col1) FROM a", diff --git a/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json b/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json index 01a19d7837c..0cb30e6eb95 100644 --- a/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json +++ b/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json @@ -254,5 +254,46 @@ "outputs": [["a", 1, 1, 2, 2], ["b", 2, 2, 2, 2], ["c", 3, 3, 3, 3]] } ] + }, + "pinot_array_agg_tests": { + "tables": { + "tbl": { + "schema": [ + {"name": "int_col", "type": "INT"}, + {"name": "long_col", "type": "LONG"}, + {"name": "double_col", "type": "DOUBLE"}, + {"name": "decimal_col", "type": "BIG_DECIMAL"}, + {"name": "string_col", "type": "STRING"}, + {"name": "json_col", "type": "JSON"}, + {"name": "int_mv_col", "type": "INT", "isSingleValue": false}, + {"name": "string_mv_col", "type": "STRING", "isSingleValue": false} + ], + "inputs": [ + [1, 10, 1.5, "1.5", "a", "{\"k\":1}", [1, 2], ["x", "y"]], + [2, 20, 2.5, "2.5", "a", "{\"k\":2}", [2, 3], ["y"]], + [2, 20, 2.5, "2.5", "b", "{\"k\":2}", [3], ["z", "x"]], + [3, 30, 3.5, "3.5", "b", "{\"k\":3}", [4, 5], ["z"]] + ] + } + }, + "queries": [ + { + "comment": "The data type argument of ARRAY_AGG is inferred when omitted", + "sql": "SELECT arraySortInt(ARRAY_AGG(int_col)), arraySortInt(ARRAY_AGG(int_col, true)), arraySortString(ARRAY_AGG(string_col, true)), arraySortString(ARRAY_AGG(json_col)), CARDINALITY(ARRAY_AGG(long_col, true)), CARDINALITY(ARRAY_AGG(double_col)), CARDINALITY(ARRAY_AGG(decimal_col, true)) FROM {tbl}", + "outputs": [[[1, 2, 2, 3], [1, 2, 3], ["a", "b"], ["{\"k\":1}", "{\"k\":2}", "{\"k\":2}", "{\"k\":3}"], 3, 4, 3]] + }, + { + "sql": "SELECT string_col, arraySortInt(ARRAY_AGG(int_col + 1, true)), arraySortInt(ARRAY_AGG(int_mv_col, true)), arraySortString(ARRAY_AGG(string_mv_col)) FROM {tbl} GROUP BY string_col", + "outputs": [["a", [2, 3], [1, 2, 3], ["x", "y", "y"]], ["b", [3, 4], [3, 4, 5], ["x", "z", "z"]]] + }, + { + "sql": "SELECT /*+ aggOptions(is_skip_leaf_stage_aggregate='true') */ string_col, arraySortInt(ARRAY_AGG(int_col + 1, true)), arraySortInt(ARRAY_AGG(int_mv_col, true)), arraySortString(ARRAY_AGG(string_mv_col)) FROM {tbl} GROUP BY string_col", + "outputs": [["a", [2, 3], [1, 2, 3], ["x", "y", "y"]], ["b", [3, 4], [3, 4, 5], ["x", "z", "z"]]] + }, + { + "sql": "SELECT arraySortInt(ARRAY_AGG(int_col, 'INT')), arraySortInt(ARRAY_AGG(int_col, 'INT', true)) FROM {tbl}", + "outputs": [[[1, 2, 2, 3], [1, 2, 3]]] + } + ] } } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java index bbd582c6cb0..d59c1ea9e6d 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java @@ -191,7 +191,7 @@ public enum AggregationFunctionType { ReturnTypes.ARG1, OperandTypes.VARIADIC, SqlTypeName.OTHER, SqlTypeName.BIGINT), // Array aggregate functions - ARRAYAGG("arrayAgg", new ArrayOfComponentReturnTypeInference(), OperandTypes.VARIADIC, SqlTypeName.OTHER), + ARRAYAGG("arrayAgg", new ArrayAggReturnTypeInference(), OperandTypes.VARIADIC, SqlTypeName.OTHER), LISTAGG("listAgg", SqlTypeName.OTHER, SqlTypeName.VARCHAR), SUMARRAYLONG("sumArrayLong", new ArrayReturnTypeInference(SqlTypeName.BIGINT), OperandTypes.ARRAY, SqlTypeName.OTHER), @@ -464,16 +464,28 @@ public enum AggregationFunctionType { /// Returns ARRAY of the component type of the first operand when the first operand is an ARRAY. /// Falls back to ARRAY of the operand type if the component type is unavailable. - private static class ArrayOfComponentReturnTypeInference implements SqlReturnTypeInference { + /// Without the data type argument, i.e. ARRAY_AGG(expression[, isDistinct]), the element type is limited to the types + /// of the explicit form. + private static class ArrayAggReturnTypeInference implements SqlReturnTypeInference { + private static final Set<SqlTypeName> INFERRED_ELEMENT_TYPES = Set.of(SqlTypeName.BOOLEAN, SqlTypeName.TINYINT, + SqlTypeName.SMALLINT, SqlTypeName.INTEGER, SqlTypeName.BIGINT, SqlTypeName.REAL, SqlTypeName.FLOAT, + SqlTypeName.DOUBLE, SqlTypeName.DECIMAL, SqlTypeName.CHAR, SqlTypeName.VARCHAR, SqlTypeName.BINARY, + SqlTypeName.VARBINARY, SqlTypeName.TIMESTAMP, SqlTypeName.UUID); + @Override public RelDataType inferReturnType(SqlOperatorBinding opBinding) { RelDataTypeFactory typeFactory = opBinding.getTypeFactory(); RelDataType operandType = opBinding.getOperandType(0); RelDataType componentType = operandType.getComponentType(); - if (componentType != null) { - return typeFactory.createArrayType(componentType, -1); + RelDataType elementType = componentType != null ? componentType : operandType; + int numOperands = opBinding.getOperandCount(); + if ((numOperands == 1 || (numOperands == 2 + && opBinding.getOperandType(1).getSqlTypeName() == SqlTypeName.BOOLEAN)) + && !INFERRED_ELEMENT_TYPES.contains(elementType.getSqlTypeName())) { + throw new IllegalArgumentException(String.format("Unsupported data type for %s(expression): %s", + opBinding.getOperator().getName(), elementType.getSqlTypeName())); } - return typeFactory.createArrayType(operandType, -1); + return typeFactory.createArrayType(elementType, -1); } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 6fc50e1c84d..db4bee7261d 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1291,7 +1291,7 @@ public class CommonConstants { public static final String AGGREGATE_UNION_TRANSPOSE = "AggregateUnionTranspose"; public static final String AGGREGATE_REDUCE_FUNCTIONS = "AggregateReduceFunctions"; public static final String AGGREGATE_FUNCTION_REWRITE = "AggregateFunctionRewrite"; - /// Supplies the data type argument of FIRST_WITH_TIME / LAST_WITH_TIME(dataColumn, timeColumn). + /// Supplies the omitted data type argument of FIRST_WITH_TIME, LAST_WITH_TIME and ARRAY_AGG. public static final String AGGREGATE_TYPE_ARGUMENT_REWRITE = "AggregateTypeArgumentRewrite"; /// Inert unless [CommonConstants.Broker#USE_APPROXIMATE_FUNCTION] is on; name it here to switch it off entirely. public static final String APPROXIMATE_AGGREGATE_REWRITE = "ApproximateAggregateRewrite"; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
