This is an automated email from the ASF dual-hosted git repository.
mihaibudiu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/calcite.git
The following commit(s) were added to refs/heads/main by this push:
new 9d5e1b28c9 [CALCITE-7640] Enumerable engine should execute aggregates
whose implementor comes from a custom RexImplementorTable
9d5e1b28c9 is described below
commit 9d5e1b28c96b3c87da17f4b0075cd39dd2e1bfca
Author: Darpan Lunagariya <[email protected]>
AuthorDate: Fri Jul 3 21:53:06 2026 +0530
[CALCITE-7640] Enumerable engine should execute aggregates whose
implementor comes from a custom RexImplementorTable
CALCITE-7631 made RexImplementorTable composable for scalar code
generation and constant reduction, but the aggregate path still resolves
implementors through the RexImpTable singleton: EnumerableAggregate's
constructor rejects any aggregate the built-ins do not know, and
AggImpState generates code against the singleton. A custom aggregate
implementor supplied through a chained table is therefore never used.
Resolve aggregate implementors through the injected table, defaulting to
the built-ins. Add RexImplementorTables.of(RelOptCluster) (reads the
planner context, else the built-in table), move the availability check
from the EnumerableAggregate constructor into EnumerableAggregateRule
(keeping only the table-independent structural checks in the node), and
thread the resolved table into AggImpState and its callers. Behaviour is
unchanged when no table is registered.
---
.../calcite/adapter/enumerable/AggImpState.java | 14 +-
.../adapter/enumerable/EnumerableAggregate.java | 10 +-
.../enumerable/EnumerableAggregateRule.java | 8 +
.../enumerable/EnumerableSortedAggregate.java | 4 +-
.../adapter/enumerable/EnumerableWindow.java | 4 +-
.../adapter/enumerable/RexImplementorTables.java | 10 +
.../apache/calcite/interpreter/AggregateNode.java | 5 +-
.../enumerable/EnumerableCustomAggregateTest.java | 214 +++++++++++++++++++++
8 files changed, 258 insertions(+), 11 deletions(-)
diff --git
a/core/src/main/java/org/apache/calcite/adapter/enumerable/AggImpState.java
b/core/src/main/java/org/apache/calcite/adapter/enumerable/AggImpState.java
index 0edc7351f0..a225837788 100644
--- a/core/src/main/java/org/apache/calcite/adapter/enumerable/AggImpState.java
+++ b/core/src/main/java/org/apache/calcite/adapter/enumerable/AggImpState.java
@@ -35,10 +35,22 @@ public class AggImpState {
public @MonotonicNonNull List<Expression> state;
public @MonotonicNonNull Expression accumulatorAdder;
+ /**
+ * Creates an AggImpState, resolving the implementor from the built-in table.
+ *
+ * @deprecated Use {@link #AggImpState(int, AggregateCall, boolean,
RexImplementorTable)}.
+ */
+ @Deprecated // to be removed before 2.0
public AggImpState(int aggIdx, AggregateCall call, boolean windowContext) {
+ this(aggIdx, call, windowContext, RexImpTable.INSTANCE);
+ }
+
+ public AggImpState(int aggIdx, AggregateCall call, boolean windowContext,
+ RexImplementorTable implementorTable) {
this.aggIdx = aggIdx;
this.call = call;
- AggImplementor implementor =
RexImpTable.INSTANCE.get(call.getAggregation(), windowContext);
+ AggImplementor implementor =
+ implementorTable.get(call.getAggregation(), windowContext);
if (implementor == null) {
throw new IllegalArgumentException(
"Unable to get aggregate implementation for aggregate "
diff --git
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregate.java
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregate.java
index d38b26c333..0429168a8f 100644
---
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregate.java
+++
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregate.java
@@ -67,12 +67,6 @@ public EnumerableAggregate(
throw new InvalidRelException(
"within-distinct aggregation not supported");
}
- AggImplementor implementor2 =
- RexImpTable.INSTANCE.get(aggCall.getAggregation(), false);
- if (implementor2 == null) {
- throw new InvalidRelException(
- "aggregation " + aggCall.getAggregation() + " not supported");
- }
}
}
@@ -182,8 +176,10 @@ public EnumerableAggregate(RelOptCluster cluster,
RelTraitSet traitSet,
final int groupCount = getGroupCount();
final List<AggImpState> aggs = new ArrayList<>(aggCalls.size());
+ final RexImplementorTable implementorTable =
+ RexImplementorTables.of(getCluster());
for (Ord<AggregateCall> call : Ord.zip(aggCalls)) {
- aggs.add(new AggImpState(call.i, call.e, false));
+ aggs.add(new AggImpState(call.i, call.e, false, implementorTable));
}
// Function0<Object[]> accumulatorInitializer =
diff --git
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregateRule.java
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregateRule.java
index db06fea558..5547408388 100644
---
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregateRule.java
+++
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAggregateRule.java
@@ -22,6 +22,7 @@
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.convert.ConverterRule;
import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.AggregateCall;
import org.apache.calcite.rel.logical.LogicalAggregate;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -48,6 +49,13 @@ protected EnumerableAggregateRule(Config config) {
final Aggregate agg = (Aggregate) rel;
final RelTraitSet traitSet = rel.getCluster()
.traitSet().replace(EnumerableConvention.INSTANCE);
+ final RexImplementorTable implementorTable =
+ RexImplementorTables.of(rel.getCluster());
+ for (AggregateCall aggCall : agg.getAggCallList()) {
+ if (implementorTable.get(aggCall.getAggregation(), false) == null) {
+ return null;
+ }
+ }
try {
return new EnumerableAggregate(
rel.getCluster(),
diff --git
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableSortedAggregate.java
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableSortedAggregate.java
index 73a4525112..29997dd3b7 100644
---
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableSortedAggregate.java
+++
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableSortedAggregate.java
@@ -136,8 +136,10 @@ public EnumerableSortedAggregate(
final int groupCount = getGroupCount();
final List<AggImpState> aggs = new ArrayList<>(aggCalls.size());
+ final RexImplementorTable implementorTable =
+ RexImplementorTables.of(getCluster());
for (Ord<AggregateCall> call : Ord.zip(aggCalls)) {
- aggs.add(new AggImpState(call.i, call.e, false));
+ aggs.add(new AggImpState(call.i, call.e, false, implementorTable));
}
// Function0<Object[]> accumulatorInitializer =
diff --git
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableWindow.java
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableWindow.java
index 78ecce8821..1b60c56126 100644
---
a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableWindow.java
+++
b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableWindow.java
@@ -224,12 +224,14 @@ private static void
sampleOfTheGeneratedWindowedAggregate() {
List<AggImpState> aggs = new ArrayList<>();
List<AggregateCall> aggregateCalls = group.getAggregateCalls(this);
+ final RexImplementorTable implementorTable =
+ RexImplementorTables.of(getCluster());
for (int aggIdx = 0; aggIdx < aggregateCalls.size(); aggIdx++) {
AggregateCall call = aggregateCalls.get(aggIdx);
if (call.ignoreNulls()) {
throw new UnsupportedOperationException("IGNORE NULLS not
supported");
}
- aggs.add(new AggImpState(aggIdx, call, true));
+ aggs.add(new AggImpState(aggIdx, call, true, implementorTable));
}
// The output from this stage is the input plus the aggregate functions.
diff --git
a/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImplementorTables.java
b/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImplementorTables.java
index 22b3340ee8..d5ff234a0e 100644
---
a/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImplementorTables.java
+++
b/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImplementorTables.java
@@ -17,6 +17,7 @@
package org.apache.calcite.adapter.enumerable;
import org.apache.calcite.adapter.enumerable.RexImpTable.RexCallImplementor;
+import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.sql.SqlAggFunction;
import org.apache.calcite.sql.SqlMatchFunction;
import org.apache.calcite.sql.SqlOperator;
@@ -36,6 +37,15 @@ public abstract class RexImplementorTables {
private RexImplementorTables() {
}
+ /** Returns the implementor table registered on the {@code cluster}'s planner
+ * {@link org.apache.calcite.plan.Context}, or the built-in
+ * {@link RexImpTable#instance()} when none is registered. */
+ public static RexImplementorTable of(RelOptCluster cluster) {
+ return cluster.getPlanner().getContext()
+ .maybeUnwrap(RexImplementorTable.class)
+ .orElse(RexImpTable.instance());
+ }
+
/** Creates a table that consults each of the given tables in turn, returning
* the first non-null implementor.
*
diff --git
a/core/src/main/java/org/apache/calcite/interpreter/AggregateNode.java
b/core/src/main/java/org/apache/calcite/interpreter/AggregateNode.java
index 7c908cc0e8..684baad8d4 100644
--- a/core/src/main/java/org/apache/calcite/interpreter/AggregateNode.java
+++ b/core/src/main/java/org/apache/calcite/interpreter/AggregateNode.java
@@ -22,6 +22,7 @@
import org.apache.calcite.adapter.enumerable.JavaRowFormat;
import org.apache.calcite.adapter.enumerable.PhysType;
import org.apache.calcite.adapter.enumerable.PhysTypeImpl;
+import org.apache.calcite.adapter.enumerable.RexImplementorTables;
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
import org.apache.calcite.adapter.enumerable.impl.AggAddContextImpl;
import org.apache.calcite.adapter.java.JavaTypeFactory;
@@ -143,7 +144,9 @@ private AccumulatorFactory getAccumulator(Compiler compiler,
final JavaTypeFactory typeFactory =
(JavaTypeFactory) rel.getCluster().getTypeFactory();
int stateOffset = 0;
- final AggImpState agg = new AggImpState(0, call, false);
+ final AggImpState agg =
+ new AggImpState(0, call, false,
+ RexImplementorTables.of(rel.getCluster()));
int stateSize = requireNonNull(agg.state, "agg.state").size();
final BlockBuilder builder2 = new BlockBuilder();
diff --git
a/core/src/test/java/org/apache/calcite/adapter/enumerable/EnumerableCustomAggregateTest.java
b/core/src/test/java/org/apache/calcite/adapter/enumerable/EnumerableCustomAggregateTest.java
new file mode 100644
index 0000000000..8e68b54757
--- /dev/null
+++
b/core/src/test/java/org/apache/calcite/adapter/enumerable/EnumerableCustomAggregateTest.java
@@ -0,0 +1,214 @@
+/*
+ * 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.calcite.adapter.enumerable;
+
+import org.apache.calcite.DataContexts;
+import org.apache.calcite.adapter.enumerable.RexImpTable.RexCallImplementor;
+import org.apache.calcite.jdbc.CalcitePrepare;
+import org.apache.calcite.linq4j.Enumerable;
+import org.apache.calcite.linq4j.Enumerator;
+import org.apache.calcite.plan.Contexts;
+import org.apache.calcite.plan.ConventionTraitDef;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRules;
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.plan.volcano.VolcanoPlanner;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelRoot;
+import org.apache.calcite.runtime.Bindable;
+import org.apache.calcite.sql.SqlAggFunction;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlMatchFunction;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlWindowTableFunction;
+import org.apache.calcite.sql.fun.SqlBasicAggFunction;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.parser.SqlParser;
+import org.apache.calcite.sql.test.SqlTestFactory;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.util.SqlOperatorTables;
+import org.apache.calcite.sql.validate.SqlValidator;
+import org.apache.calcite.sql2rel.SqlToRelConverter;
+import org.apache.calcite.tools.FrameworkConfig;
+import org.apache.calcite.tools.Frameworks;
+import org.apache.calcite.tools.Planner;
+import org.apache.calcite.tools.Programs;
+
+import com.google.common.collect.ImmutableMap;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Tests that an aggregate implementor supplied by a custom
+ * {@link RexImplementorTable} registered on the planner
+ * {@link org.apache.calcite.plan.Context} is used both to plan and to execute
+ * an {@link EnumerableAggregate}.
+ */
+class EnumerableCustomAggregateTest {
+ /** An aggregate function with no built-in implementor. */
+ private static final SqlAggFunction MY_AGG =
+ SqlBasicAggFunction.create("MY_CUSTOM_AGG", SqlKind.OTHER_FUNCTION,
+ ReturnTypes.BIGINT, OperandTypes.ANY);
+
+ private static final String SQL = "SELECT g, MY_CUSTOM_AGG(x) AS c\n"
+ + "FROM (VALUES ('a', 1), ('a', 2), ('b', 3)) AS t (g, x)\n"
+ + "GROUP BY g";
+
+ /** Returns a table that maps {@link #MY_AGG} to the built-in {@code COUNT}
+ * implementor, so the custom aggregate counts its (non-null) input. */
+ private static RexImplementorTable myAggTable() {
+ final AggImplementor countImplementor =
+ requireNonNull(
+ RexImpTable.instance().get(SqlStdOperatorTable.COUNT, false),
+ "COUNT implementor");
+ return new RexImplementorTable() {
+ @Override public @Nullable RexCallImplementor get(SqlOperator operator) {
+ return null;
+ }
+
+ @Override public @Nullable AggImplementor get(SqlAggFunction aggregation,
+ boolean forWindowAggregate) {
+ return aggregation == MY_AGG ? countImplementor : null;
+ }
+
+ @Override public @Nullable MatchImplementor get(SqlMatchFunction
function) {
+ return null;
+ }
+
+ @Override public @Nullable TableFunctionCallImplementor get(
+ SqlWindowTableFunction operator) {
+ return null;
+ }
+ };
+ }
+
+ /** Plans a {@code GROUP BY} query through the {@link Frameworks} API. */
+ private static RelNode planWithFrameworks(
+ @Nullable RexImplementorTable implementorTable) throws Exception {
+ final List<RelOptRule> rules = new ArrayList<>(EnumerableRules.rules());
+ rules.addAll(RelOptRules.CALC_RULES);
+ rules.remove(EnumerableRules.ENUMERABLE_PROJECT_RULE);
+ Frameworks.ConfigBuilder configBuilder = Frameworks.newConfigBuilder()
+ .parserConfig(SqlParser.Config.DEFAULT)
+ .defaultSchema(Frameworks.createRootSchema(true))
+ .operatorTable(
+ SqlOperatorTables.chain(SqlStdOperatorTable.instance(),
+ SqlOperatorTables.of(MY_AGG)))
+ .programs(Programs.ofRules(rules));
+ if (implementorTable != null) {
+ configBuilder = configBuilder.context(Contexts.of(implementorTable));
+ }
+ final FrameworkConfig config = configBuilder.build();
+
+ final Planner planner = Frameworks.getPlanner(config);
+ final SqlNode parse = planner.parse(SQL);
+ final RelNode logical = planner.rel(planner.validate(parse)).project();
+ final RelTraitSet traitSet =
+ logical.getTraitSet().replace(EnumerableConvention.INSTANCE);
+ return planner.transform(0, traitSet, logical);
+ }
+
+ private static EnumerableRel planDirectly(RexImplementorTable
implementorTable)
+ throws Exception {
+ final VolcanoPlanner planner = new
VolcanoPlanner(Contexts.of(implementorTable));
+ planner.addRelTraitDef(ConventionTraitDef.INSTANCE);
+ for (RelOptRule rule : EnumerableRules.rules()) {
+ planner.addRule(rule);
+ }
+ for (RelOptRule rule : RelOptRules.CALC_RULES) {
+ planner.addRule(rule);
+ }
+
+ final SqlTestFactory factory = SqlTestFactory.INSTANCE
+ .withPlannerFactory(context -> planner)
+ .withOperatorTable(opTab ->
+ SqlOperatorTables.chain(opTab, SqlOperatorTables.of(MY_AGG)));
+ final SqlNode parse = factory.createParser(SQL).parseQuery();
+ final SqlToRelConverter converter = factory.createSqlToRelConverter();
+ final SqlValidator validator = converter.validator;
+ final RelRoot root =
+ converter.convertQuery(validator.validate(parse), false, true);
+ final RelTraitSet traitSet =
+ root.rel.getTraitSet().replace(EnumerableConvention.INSTANCE);
+ final RelNode rel = planner.changeTraits(root.rel, traitSet);
+ planner.setRoot(rel);
+ return (EnumerableRel) planner.findBestExp();
+ }
+
+ private static Map<String, Long> execute(EnumerableRel physical) {
+ final Bindable bindable =
+ EnumerableInterpretable.toBindable(Collections.emptyMap(),
+ CalcitePrepare.Dummy.getSparkHandler(false),
+ physical, EnumerableRel.Prefer.ARRAY);
+ final Enumerable<Object[]> result = bindable.bind(DataContexts.EMPTY);
+ final Map<String, Long> actual = new HashMap<>();
+ try (Enumerator<Object[]> enumerator = result.enumerator()) {
+ while (enumerator.moveNext()) {
+ final Object[] row = enumerator.current();
+ actual.put((String) row[0], ((Number) row[1]).longValue());
+ }
+ }
+ return actual;
+ }
+
+ /** With the custom table on the planner context, the aggregate is planned as
+ * an {@link EnumerableAggregate} and executes to the expected values. */
+ @Test void customAggregateImplementorIsUsedEndToEnd() throws Exception {
+ final RelNode physical =
+ planWithFrameworks(
+ RexImplementorTables.chain(myAggTable(), RexImpTable.instance()));
+ assertThat(RelOptUtil.toString(physical),
+ containsString("EnumerableAggregate"));
+
+ final Map<String, Long> actual = execute((EnumerableRel) physical);
+ assertThat(actual, is(ImmutableMap.of("a", 2L, "b", 1L)));
+ }
+
+ @Test void directPlannerUsesCustomTable() throws Exception {
+ final RexImplementorTable implementors =
+ RexImplementorTables.chain(myAggTable(), RexImpTable.instance());
+ final EnumerableRel physical = planDirectly(implementors);
+
+ assertThat(RelOptUtil.toString(physical),
+ containsString("EnumerableAggregate"));
+ assertThat(execute(physical), is(ImmutableMap.of("a", 2L, "b", 1L)));
+ }
+
+ /** Without the custom table the built-in table has no implementor for the
+ * aggregate, so {@link EnumerableAggregateRule} declines and planning
fails. */
+ @Test void unknownAggregateIsRejectedWithoutCustomTable() {
+ assertThrows(RuntimeException.class,
+ () -> planWithFrameworks(null));
+ }
+}
