This is an automated email from the ASF dual-hosted git repository.
gortiz 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 a6d16ef3772 Fix ClassCastException in skip-leaf direct aggregation
with deduplicated agg calls (#19023)
a6d16ef3772 is described below
commit a6d16ef377265b060635407e71e872fe3376aef6
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Tue Jul 21 12:01:46 2026 +0200
Fix ClassCastException in skip-leaf direct aggregation with deduplicated
agg calls (#19023)
---
.../PinotAggregateExchangeNodeInsertRule.java | 38 +++++--
.../PinotAggregateExchangeNodeInsertRuleTest.java | 121 +++++++++++++++++++++
2 files changed, 152 insertions(+), 7 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 fb3d0f437c4..96f9c422668 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
@@ -160,7 +160,7 @@ public class PinotAggregateExchangeNodeInsertRule {
return;
}
- PinotLogicalAggregate newAggRel = createPlan(call, aggRel, true,
hintOptions, newCollations, limit);
+ RelNode newAggRel = createPlan(call, aggRel, true, hintOptions,
newCollations, limit);
RelNode newProjectRel = projectRel.copy(projectRel.getTraitSet(),
List.of(newAggRel));
call.transformTo(sortRel.copy(sortRel.getTraitSet(),
List.of(newProjectRel)));
}
@@ -202,7 +202,7 @@ public class PinotAggregateExchangeNodeInsertRule {
return;
}
- PinotLogicalAggregate newAggRel = createPlan(call, aggRel, true,
hintOptions, collations, limit);
+ RelNode newAggRel = createPlan(call, aggRel, true, hintOptions,
collations, limit);
call.transformTo(sortRel.copy(sortRel.getTraitSet(),
List.of(newAggRel)));
}
}
@@ -227,7 +227,7 @@ public class PinotAggregateExchangeNodeInsertRule {
}
}
- private static PinotLogicalAggregate createPlan(RelOptRuleCall call,
Aggregate aggRel, boolean hasGroupBy,
+ private static RelNode createPlan(RelOptRuleCall call, Aggregate aggRel,
boolean hasGroupBy,
Map<String, String> hintOptions, @Nullable List<RelFieldCollation>
collations, int limit) {
// WITHIN GROUP collation is not supported in leaf stage aggregation.
RelCollation withinGroupCollation = extractWithinGroupCollation(aggRel);
@@ -261,12 +261,28 @@ public class PinotAggregateExchangeNodeInsertRule {
* Use this group by optimization to skip leaf stage aggregation when
aggregating at leaf level is not desired. Many
* situation could be wasted effort to do group-by on leaf, eg: when
cardinality of group by column is very high.
*/
- private static PinotLogicalAggregate
createPlanWithExchangeDirectAggregation(RelOptRuleCall call, Aggregate aggRel,
+ private static RelNode
createPlanWithExchangeDirectAggregation(RelOptRuleCall call, Aggregate aggRel,
@Nullable RelCollation withinGroupCollation, @Nullable
List<RelFieldCollation> collations, int limit) {
RelNode input = aggRel.getInput();
+ // A layout-restoring Project that RelBuilder may add over the rebuilt
aggregate (see below); when
+ // present it is re-applied over the direct aggregate so the aggregate's
output columns are
+ // unchanged for the parent.
+ Project topProject = null;
// Create Project when there's none below the aggregate.
if (!(PinotRuleUtils.unboxRel(input) instanceof Project)) {
- aggRel = (Aggregate) generateProjectUnderAggregate(call, aggRel);
+ RelNode rewritten =
PinotRuleUtils.unboxRel(generateProjectUnderAggregate(call, aggRel));
+ if (rewritten instanceof Project) {
+ // generateProjectUnderAggregate rebuilds the aggregate with
RelBuilder, which deduplicates
+ // identical aggregate calls and, when it does, adds a top Project to
restore the original
+ // output layout (a column referenced more than once — e.g. a SUM
reused by an AVG reduction —
+ // is projected twice). Keep that Project; the aggregate is its input,
and we re-apply it over
+ // the direct aggregate below. Without this, casting the Project to
Aggregate threw
+ // "LogicalProject cannot be cast to Aggregate".
+ topProject = (Project) rewritten;
+ aggRel = (Aggregate) PinotRuleUtils.unboxRel(topProject.getInput());
+ } else {
+ aggRel = (Aggregate) rewritten;
+ }
input = aggRel.getInput();
}
@@ -280,8 +296,16 @@ public class PinotAggregateExchangeNodeInsertRule {
exchange = PinotLogicalExchange.create(input, distribution);
}
- return new PinotLogicalAggregate(aggRel, exchange, buildAggCalls(aggRel,
AggType.DIRECT, false), AggType.DIRECT,
- false, collations, limit);
+ PinotLogicalAggregate directAggregate = new PinotLogicalAggregate(aggRel,
exchange,
+ buildAggCalls(aggRel, AggType.DIRECT, false), AggType.DIRECT, false,
collations, limit);
+ if (topProject == null) {
+ return directAggregate;
+ }
+ // Re-apply the layout-restoring Project over the direct aggregate (same
projects and row type,
+ // new input). The direct aggregate keeps the rebuilt aggregate's output
columns, so the Project's
+ // references remain valid.
+ return topProject.copy(topProject.getTraitSet(), directAggregate,
topProject.getProjects(),
+ topProject.getRowType());
}
/**
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRuleTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRuleTest.java
new file mode 100644
index 00000000000..9f3dc888fac
--- /dev/null
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRuleTest.java
@@ -0,0 +1,121 @@
+/**
+ * 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 com.google.common.collect.ImmutableList;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.plan.hep.HepPlanner;
+import org.apache.calcite.plan.hep.HepProgramBuilder;
+import org.apache.calcite.plan.hep.HepRelVertex;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.rel.hint.HintPredicates;
+import org.apache.calcite.rel.hint.HintStrategyTable;
+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.logical.LogicalValues;
+import org.apache.calcite.rel.metadata.DefaultRelMetadataProvider;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalAggregate;
+import org.apache.pinot.query.type.TypeFactory;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// Unit tests for {@link PinotAggregateExchangeNodeInsertRule}, focused on
the skip-leaf-stage
+/// direct-aggregation path.
+public class PinotAggregateExchangeNodeInsertRuleTest {
+ private static final TypeFactory TYPE_FACTORY = TypeFactory.INSTANCE;
+ private static final RexBuilder REX_BUILDER = RexBuilder.DEFAULT;
+
+ private static AggregateCall count(String name) {
+ return AggregateCall.create(SqlStdOperatorTable.COUNT, false, List.of(1),
-1,
+ TYPE_FACTORY.createSqlType(SqlTypeName.BIGINT), name);
+ }
+
+ /// Builds `Aggregate[hint=is_skip_leaf_stage_group_by]( <count()>*n ) over
a two-column Values`,
+ /// runs the WithoutSort rule, and returns the (unwrapped) transformed root.
+ private static RelNode runSkipLeaf(int numCountCalls) {
+ HepProgramBuilder program = new HepProgramBuilder();
+
program.addRuleInstance(PinotAggregateExchangeNodeInsertRule.WithoutSort.INSTANCE);
+ HepPlanner planner = new HepPlanner(program.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+ cluster.setHintStrategies(HintStrategyTable.builder()
+ .hintStrategy(PinotHintOptions.AGGREGATE_HINT_OPTIONS,
HintPredicates.AGGREGATE)
+ .build());
+ RelDataType rowType = TYPE_FACTORY.builder()
+ .add("g", SqlTypeName.INTEGER)
+ .add("v", SqlTypeName.INTEGER)
+ .build();
+ LogicalValues input = new LogicalValues(cluster,
RelTraitSet.createEmpty(), rowType, ImmutableList.of());
+ List<AggregateCall> aggCalls = new ArrayList<>();
+ for (int i = 0; i < numCountCalls; i++) {
+ aggCalls.add(count("c" + i));
+ }
+ LogicalAggregate aggregate = (LogicalAggregate) LogicalAggregate.create(
+ input, ImmutableBitSet.of(0), List.of(ImmutableBitSet.of(0)),
aggCalls)
+
.withHints(List.of(RelHint.builder(PinotHintOptions.AGGREGATE_HINT_OPTIONS)
+
.hintOption(PinotHintOptions.AggregateOptions.IS_SKIP_LEAF_STAGE_GROUP_BY,
"true")
+ .build()));
+ planner.setRoot(aggregate);
+ return unwrap(planner.findBestExp());
+ }
+
+ private static RelNode unwrap(RelNode node) {
+ return node instanceof HepRelVertex ? ((HepRelVertex)
node).getCurrentRel() : node;
+ }
+
+ /// Regression: two identical `COUNT` calls make
`generateProjectUnderAggregate` deduplicate them
+ /// and return a `Project`-over-`Aggregate` (the top `Project` re-references
the deduplicated column
+ /// to restore the original layout). The rule used to cast that result to
`Aggregate` and throw
+ /// `LogicalProject cannot be cast to Aggregate`. It must now preserve the
`Project` over a direct
+ /// aggregate.
+ @Test
+ public void skipLeafWithDuplicateAggregateCallsPreservesLayoutProject() {
+ RelNode result = runSkipLeaf(2);
+ assertTrue(result instanceof LogicalProject,
+ "expected a layout-restoring Project at the root, got " +
result.getClass().getSimpleName());
+ RelNode child = unwrap(result.getInput(0));
+ assertTrue(child instanceof PinotLogicalAggregate,
+ "expected a PinotLogicalAggregate under the Project, got " +
child.getClass().getSimpleName());
+ // group key + the two (deduplicated then re-expanded) count columns.
+ assertEquals(result.getRowType().getFieldCount(), 3);
+ }
+
+ /// The common case (no duplicate agg calls) still produces a bare direct
aggregate — no extra
+ /// Project.
+ @Test
+ public void skipLeafWithoutDuplicatesProducesDirectAggregate() {
+ RelNode result = runSkipLeaf(1);
+ assertTrue(result instanceof PinotLogicalAggregate,
+ "expected a PinotLogicalAggregate, got " +
result.getClass().getSimpleName());
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
