This is an automated email from the ASF dual-hosted git repository.
jackie 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 e80d95fed6 Handle Agg Functions with Literal Args When Used with a
Union (#13561)
e80d95fed6 is described below
commit e80d95fed6db890d34bac7b35a8d575971bdec65
Author: Ankit Sultana <ankitsult...@uber.com>
AuthorDate: Tue Jul 9 20:07:37 2024 -0500
Handle Agg Functions with Literal Args When Used with a Union (#13561)
---
.../PinotAggregateExchangeNodeInsertRule.java | 20 ++++++++++++++----
.../src/test/resources/queries/AggregatePlans.json | 24 ++++++++++++++++++++++
2 files changed, 40 insertions(+), 4 deletions(-)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java
index 58d423842a..8a9f179179 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java
@@ -31,6 +31,7 @@ import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Aggregate;
import org.apache.calcite.rel.core.AggregateCall;
import org.apache.calcite.rel.core.Project;
+import org.apache.calcite.rel.core.Union;
import org.apache.calcite.rel.hint.RelHint;
import org.apache.calcite.rel.logical.LogicalAggregate;
import org.apache.calcite.rel.rules.AggregateExtractProjectRule;
@@ -219,8 +220,8 @@ public class PinotAggregateExchangeNodeInsertRule extends
RelOptRule {
private static PinotLogicalAggregate
convertAggFromIntermediateInput(RelOptRuleCall call,
PinotLogicalExchange exchange, AggType aggType) {
Aggregate aggRel = call.rel(0);
- RelNode input = PinotRuleUtils.unboxRel(aggRel.getInput());
- List<RexNode> projects = (input instanceof Project) ? ((Project)
input).getProjects() : null;
+ RelNode input = aggRel.getInput();
+ List<RexNode> projects = findImmediateProjects(input);
// Create new AggregateCalls from exchange input. Exchange produces
results with group keys followed by intermediate
// aggregate results.
@@ -258,8 +259,8 @@ public class PinotAggregateExchangeNodeInsertRule extends
RelOptRule {
}
private static List<AggregateCall> buildAggCalls(Aggregate aggRel, AggType
aggType) {
- RelNode input = PinotRuleUtils.unboxRel(aggRel.getInput());
- List<RexNode> projects = (input instanceof Project) ? ((Project)
input).getProjects() : null;
+ RelNode input = aggRel.getInput();
+ List<RexNode> projects = findImmediateProjects(input);
List<AggregateCall> orgAggCalls = aggRel.getAggCallList();
List<AggregateCall> aggCalls = new ArrayList<>(orgAggCalls.size());
for (AggregateCall orgAggCall : orgAggCalls) {
@@ -330,4 +331,15 @@ public class PinotAggregateExchangeNodeInsertRule extends
RelOptRule {
ImmutableList.of(), aggType.isInputIntermediateFormat() ? -1 :
orgAggCall.filterArg, orgAggCall.distinctKeys,
orgAggCall.collation, numGroups, input, null, null);
}
+
+ @Nullable
+ private static List<RexNode> findImmediateProjects(RelNode relNode) {
+ relNode = PinotRuleUtils.unboxRel(relNode);
+ if (relNode instanceof Project) {
+ return ((Project) relNode).getProjects();
+ } else if (relNode instanceof Union) {
+ return findImmediateProjects(relNode.getInput(0));
+ }
+ return null;
+ }
}
diff --git a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
index 84f1f6533a..02c91a6490 100644
--- a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
+++ b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
@@ -133,6 +133,30 @@
"\n LogicalTableScan(table=[[default, a]])",
"\n"
]
+ },
+ {
+ "description": "Select aggregates with literals on top of a union",
+ "sql": "EXPLAIN PLAN FOR with teamOne as (select /*+
aggOptions(is_skip_leaf_stage_group_by='true') */ col2, percentile(col3, 50) as
sum_of_runs from a group by col2), teamTwo as (select /*+
aggOptions(is_skip_leaf_stage_group_by='true') */ col2, percentile(col3, 50) as
sum_of_runs from a group by col2), all as (select col2, sum_of_runs from
teamOne union all select col2, sum_of_runs from teamTwo) select /*+
aggOption(is_skip_leaf_stage_group_by='true') */ col2, percentile(sum_of [...]
+ "output": [
+ "Execution Plan",
+ "\nPinotLogicalAggregate(group=[{0}], agg#0=[PERCENTILE($1, 50)])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n PinotLogicalAggregate(group=[{0}], agg#0=[PERCENTILE($1,
50)])",
+ "\n LogicalUnion(all=[true])",
+ "\n PinotLogicalExchange(distribution=[hash[0, 1, 2]])",
+ "\n LogicalProject(col2=[$0], sum_of_runs=[$1], $f2=[50])",
+ "\n PinotLogicalAggregate(group=[{0}],
agg#0=[PERCENTILE($1, 50)])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col2=[$1], col3=[$2], $f2=[50])",
+ "\n LogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[hash[0, 1, 2]])",
+ "\n LogicalProject(col2=[$0], sum_of_runs=[$1], $f2=[50])",
+ "\n PinotLogicalAggregate(group=[{0}],
agg#0=[PERCENTILE($1, 50)])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col2=[$1], col3=[$2], $f2=[50])",
+ "\n LogicalTableScan(table=[[default, a]])",
+ "\n"
+ ]
}
]
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org
