This is an automated email from the ASF dual-hosted git repository.
yashmayya 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 f8352be252f Preserve aggregate hints through ProjectAggregateMergeRule
(#19324)
f8352be252f is described below
commit f8352be252fc1a6cf099480a7812d1555d264e0e
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Aug 20 17:59:59 2026 -0400
Preserve aggregate hints through ProjectAggregateMergeRule (#19324)
---
.../rel/rules/PinotProjectAggregateMergeRule.java | 135 +++++++++++++++++++++
.../calcite/rel/rules/PinotQueryRuleSets.java | 7 +-
.../apache/pinot/query/QueryCompilationTest.java | 49 ++++++++
3 files changed, 188 insertions(+), 3 deletions(-)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotProjectAggregateMergeRule.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotProjectAggregateMergeRule.java
new file mode 100644
index 00000000000..f76bc899b38
--- /dev/null
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotProjectAggregateMergeRule.java
@@ -0,0 +1,135 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.calcite.rel.rules;
+
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.calcite.plan.RelHintsPropagator;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelShuttleImpl;
+import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.hint.RelHint;
+import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.rules.ProjectAggregateMergeRule;
+
+
+/// Pinot customized version of [ProjectAggregateMergeRule] that preserves the
matched [Aggregate]'s hints.
+///
+/// [ProjectAggregateMergeRule] rebuilds the aggregate from scratch with a
`RelBuilder`, so the rebuilt node
+/// starts out with no hints. Calcite normally repairs this automatically:
[RelOptRuleCall#transformTo(RelNode)]
+/// defaults to propagating the hints of `rels[0]` (the node the rule matched
on) into the new sub-tree. That
+/// repair does not apply here, because `rels[0]` is the `Project`, and
`aggOptions` hints are attached to the
+/// `Aggregate` (see `HintPredicates.AGGREGATE` in
+/// [org.apache.pinot.calcite.rel.hint.PinotHintStrategyTable]). The
aggregate's hints are therefore dropped
+/// silently, and every `aggOptions` option on it is lost — including
`is_partitioned_by_group_by_keys`,
+/// `is_skip_leaf_stage_group_by`, `is_leaf_return_final_result` and the
group-trim options.
+///
+/// The user-visible symptom is a colocation hint that stops working as soon
as a `Project` lands directly above
+/// the aggregate. The most common way that happens is a `SUM` over a
*nullable* argument:
+/// [PinotAggregateReduceFunctionsRule] rewrites it to `$SUM0(x) + COUNT(x)`
plus a
+/// `CASE(COUNT(x) = 0, NULL, $SUM0(x))` project, which is exactly the pattern
this rule matches. With a
+/// non-nullable argument the reduction collapses to a bare `$SUM0` with no
project, no match, and the hint
+/// survives — which is why the problem only shows up on some queries.
+///
+/// This rule reuses Calcite's transformation verbatim (no forked rule body to
drift out of sync) and only
+/// re-attaches the hints afterwards. The rule never changes the aggregate's
group set, so its hints remain
+/// valid for the rebuilt node.
+public class PinotProjectAggregateMergeRule extends ProjectAggregateMergeRule {
+
+ public static PinotProjectAggregateMergeRule instanceWithDescription(String
description) {
+ return new PinotProjectAggregateMergeRule((Config)
Config.DEFAULT.withDescription(description));
+ }
+
+ private PinotProjectAggregateMergeRule(Config config) {
+ super(config);
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ Aggregate aggregate = call.rel(1);
+ if (aggregate.getHints().isEmpty()) {
+ super.onMatch(call);
+ } else {
+ super.onMatch(new HintPreservingRuleCall(call, aggregate));
+ }
+ }
+
+ /// Delegating [RelOptRuleCall] that re-attaches the matched aggregate's
hints to the rebuilt aggregate before
+ /// handing the rule's output to the real call.
+ ///
+ /// The hints are copied verbatim rather than through
+ /// [org.apache.calcite.plan.RelOptUtil#propagateRelHints(RelNode,
RelNode)]: that helper appends the child
+ /// index to each hint's [RelHint#inheritPath] as it descends, so
re-propagating on every rule application grows
+ /// the inherit path without bound. The rebuilt node then never compares
equal to the previous one, the rule
+ /// keeps re-firing on its own output, and planning dies with a
`StackOverflowError`. Copying the hint list
+ /// unchanged makes the rewrite a fixpoint, so the planner terminates after
it has been applied once.
+ private static class HintPreservingRuleCall extends RelOptRuleCall {
+ private final RelOptRuleCall _delegate;
+ private final Aggregate _aggregate;
+
+ HintPreservingRuleCall(RelOptRuleCall delegate, Aggregate aggregate) {
+ super(delegate.getPlanner(), delegate.getOperand0(), delegate.getRels(),
Map.of(), delegate.getParents());
+ _delegate = delegate;
+ _aggregate = aggregate;
+ }
+
+ @Override
+ public void transformTo(RelNode rel, Map<RelNode, RelNode> equiv,
RelHintsPropagator handler) {
+ _delegate.transformTo(rel.accept(new
RestoreAggregateHintsShuttle(_aggregate.getHints())), equiv, handler);
+ }
+
+ @Nullable
+ @Override
+ public List<RelNode> getChildRels(RelNode rel) {
+ return _delegate.getChildRels(rel);
+ }
+ }
+
+ /// Attaches the given hints to the topmost [Aggregate] of the tree it
visits, which is the aggregate
+ /// [ProjectAggregateMergeRule] rebuilt. Nothing below it is touched: the
rule reuses the original aggregate's
+ /// input unchanged, so the only hintable node it recreates is the aggregate
itself.
+ private static class RestoreAggregateHintsShuttle extends RelShuttleImpl {
+ private final List<RelHint> _hints;
+ private boolean _restored;
+
+ RestoreAggregateHintsShuttle(List<RelHint> hints) {
+ _hints = hints;
+ }
+
+ @Override
+ public RelNode visit(LogicalAggregate aggregate) {
+ return restoreHints(aggregate);
+ }
+
+ @Override
+ public RelNode visit(RelNode other) {
+ return other instanceof Aggregate ? restoreHints((Aggregate) other) :
super.visit(other);
+ }
+
+ private RelNode restoreHints(Aggregate aggregate) {
+ if (_restored) {
+ return aggregate;
+ }
+ _restored = true;
+ return aggregate.withHints(_hints);
+ }
+ }
+}
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
index e3f54b0bd58..38822d4dfad 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
@@ -32,7 +32,6 @@ import org.apache.calcite.rel.rules.FilterMergeRule;
import org.apache.calcite.rel.rules.FilterProjectTransposeRule;
import org.apache.calcite.rel.rules.FilterSetOpTransposeRule;
import org.apache.calcite.rel.rules.JoinPushExpressionsRule;
-import org.apache.calcite.rel.rules.ProjectAggregateMergeRule;
import org.apache.calcite.rel.rules.ProjectFilterTransposeRule;
import org.apache.calcite.rel.rules.ProjectMergeRule;
import org.apache.calcite.rel.rules.ProjectRemoveRule;
@@ -238,8 +237,10 @@ public class PinotQueryRuleSets {
UnionMergeRule.Config.DEFAULT
.withDescription(PlannerRuleNames.UNION_MERGE).toRule(),
// Drop unused aggregate calls when a Project on top of the Aggregate
doesn't reference them. Default-on.
- ProjectAggregateMergeRule.Config.DEFAULT
- .withDescription(PlannerRuleNames.PROJECT_AGGREGATE_MERGE).toRule(),
+ // Pinot fork of Calcite's ProjectAggregateMergeRule: the stock rule
rebuilds the aggregate and silently
+ // drops its aggOptions hints, because Calcite only restores the hints
of the node the rule matched on
+ // (the Project). See PinotProjectAggregateMergeRule.
+
PinotProjectAggregateMergeRule.instanceWithDescription(PlannerRuleNames.PROJECT_AGGREGATE_MERGE),
PruneEmptyRules.CorrelateLeftEmptyRuleConfig.DEFAULT
.withDescription(PlannerRuleNames.PRUNE_EMPTY_CORRELATE_LEFT).toRule(),
PruneEmptyRules.CorrelateRightEmptyRuleConfig.DEFAULT
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
index 54e916982c5..66391119edd 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
@@ -772,6 +772,55 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
}
}
+ /// `ProjectAggregateMergeRule` rebuilds the aggregate with a `RelBuilder`,
and Calcite's automatic hint
+ /// propagation only restores the hints of the node the rule matched on --
the `Project`, which never carries
+ /// `aggOptions`. Without
[org.apache.pinot.calcite.rel.rules.PinotProjectAggregateMergeRule] the
aggregate's
+ /// hints are dropped and the aggregate is split into LEAF + exchange +
FINAL despite the colocation hint.
+ ///
+ /// A `Project` lands directly above the aggregate here because the `SUM`
argument is nullable:
+ /// `PinotAggregateReduceFunctionsRule` rewrites `SUM(x)` into `$SUM0(x) +
COUNT(x)` plus a
+ /// `CASE(COUNT(x) = 0, NULL, $SUM0(x))` project. With a non-nullable
argument the reduction collapses to a
+ /// bare `$SUM0` with no project, the rule does not match, and the hint
survives either way.
+ @Test
+ public void testAggregateHintSurvivesProjectAggregateMerge() {
+ String query = "EXPLAIN PLAN FOR SELECT /*+
aggOptions(is_partitioned_by_group_by_keys='true') */ "
+ + "col1, SUM(CASE WHEN col3 > 5 THEN col3 ELSE NULL END) FROM b GROUP
BY col1";
+ String explain = _queryEnvironment.explainQuery(query,
RANDOM_REQUEST_ID_GEN.nextLong());
+ assertTrue(explain.contains("aggType=[DIRECT]"),
+ "is_partitioned_by_group_by_keys should produce a DIRECT aggregate,
but got:\n" + explain);
+ assertFalse(explain.contains("PinotLogicalExchange"),
+ "A colocated aggregate must not have an exchange below it, but got:\n"
+ explain);
+ }
+
+ /// Same defect reached through a window function, which is how it shows up
in practice: `LAG` is nullable, so
+ /// any expression derived from it is nullable, so the `SUM` over it takes
the reduction path above. The window
+ /// itself is colocated by `windowOptions`; only the aggregate above it used
to lose its hint.
+ @Test
+ public void testAggregateHintSurvivesProjectAggregateMergeAboveWindow() {
+ String query = "EXPLAIN PLAN FOR WITH w AS ("
+ + "SELECT /*+ windowOptions(is_partitioned_by_window_keys='true') */
col1, col3, ts, "
+ + "LAG(col3, 1) OVER (PARTITION BY col1 ORDER BY ts) AS prev FROM b) "
+ + "SELECT /*+ aggOptions(is_partitioned_by_group_by_keys='true') */ "
+ + "col1, SUM(CASE WHEN prev IS NULL THEN 0 ELSE col3 - prev END) FROM
w GROUP BY col1";
+ String explain = _queryEnvironment.explainQuery(query,
RANDOM_REQUEST_ID_GEN.nextLong());
+ assertTrue(explain.contains("aggType=[DIRECT]"),
+ "is_partitioned_by_group_by_keys should produce a DIRECT aggregate,
but got:\n" + explain);
+ }
+
+ /// The other `aggOptions` options travel on the same hint and were lost the
same way.
+ /// `is_skip_leaf_stage_group_by` must still push the aggregate above the
exchange.
+ @Test
+ public void testSkipLeafStageGroupByHintSurvivesProjectAggregateMerge() {
+ String query = "EXPLAIN PLAN FOR SELECT /*+
aggOptions(is_skip_leaf_stage_group_by='true') */ "
+ + "col1, SUM(CASE WHEN col3 > 5 THEN col3 ELSE NULL END) FROM b GROUP
BY col1";
+ String explain = _queryEnvironment.explainQuery(query,
RANDOM_REQUEST_ID_GEN.nextLong());
+ assertTrue(explain.contains("aggType=[DIRECT]"),
+ "is_skip_leaf_stage_group_by should produce a single DIRECT aggregate
above the exchange, but got:\n"
+ + explain);
+ assertFalse(explain.contains("aggType=[LEAF]"),
+ "is_skip_leaf_stage_group_by must not leave a LEAF aggregate, but
got:\n" + explain);
+ }
+
@Test
public void testQueryWithHint() {
// Hinting the query to use final stage aggregation makes server directly
return final result
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
