mihaibudiu commented on code in PR #5074: URL: https://github.com/apache/calcite/pull/5074#discussion_r3521715537
########## core/src/test/java/org/apache/calcite/adapter/enumerable/EnumerableCustomAggregateTest.java: ########## @@ -0,0 +1,227 @@ +/* + * 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.jdbc.JavaTypeFactoryImpl; +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.RelOptCluster; +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.RelCollationTraitDef; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.RelFactories; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +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.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.util.SqlOperatorTables; +import org.apache.calcite.tools.FrameworkConfig; +import org.apache.calcite.tools.Frameworks; +import org.apache.calcite.tools.Planner; +import org.apache.calcite.tools.Program; +import org.apache.calcite.tools.Programs; +import org.apache.calcite.tools.RelBuilder; + +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); + + /** 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; + } + }; + } + + /** Builds {@code SELECT g, MY_CUSTOM_AGG(x) FROM (VALUES ...) GROUP BY g} on a + * planner whose context carries {@code implementorTable} (or the built-in + * table when null), and optimizes it to the enumerable convention. */ + private static RelNode planAggregate( Review Comment: This function should have a more precise name, like createAndOptimizeTestProgram. ########## core/src/test/java/org/apache/calcite/adapter/enumerable/EnumerableCustomAggregateTest.java: ########## @@ -0,0 +1,227 @@ +/* + * 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.jdbc.JavaTypeFactoryImpl; +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.RelOptCluster; +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.RelCollationTraitDef; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.RelFactories; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +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.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.util.SqlOperatorTables; +import org.apache.calcite.tools.FrameworkConfig; +import org.apache.calcite.tools.Frameworks; +import org.apache.calcite.tools.Planner; +import org.apache.calcite.tools.Program; +import org.apache.calcite.tools.Programs; +import org.apache.calcite.tools.RelBuilder; + +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); + + /** 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; + } + }; + } + + /** Builds {@code SELECT g, MY_CUSTOM_AGG(x) FROM (VALUES ...) GROUP BY g} on a Review Comment: Maybe you can show the full query here ########## core/src/test/java/org/apache/calcite/adapter/enumerable/EnumerableCustomAggregateTest.java: ########## @@ -0,0 +1,227 @@ +/* + * 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.jdbc.JavaTypeFactoryImpl; +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.RelOptCluster; +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.RelCollationTraitDef; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.RelFactories; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +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.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.util.SqlOperatorTables; +import org.apache.calcite.tools.FrameworkConfig; +import org.apache.calcite.tools.Frameworks; +import org.apache.calcite.tools.Planner; +import org.apache.calcite.tools.Program; +import org.apache.calcite.tools.Programs; +import org.apache.calcite.tools.RelBuilder; + +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); + + /** 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; + } + }; + } + + /** Builds {@code SELECT g, MY_CUSTOM_AGG(x) FROM (VALUES ...) GROUP BY g} on a + * planner whose context carries {@code implementorTable} (or the built-in + * table when null), and optimizes it to the enumerable convention. */ + private static RelNode planAggregate( + @Nullable RexImplementorTable implementorTable) { + final VolcanoPlanner planner = + new VolcanoPlanner(implementorTable == null + ? Contexts.empty() : Contexts.of(implementorTable)); + planner.addRelTraitDef(ConventionTraitDef.INSTANCE); + planner.addRelTraitDef(RelCollationTraitDef.INSTANCE); + for (RelOptRule rule : EnumerableRules.rules()) { + planner.addRule(rule); + } + for (RelOptRule rule : RelOptRules.CALC_RULES) { + planner.addRule(rule); + } + planner.removeRule(EnumerableRules.ENUMERABLE_PROJECT_RULE); + + final RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + final RelOptCluster cluster = + RelOptCluster.create(planner, new RexBuilder(typeFactory)); + final RelBuilder builder = + RelFactories.LOGICAL_BUILDER.create(cluster, null); + final RelNode logical = builder + .values(new String[]{"g", "x"}, "a", 1, "a", 2, "b", 3) + .aggregate(builder.groupKey("g"), + builder.aggregateCall(MY_AGG, builder.field("x")).as("c")) + .build(); + final RelNode enumRoot = + planner.changeTraits(logical, + cluster.traitSet().replace(EnumerableConvention.INSTANCE)); + planner.setRoot(enumRoot); + return planner.findBestExp(); + } + + /** 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() { + final RexImplementorTable table = + RexImplementorTables.chain(myAggTable(), RexImpTable.instance()); + final RelNode physical = planAggregate(table); + assertThat(RelOptUtil.toString(physical), + containsString("EnumerableAggregate")); + + final Bindable bindable = + EnumerableInterpretable.toBindable(Collections.emptyMap(), + CalcitePrepare.Dummy.getSparkHandler(false), + (EnumerableRel) 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()); + } + } + // group "a" has two rows, group "b" has one. + assertThat(actual, is(ImmutableMap.of("a", 2L, "b", 1L))); + } + + /** Registers the custom aggregate (operator table) and its implementor + * (planner context) through the {@link Frameworks} embedding API and runs a + * SQL query end-to-end, checking the actual result values. */ + @Test void customAggregateExecutesFromSql() throws Exception { + final List<RelOptRule> rules = new ArrayList<>(EnumerableRules.rules()); + rules.addAll(RelOptRules.CALC_RULES); + rules.remove(EnumerableRules.ENUMERABLE_PROJECT_RULE); + final Program program = Programs.ofRules(rules); + final FrameworkConfig config = Frameworks.newConfigBuilder() + .parserConfig(SqlParser.Config.DEFAULT) + .defaultSchema(Frameworks.createRootSchema(true)) + .operatorTable( + SqlOperatorTables.chain(SqlStdOperatorTable.instance(), + SqlOperatorTables.of(MY_AGG))) + .context( + Contexts.of( + RexImplementorTables.chain(myAggTable(), RexImpTable.instance()))) + .programs(program) + .build(); + + final Planner planner = Frameworks.getPlanner(config); + final SqlNode parse = planner.parse("SELECT g, MY_CUSTOM_AGG(x) AS c\n" Review Comment: why here use the parser and above build the plan by hand? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
