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
The following commit(s) were added to
refs/heads/xiangfu0/agg-rewrite/02-array-agg by this push:
new 1a5b828a80a Limit early inference to ARRAY_AGG and preserve union
column semantics
1a5b828a80a is described below
commit 1a5b828a80a0a7682d5711c65e6e1a00732f131d
Author: Xiang Fu <[email protected]>
AuthorDate: Tue Oct 27 13:37:02 2026 -0700
Limit early inference to ARRAY_AGG and preserve union column semantics
Keep MODE CASE-to-FILTER conversion intact. Append aggregate literals above
unions and lift an existing column only when all branches project the same
literal. Reject DISTINCT ARRAY_AGG explicitly in SSE instead of silently
dropping DISTINCT.
---
.../apache/pinot/sql/parsers/CalciteSqlParser.java | 4 +++
.../AggregateFunctionRewriteOptimizerTest.java | 9 +++++
.../PinotAggregateTypeArgumentRewriteRule.java | 18 +++++++---
.../pinot/calcite/rel/rules/PinotRuleUtils.java | 38 ++++++++++++----------
.../org/apache/pinot/query/QueryEnvironment.java | 2 +-
.../query/ApproximateAggregateRewriteTest.java | 12 +++++++
.../src/test/resources/queries/AggregatePlans.json | 19 ++++++-----
.../src/test/resources/queries/UDFAggregates.json | 5 +++
8 files changed, 75 insertions(+), 32 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
index 274b6be95cf..9039445b360 100644
---
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
+++
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
@@ -1242,6 +1242,10 @@ public class CalciteSqlParser {
private static Expression compileFunctionExpression(SqlBasicCall
functionNode) {
SqlKind functionKind = functionNode.getKind();
+ if (functionKind == SqlKind.ARRAY_AGG &&
functionNode.getFunctionQuantifier() != null
+ && "DISTINCT".equals(functionNode.getFunctionQuantifier().toString()))
{
+ throw new SqlCompilationException("Function 'ARRAY_AGG' on DISTINCT is
not supported.");
+ }
boolean negated = false;
String canonicalName;
switch (functionKind) {
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 fc9828c9f33..3e324ce47f3 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
@@ -25,6 +25,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.exception.BadQueryRequestException;
import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.apache.pinot.sql.parsers.SqlCompilationException;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
@@ -109,6 +110,14 @@ public class AggregateFunctionRewriteOptimizerTest {
+ "FROM testTable");
}
+ @Test
+ public void testArrayAggRejectsDistinctKeyword() {
+ for (String arguments : new String[]{"intCol", "intCol, false", "intCol,
'INT'"}) {
+ expectThrows(SqlCompilationException.class,
+ () -> CalciteSqlParser.compileToPinotQuery("SELECT
ARRAY_AGG(DISTINCT " + arguments + ") FROM testTable"));
+ }
+ }
+
@Test
public void testArrayAggRejectsUnsupportedColumn() {
PinotQuery query = CalciteSqlParser.compileToPinotQuery("SELECT
ARRAY_AGG(mapCol) FROM testTable");
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 332f92ced98..324d10fb5c9 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
@@ -64,20 +64,27 @@ public class PinotAggregateTypeArgumentRewriteRule extends
RelOptRule {
ColumnDataType.BYTES, ColumnDataType.UUID);
public static final PinotAggregateTypeArgumentRewriteRule INSTANCE =
- new PinotAggregateTypeArgumentRewriteRule(true,
PinotRuleUtils.PINOT_REL_FACTORY, null);
+ new PinotAggregateTypeArgumentRewriteRule(true, false,
PinotRuleUtils.PINOT_REL_FACTORY, null);
/// Leaves ANY_VALUE without the data type argument.
public static final PinotAggregateTypeArgumentRewriteRule WITHOUT_ANY_VALUE =
- new PinotAggregateTypeArgumentRewriteRule(false,
PinotRuleUtils.PINOT_REL_FACTORY,
+ new PinotAggregateTypeArgumentRewriteRule(false, false,
PinotRuleUtils.PINOT_REL_FACTORY,
"PinotAggregateTypeArgumentRewriteRuleWithoutAnyValue");
+ /// Runs before CASE-to-FILTER without changing other aggregation functions.
+ public static final PinotAggregateTypeArgumentRewriteRule ARRAY_AGG_ONLY =
+ new PinotAggregateTypeArgumentRewriteRule(false, true,
PinotRuleUtils.PINOT_REL_FACTORY,
+ "PinotAggregateTypeArgumentRewriteRuleArrayAggOnly");
+
+ private final boolean _arrayAggOnly;
private final boolean _rewriteAnyValue;
// Matches with RelOptRule operands, like the other aggregate rewrite rules,
instead of a RelRule config.
@SuppressWarnings("deprecation")
- private PinotAggregateTypeArgumentRewriteRule(boolean rewriteAnyValue,
RelBuilderFactory factory,
- @Nullable String description) {
+ private PinotAggregateTypeArgumentRewriteRule(boolean rewriteAnyValue,
boolean arrayAggOnly,
+ RelBuilderFactory factory, @Nullable String description) {
super(operand(LogicalAggregate.class, any()), factory, description);
_rewriteAnyValue = rewriteAnyValue;
+ _arrayAggOnly = arrayAggOnly;
}
@Override
@@ -142,6 +149,9 @@ public class PinotAggregateTypeArgumentRewriteRule extends
RelOptRule {
}
return List.of(getDataTypeArgument(name, dataType.getSqlTypeName()));
}
+ if (_arrayAggOnly) {
+ return null;
+ }
if (name.equals(AggregationFunctionType.MODE.name())) {
// BOOLEAN and TIMESTAMP already return their own type, cast from the
DOUBLE result.
return argList.size() == 1 &&
SqlTypeName.CHAR_TYPES.contains(getSqlTypeName(input, argList.get(0)))
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java
index 2658d210d51..7091c3f7cd9 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java
@@ -89,22 +89,10 @@ public class PinotRuleUtils {
}
}
- /// Appends aggregate literal arguments without hiding literals already
projected beneath a union. If a union has
- /// no projected literals, keeps one project above it; otherwise recurses
through nested unions to their projects.
+ /// Appends literal arguments in one project, retaining a union column as a
literal only when every branch projects
+ /// that same literal. Branch-specific constants remain input references.
public static RelNode appendAggregateLiterals(RelNode input, List<RexNode>
literals) {
- return appendAggregateLiterals(input, literals, false);
- }
-
- private static RelNode appendAggregateLiterals(RelNode input, List<RexNode>
literals, boolean insideUnion) {
RelNode rel = unboxRel(input);
- if (rel instanceof Union && (insideUnion || hasProjectedLiterals(rel))) {
- Union union = (Union) rel;
- List<RelNode> inputs = new ArrayList<>(union.getInputs().size());
- for (RelNode unionInput : union.getInputs()) {
- inputs.add(appendAggregateLiterals(unionInput, literals, true));
- }
- return union.copy(union.getTraitSet(), inputs);
- }
if (rel instanceof Project) {
Project project = (Project) rel;
List<RexNode> projects = new ArrayList<>(project.getProjects());
@@ -114,18 +102,32 @@ public class PinotRuleUtils {
RexBuilder rexBuilder = input.getCluster().getRexBuilder();
List<RexNode> projects = new
ArrayList<>(input.getRowType().getFieldCount() + literals.size());
for (int i = 0; i < input.getRowType().getFieldCount(); i++) {
- projects.add(rexBuilder.makeInputRef(input, i));
+ RexLiteral literal = rel instanceof Union ? commonLiteral(rel, i) : null;
+ projects.add(literal != null ? literal : rexBuilder.makeInputRef(input,
i));
}
projects.addAll(literals);
return LogicalProject.create(input, List.of(), projects, (List<String>)
null, Set.of());
}
- private static boolean hasProjectedLiterals(RelNode input) {
+ @Nullable
+ private static RexLiteral commonLiteral(RelNode input, int index) {
RelNode rel = unboxRel(input);
if (rel instanceof Project) {
- return ((Project) rel).getProjects().stream().anyMatch(project ->
project instanceof RexLiteral);
+ RexNode expression = ((Project) rel).getProjects().get(index);
+ return expression instanceof RexLiteral ? (RexLiteral) expression : null;
+ }
+ if (rel instanceof Union) {
+ RexLiteral common = null;
+ for (RelNode branch : rel.getInputs()) {
+ RexLiteral literal = commonLiteral(branch, index);
+ if (literal == null || (common != null && !common.equals(literal))) {
+ return null;
+ }
+ common = literal;
+ }
+ return common;
}
- return rel instanceof Union &&
rel.getInputs().stream().anyMatch(PinotRuleUtils::hasProjectedLiterals);
+ return null;
}
public static boolean isExchange(RelNode rel) {
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
index 44501f7d439..95e9b4f102e 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
@@ -613,7 +613,7 @@ public class QueryEnvironment {
// Supplies the data type argument before AggregateCaseToFilter, which
rewrites only one-argument calls: its
// ARRAY_AGG(expression) FILTER (WHERE condition) form of ARRAY_AGG(CASE
WHEN ...) fails at runtime.
if
(relOptRule.toString().equals(CommonConstants.Broker.PlannerRuleNames.AGGREGATE_CASE_TO_FILTER))
{
-
hepProgramBuilder.addRuleInstance(PinotAggregateTypeArgumentRewriteRule.WITHOUT_ANY_VALUE);
+
hepProgramBuilder.addRuleInstance(PinotAggregateTypeArgumentRewriteRule.ARRAY_AGG_ONLY);
}
hepProgramBuilder.addRuleInstance(relOptRule);
}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
index cea95cc294d..f580c196b6f 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
@@ -130,6 +130,18 @@ public class ApproximateAggregateRewriteTest extends
QueryEnvironmentTestBase {
}
}
+ @Test
+ public void testApproximateParamsKeepBranchSpecificUnionColumns() {
+ QueryEnvironment env = newQueryEnvironment(true, "threshold=17",
"threshold=23");
+ String query = "EXPLAIN PLAN FOR SELECT DISTINCTCOUNT(x), COVARPOP(x, y)
FROM "
+ + "(SELECT col3 x, 1.0 y FROM a UNION ALL SELECT col3, CAST(col7 AS
DOUBLE) FROM b)";
+ String plan = env.compile(query,
CalciteSqlParser.compileToSqlNodeAndOptions(query))
+ .explain(RANDOM_REQUEST_ID_GEN.nextLong(), null).getExplainPlan();
+ String leaf = plan.lines().filter(line ->
line.contains("aggType=[LEAF]")).findFirst().orElse("");
+ assertTrue(leaf.contains("_UTF-8'threshold=17'"), plan);
+ assertFalse(leaf.contains("1.0"), plan);
+ }
+
/// `AggregateCaseToFilterRule` turns `COUNT(DISTINCT CASE WHEN c THEN x
END)` into `COUNT(DISTINCT x) FILTER (WHERE
/// c)`, so the rows that fail `c` add nothing to the count. It only matches
a single-argument call, so the rewrite
/// has to run after it: with the parameters appended first, the CASE stays
and its NULL counts as one more value.
diff --git a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
index fe7e4caf100..e8442d6d3a0 100644
--- a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
+++ b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
@@ -374,20 +374,21 @@
]
},
{
- "description": "The data type argument goes into every input of a
union, next to the other literal arguments",
+ "description": "The data type argument stays above a union alongside
literals common to every branch",
"sql": "EXPLAIN PLAN FOR SELECT ARRAY_AGG(x), PERCENTILE(x, 50),
ARRAY_AGG(DISTINCT x, 'INT') FROM (SELECT a.col3 x FROM a UNION ALL SELECT
b.col3 FROM b)",
"output": [
"Execution Plan",
"\nPinotLogicalAggregate(group=[{}], agg#0=[ARRAYAGG($0,
_UTF-8'INT')], agg#1=[PERCENTILE($1, 50)], agg#2=[ARRAYAGG($2, _UTF-8'INT',
true)], aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash])",
"\n PinotLogicalAggregate(group=[{}], agg#0=[ARRAYAGG($0,
_UTF-8'INT')], agg#1=[PERCENTILE($0, 50)], agg#2=[ARRAYAGG($0, _UTF-8'INT',
true)], aggType=[LEAF])",
- "\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[single])",
- "\n LogicalProject($f0=[$2], $f1=[50], $f2=[_UTF-8'INT'],
$f3=[_UTF-8'INT'])",
- "\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[single])",
- "\n LogicalProject($f0=[$2], $f1=[50], $f2=[_UTF-8'INT'],
$f3=[_UTF-8'INT'])",
- "\n PinotLogicalTableScan(table=[[default, b]])",
+ "\n LogicalProject($f0=[$0], $f1=[50], $f2=[_UTF-8'INT'],
$f3=[_UTF-8'INT'])",
+ "\n LogicalUnion(all=[true])",
+ "\n PinotLogicalExchange(distribution=[single])",
+ "\n LogicalProject(x=[$2], $f1=[50], $f2=[_UTF-8'INT'])",
+ "\n PinotLogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[single])",
+ "\n LogicalProject(x=[$2], $f1=[50], $f2=[_UTF-8'INT'])",
+ "\n PinotLogicalTableScan(table=[[default, b]])",
"\n"
]
},
@@ -399,7 +400,7 @@
"\nPinotLogicalAggregate(group=[{}], agg#0=[FIRSTWITHTIME($0, $0,
_UTF-8'INT')], agg#1=[ARRAYAGG($1, _UTF-8'INT')], aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash])",
"\n PinotLogicalAggregate(group=[{}], agg#0=[FIRSTWITHTIME($0,
$1, _UTF-8'INT')], agg#1=[ARRAYAGG($0, _UTF-8'INT')], aggType=[LEAF])",
- "\n LogicalProject($f0=[$0], $f1=[$1], $f2=[_UTF-8'INT'])",
+ "\n LogicalProject($f0=[$0], $f1=[$1], $f2=[_UTF-8'INT'],
$f3=[_UTF-8'INT'])",
"\n LogicalUnion(all=[true])",
"\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col3=[$2], ts=[$7])",
diff --git a/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json
b/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json
index c405e55daa2..f1ca915df48 100644
--- a/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json
+++ b/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json
@@ -45,6 +45,11 @@
}
},
"queries": [
+ {"sql": "SELECT MODE(CASE WHEN int_col > 3 THEN string_col END) FROM
{tbl}", "outputs": [["c"]]},
+ {"sql": "SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */
bool_col, MODE(CASE WHEN int_col > 3 THEN string_col END) FROM {tbl} GROUP BY
bool_col", "outputs": [[true, "c"], [false, "c"]]},
+ {"sql": "SET inferAnyValueType=true; SELECT src, ANY_VALUE(x) FROM
(SELECT 'a' src, 1 x FROM {tbl} UNION ALL SELECT 'b', 2 FROM {tbl}) GROUP BY
src", "outputs": [["a", 1], ["b", 2]]},
+ {"sql": "SELECT LASTWITHTIME(x, t) FROM (SELECT int_col x,
CAST('2019-01-01 00:00:00' AS TIMESTAMP) t FROM {tbl} UNION ALL SELECT int_col,
timestamp_col FROM {tbl})", "outputs": [[101]]},
+ {"sql": "SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */
src, k, MODE(s) FROM (SELECT 'a' src, int_col k, string_col s FROM {tbl} UNION
ALL SELECT 'b', int_col, string_col FROM {tbl}) WHERE k = 2 GROUP BY src, k",
"outputs": [["a", 2, "a"], ["b", 2, "a"]]},
{
"sql": "SELECT FIRSTWITHTIME(int_col, timestamp_col, 'INT'),
FIRSTWITHTIME(string_col, timestamp_col, 'STRING'), first_with_time(bool_col,
CAST(long_col AS TIMESTAMP), 'BOOLEAN') FROM {tbl}",
"outputs": [[2, "a", true]]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
