You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/10/21 01:11:07 UTC
(pinot) 01/01: Support ARRAY_AGG without a data type
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 <xi...@apache.org>
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: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org