Jackie-Jiang commented on code in PR #19986: URL: https://github.com/apache/pinot/pull/19986#discussion_r4364440737
########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java: ########## @@ -0,0 +1,173 @@ +/** + * 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.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +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.hint.RelHint; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.tools.RelBuilderFactory; +import org.apache.pinot.segment.spi.AggregationFunctionType; + + +/// Supplies the explicit data type argument from the type of the data argument: +/// - FIRST_WITH_TIME(dataColumn, timeColumn) -> FIRST_WITH_TIME(dataColumn, timeColumn, 'dataType') +/// - LAST_WITH_TIME(dataColumn, timeColumn) -> LAST_WITH_TIME(dataColumn, timeColumn, 'dataType') +/// +/// The result type of the call is unchanged. The literal is projected directly beneath the aggregate, appended to the +/// existing project when there is one, because [PinotAggregateExchangeNodeInsertRule] and `AggregatePushdownRule` +/// inline literal arguments only from the immediate project. +/// +/// The servers require the data type argument, so the rule is required rather than an optimization: `QueryEnvironment` +/// always applies it, and `skipPlannerRules` cannot turn it off. +public class PinotAggregateTypeArgumentRewriteRule extends RelOptRule { + public static final PinotAggregateTypeArgumentRewriteRule INSTANCE = + new PinotAggregateTypeArgumentRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY); + + // Matches with RelOptRule operands, like the other aggregate rewrite rules, instead of a RelRule config. + @SuppressWarnings("deprecation") + private PinotAggregateTypeArgumentRewriteRule(RelBuilderFactory factory) { + super(operand(LogicalAggregate.class, any()), factory, null); + } + + @Override + public boolean matches(RelOptRuleCall call) { + // Requiring a call without its data type is also what makes the rule terminate. + Aggregate aggRel = call.rel(0); + RelNode input = aggRel.getInput(); + return aggRel.getAggCallList().stream().anyMatch(aggCall -> getDataTypeArgument(aggCall, input) != null); + } + + @Override + public void onMatch(RelOptRuleCall call) { + Aggregate aggRel = call.rel(0); + RelNode input = aggRel.getInput(); + RelNode unboxedInput = PinotRuleUtils.unboxRel(input); + RexBuilder rexBuilder = aggRel.getCluster().getRexBuilder(); + + List<RexNode> projects; + RelNode projectInput; + List<RelHint> hints; + if (unboxedInput instanceof Project) { + Project project = (Project) unboxedInput; + projects = new ArrayList<>(project.getProjects()); + projectInput = project.getInput(); + hints = project.getHints(); + } else { + int numFields = input.getRowType().getFieldCount(); + projects = new ArrayList<>(numFields + 1); + for (int i = 0; i < numFields; i++) { + projects.add(rexBuilder.makeInputRef(input, i)); + } + projectInput = input; + hints = List.of(); + } + + // Appending keeps the existing field indexes, so the group set and the other calls are unaffected. + List<AggregateCall> orgAggCalls = aggRel.getAggCallList(); + int numAggCalls = orgAggCalls.size(); + int[] literalIndexes = new int[numAggCalls]; Review Comment: Can we build the rewritten aggregate calls in the same loop that appends the type literals, then create the `LogicalProject` afterward? That would eliminate `literalIndexes`, the `-1` sentinel, and the second loop while preserving literal deduplication and the original aggregate result types. ########## pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java: ########## @@ -426,6 +427,46 @@ public static AggregationFunctionType getAggregationFunctionType(String function } } + /// Accepts `(dataColumn, timeColumn)` and `(dataColumn, timeColumn, 'dataType')`. The two-argument form is limited to + /// the types of the explicit form, since its data type argument is inferred from the data argument. + /// + /// `INFERRED_DATA_TYPES` matches the types that `PinotAggregateTypeArgumentRewriteRule` maps to a data type argument, + /// and the single-stage `AggregateFunctionRewriteOptimizer` supports the corresponding schema types. + private static class ValueWithTimeOperandTypeChecker implements SqlOperandTypeChecker { + private static final SqlOperandTypeChecker DELEGATE = OperandTypes.or( + OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.ANY), + OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)); + private static final Set<SqlTypeName> INFERRED_DATA_TYPES = Set.of(SqlTypeName.TINYINT, SqlTypeName.SMALLINT, Review Comment: Can we remove `INFERRED_DATA_TYPES`? The rewrite rule already maps supported types and throws for unsupported ones, so this duplicates its type restrictions. Failing during planning seems sufficient here. Without the whitelist, we can also remove `ValueWithTimeOperandTypeChecker` and use `OperandTypes.or(...)` directly. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
