gortiz commented on code in PR #19986:
URL: https://github.com/apache/pinot/pull/19986#discussion_r4336926293
##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java:
##########
@@ -426,6 +428,26 @@ public static AggregationFunctionType
getAggregationFunctionType(String function
}
}
+ /// Returns the type of the data argument. The two-argument form is limited
to the types of the explicit form.
+ private static class ValueWithTimeReturnTypeInference implements
SqlReturnTypeInference {
+ private static final Set<SqlTypeName> INFERRED_DATA_TYPES =
Set.of(SqlTypeName.TINYINT, SqlTypeName.SMALLINT,
+ SqlTypeName.INTEGER, SqlTypeName.BIGINT, SqlTypeName.REAL,
SqlTypeName.FLOAT, SqlTypeName.DOUBLE,
+ SqlTypeName.DECIMAL, SqlTypeName.CHAR, SqlTypeName.VARCHAR,
SqlTypeName.BOOLEAN, SqlTypeName.TIMESTAMP,
+ SqlTypeName.BINARY, SqlTypeName.VARBINARY, SqlTypeName.UUID);
+
+ @Override
+ public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
+ if (opBinding.getOperandCount() == 2) {
+ SqlTypeName dataType = opBinding.getOperandType(0).getSqlTypeName();
+ if (!INFERRED_DATA_TYPES.contains(dataType)) {
+ throw new IllegalArgumentException(String.format("Unsupported data
type for %s(dataColumn, timeColumn): %s",
Review Comment:
Throwing an IllegalArgumentException from the return type inference works,
but it skips Calcite's validation path. The user gets no position info, and the
error comes back as a generic planning failure instead of a validation error.
Can you move this check into the operand type checker? A custom
SqlOperandTypeChecker that wraps the current `OperandTypes.or(...)` can call
`callBinding.newValidationSignatureError()` for the 2-arg form when the data
argument type is not supported. Then the return type inference stays plain
`ARG0`. Please also make the plan test assert the error type, not only the
message.
##########
pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java:
##########
@@ -0,0 +1,171 @@
+/**
+ * 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.
+public class PinotAggregateTypeArgumentRewriteRule extends RelOptRule {
+ public static final PinotAggregateTypeArgumentRewriteRule INSTANCE =
+ new
PinotAggregateTypeArgumentRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY, null);
+
+ public static PinotAggregateTypeArgumentRewriteRule
instanceWithDescription(String description) {
+ return new
PinotAggregateTypeArgumentRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY,
description);
+ }
+
+ // Matches with RelOptRule operands, like the other aggregate rewrite rules,
instead of a RelRule config.
+ @SuppressWarnings("deprecation")
+ private PinotAggregateTypeArgumentRewriteRule(RelBuilderFactory factory,
@Nullable String description) {
+ super(operand(LogicalAggregate.class, any()), factory, description);
+ }
+
+ @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];
+ Map<String, Integer> literalIndexByDataType = new HashMap<>();
+ for (int i = 0; i < numAggCalls; i++) {
+ String dataType = getDataTypeArgument(orgAggCalls.get(i), input);
+ literalIndexes[i] = dataType == null ? -1 :
literalIndexByDataType.computeIfAbsent(dataType, type -> {
+ projects.add(rexBuilder.makeLiteral(type));
+ return projects.size() - 1;
+ });
+ }
+ RelNode newInput = LogicalProject.create(projectInput, hints, projects,
(List<String>) null, Set.of());
+
+ List<AggregateCall> aggCalls = new ArrayList<>(numAggCalls);
+ for (int i = 0; i < numAggCalls; i++) {
+ AggregateCall aggCall = orgAggCalls.get(i);
+ if (literalIndexes[i] < 0) {
+ aggCalls.add(aggCall);
+ continue;
+ }
+ List<Integer> argList = new ArrayList<>(aggCall.getArgList());
+ argList.add(literalIndexes[i]);
+ // Keeps the original result type, so the rewrite does not change the
result schema.
+ aggCalls.add(aggCall.withArgList(argList));
+ }
+ call.transformTo(
+ aggRel.copy(aggRel.getTraitSet(), newInput, aggRel.getGroupSet(),
aggRel.getGroupSets(), aggCalls));
+ }
+
+ /// Returns the data type argument to append, or `null` when the call
already has one or is not such a call.
+ @Nullable
+ private static String getDataTypeArgument(AggregateCall aggCall, RelNode
input) {
+ if (aggCall.getArgList().size() != 2) {
+ return null;
+ }
+ String name =
AggregationFunctionType.getNormalizedAggregationFunctionName(aggCall.getAggregation().getName());
+ if (!name.equals(AggregationFunctionType.FIRSTWITHTIME.name())
+ && !name.equals(AggregationFunctionType.LASTWITHTIME.name())) {
+ return null;
+ }
+ SqlTypeName dataType =
+
input.getRowType().getFieldList().get(aggCall.getArgList().get(0)).getType().getSqlTypeName();
+ switch (dataType) {
+ case TINYINT:
Review Comment:
Small question on the narrow integer types. For TINYINT and SMALLINT the
call type stays `ARG0`, but the literal is `'INT'`, so the function produces an
INT. `RelToPlanNodeConverter` maps both to INT, so I think this is consistent
at runtime. Can you add a runtime case with `CAST(x AS SMALLINT)` to lock that
in?
##########
pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java:
##########
@@ -175,6 +175,10 @@ private PinotQueryRuleSets() {
PinotAggregateFunctionRewriteRule
.instanceWithDescription(PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE),
+ // Must stay ahead of PinotAggregateExchangeNodeInsertRule
(POST_LOGICAL), which inlines the literal argument.
+ PinotAggregateTypeArgumentRewriteRule
+
.instanceWithDescription(PlannerRuleNames.AGGREGATE_TYPE_ARGUMENT_REWRITE),
Review Comment:
This rule is needed for correctness, not only for optimization. Without it,
every 2-arg FIRST_WITH_TIME / LAST_WITH_TIME reaches the leaf unchanged and
fails with "expects 3 arguments". The rule runs in an opt-program phase and has
a `PlannerRuleNames` entry, so `skipPlannerRules=AggregateTypeArgumentRewrite`
turns it off. Can we make it non-skippable? I see two options. One: register it
like the POST_LOGICAL rules, which `skipPlannerRules` does not filter. Two: add
the literal in `PinotAggregateExchangeNodeInsertRule`, which already inlines
the literal arguments. Either way, please add a test that runs with the rule
skipped, to document what happens.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java:
##########
@@ -46,62 +53,83 @@ public void optimize(PinotQuery pinotQuery, @Nullable
Schema schema) {
return;
}
- // Only perform auto rewrite when enabled through query option.
- if (pinotQuery.getQueryOptions() == null ||
!Boolean.parseBoolean(pinotQuery.getQueryOptions().get(
-
CommonConstants.Broker.Request.QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE))) {
- return;
- }
-
- List<Expression> selectList = pinotQuery.getSelectList();
- if (selectList != null) {
- for (Expression expression : selectList) {
- maybeRewriteAggregateFunction(expression, schema);
- }
- }
-
- List<Expression> groupByList = pinotQuery.getGroupByList();
- if (groupByList != null) {
- for (Expression expression : groupByList) {
- maybeRewriteAggregateFunction(expression, schema);
- }
- }
+ boolean rewriteToTypedVariants = pinotQuery.getQueryOptions() != null &&
Boolean.parseBoolean(
+
pinotQuery.getQueryOptions().get(CommonConstants.Broker.Request.QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE));
+ rewrite(pinotQuery.getSelectList(), schema, rewriteToTypedVariants);
+ rewrite(pinotQuery.getGroupByList(), schema, rewriteToTypedVariants);
+ rewrite(pinotQuery.getOrderByList(), schema, rewriteToTypedVariants);
+ rewrite(pinotQuery.getFilterExpression(), schema, rewriteToTypedVariants);
+ rewrite(pinotQuery.getHavingExpression(), schema, rewriteToTypedVariants);
+ }
- List<Expression> orderByList = pinotQuery.getOrderByList();
- if (orderByList != null) {
- for (Expression expression : orderByList) {
- maybeRewriteAggregateFunction(expression, schema);
+ private static void rewrite(@Nullable List<Expression> expressions, Schema
schema, boolean rewriteToTypedVariants) {
+ if (expressions != null) {
+ for (Expression expression : expressions) {
+ rewrite(expression, schema, rewriteToTypedVariants);
}
}
-
- maybeRewriteAggregateFunction(pinotQuery.getFilterExpression(), schema);
- maybeRewriteAggregateFunction(pinotQuery.getHavingExpression(), schema);
}
- private void maybeRewriteAggregateFunction(@Nullable Expression expression,
Schema schema) {
+ private static void rewrite(@Nullable Expression expression, Schema schema,
boolean rewriteToTypedVariants) {
if (expression == null || !expression.isSetFunctionCall()) {
return;
}
Function function = expression.getFunctionCall();
String functionName = function.getOperator();
if (!AggregationFunctionType.isAggregationFunction(functionName)) {
+ rewrite(function.getOperands(), schema, rewriteToTypedVariants);
Review Comment:
This recursion also changes what `autoRewriteAggregationType` does today.
Before, the optimizer rewrote only top-level calls, so `MIN(longCol) AS m`,
`MAX(longCol) + 1` and `HAVING MIN(longCol) > 5` kept MIN/MAX. With this change
they become MINLONG/MAXLONG, and the result type changes from DOUBLE to LONG. I
think the new behavior is right: before, the select list and HAVING could end
up as two different aggregations of the same expression. But users who already
set the option will see the change, so please mention it in the PR description.
Please also add a focused test for it. Today
`testTypedVariantsRequireQueryOption` covers it only through the `AS m` alias.
--
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]
