Jackie-Jiang commented on code in PR #19489:
URL: https://github.com/apache/pinot/pull/19489#discussion_r3994548230
##########
pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java:
##########
@@ -46,48 +56,69 @@ 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;
- }
+ boolean autoRewrite = pinotQuery.getQueryOptions() != null &&
Boolean.parseBoolean(pinotQuery.getQueryOptions().get(
+
CommonConstants.Broker.Request.QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE));
List<Expression> selectList = pinotQuery.getSelectList();
if (selectList != null) {
for (Expression expression : selectList) {
- maybeRewriteAggregateFunction(expression, schema);
+ maybeRewriteAggregateFunction(expression, schema, autoRewrite);
}
}
List<Expression> groupByList = pinotQuery.getGroupByList();
if (groupByList != null) {
for (Expression expression : groupByList) {
- maybeRewriteAggregateFunction(expression, schema);
+ maybeRewriteAggregateFunction(expression, schema, autoRewrite);
}
}
List<Expression> orderByList = pinotQuery.getOrderByList();
if (orderByList != null) {
for (Expression expression : orderByList) {
- maybeRewriteAggregateFunction(expression, schema);
+ maybeRewriteAggregateFunction(expression, schema, autoRewrite);
}
}
- maybeRewriteAggregateFunction(pinotQuery.getFilterExpression(), schema);
- maybeRewriteAggregateFunction(pinotQuery.getHavingExpression(), schema);
+ maybeRewriteAggregateFunction(pinotQuery.getFilterExpression(), schema,
autoRewrite);
+ maybeRewriteAggregateFunction(pinotQuery.getHavingExpression(), schema,
autoRewrite);
}
- private void maybeRewriteAggregateFunction(@Nullable Expression expression,
Schema schema) {
+ private void maybeRewriteAggregateFunction(@Nullable Expression expression,
Schema schema, boolean autoRewrite) {
if (expression == null || !expression.isSetFunctionCall()) {
return;
}
Function function = expression.getFunctionCall();
+ List<Expression> operands = function.getOperands();
+ for (Expression operand : operands) {
+ // Infer arguments within scalar expressions while preserving the
existing top-level variant rewrites.
+ maybeRewriteAggregateFunction(operand, schema, false);
+ }
String functionName = function.getOperator();
if (!AggregationFunctionType.isAggregationFunction(functionName)) {
return;
}
+ List<String> inferredArguments =
AggregationFunctionType.getAggregationFunctionType(functionName)
+ .inferArguments(operands.size(), i -> {
+ ColumnDataType type = getOperandType(operands.get(i), schema);
+ if (type == null || type.isArray() || type == ColumnDataType.OBJECT)
{
+ return DataType.UNKNOWN;
+ }
+ return type == ColumnDataType.MAP ? DataType.MAP : type.toDataType();
+ });
Review Comment:
**[Critical] Gate automatic timestamp inference during rollout**
Existing `MODE(timestampCol)` queries now gain a third `TIMESTAMP` argument
unconditionally, even when `autoRewriteAggregationType=false`. During a
broker-first upgrade, older servers reject these previously valid queries with
`Mode expects at most 2 arguments, got: 3`. Existing timestamp `AVG` calls also
become invalid on a fully upgraded cluster.
Please gate automatic conversion behind an initially disabled rollout option
and preserve legacy behavior until it is enabled after server upgrades.
Explicit typed calls can remain available for the new behavior.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java:
##########
@@ -133,4 +164,96 @@ private void maybeRewriteAggregateFunction(@Nullable
Expression expression, Sche
}
}
}
+
+ @Nullable
+ private static ColumnDataType getOperandType(Expression operand, Schema
schema) {
+ if (operand.isSetIdentifier()) {
+ FieldSpec fieldSpec =
schema.getFieldSpecFor(operand.getIdentifier().getName());
+ return fieldSpec != null
+ ? ColumnDataType.fromDataType(fieldSpec.getDataType(),
fieldSpec.isSingleValueField())
+ : null;
+ }
+ if (operand.isSetLiteral()) {
+ return
RequestUtils.getLiteralTypeAndValue(operand.getLiteral()).getLeft();
+ }
+ if (!operand.isSetFunctionCall()) {
+ return null;
+ }
+ Function function = operand.getFunctionCall();
+ List<Expression> arguments = function.getOperands();
+ String name = FunctionRegistry.canonicalize(function.getOperator());
+ switch (name) {
+ case "cast":
+ return literalType(arguments, 1);
+ case "jsonextractscalar":
+ case "jsonextractscalarfast":
+ case "jsonextractscalarfirstmatch":
+ case "jsonextractscalarfory":
+ return literalType(arguments, 2);
+ case "case":
+ // CASE stores alternating condition/result pairs followed by an
optional ELSE result.
+ ColumnDataType resultType = ColumnDataType.UNKNOWN;
+ for (int i = 1; i < arguments.size(); i += 2) {
+ resultType = commonType(resultType, getOperandType(arguments.get(i),
schema));
+ }
+ if (arguments.size() % 2 == 1) {
+ resultType = commonType(resultType,
getOperandType(arguments.get(arguments.size() - 1), schema));
+ }
+ return resultType;
+ case "datetimeconvert":
+ if (arguments.size() < 3 || !arguments.get(2).isSetLiteral()
+ || !arguments.get(2).getLiteral().isSetStringValue()) {
+ return null;
+ }
+ DateTimeFieldSpec.TimeFormat format =
+ new
DateTimeFormatSpec(arguments.get(2).getLiteral().getStringValue()).getTimeFormat();
+ return format == DateTimeFieldSpec.TimeFormat.EPOCH || format ==
DateTimeFieldSpec.TimeFormat.TIMESTAMP
+ ? ColumnDataType.LONG
+ : ColumnDataType.STRING;
+ default:
+ ColumnDataType[] argumentTypes = new ColumnDataType[arguments.size()];
+ for (int i = 0; i < arguments.size(); i++) {
+ argumentTypes[i] = getOperandType(arguments.get(i), schema);
+ if (argumentTypes[i] == null) {
+ return null;
+ }
+ }
+ FunctionInfo functionInfo = FunctionRegistry.lookupFunctionInfo(name,
argumentTypes);
+ return functionInfo != null ?
FunctionUtils.getColumnDataType(functionInfo.getMethod().getReturnType()) :
null;
Review Comment:
**[Major] Infer COALESCE from its native transform**
`MODE(COALESCE(stringCol, ''))` falls back to the registry function's
`Object` return type, although the native single-stage
`CoalesceTransformFunction` returns `STRING` for these operands. MODE therefore
remains configured as `DOUBLE` and cannot serialize or finalize its string
frequency map.
Please infer polymorphic native functions from their operands using the
execution implementation's typing rules. A targeted optimizer test expecting
the appended `'MIN', 'STRING'` arguments fails for this expression.
##########
pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunctionTest.java:
##########
@@ -40,6 +51,58 @@ Object[] scenarios() {
};
}
+ @DataProvider
+ Object[] typedScenarios() {
+ return new Object[]{new Scenario(DataType.STRING, true), new
Scenario(DataType.STRING, false),
+ new Scenario(DataType.TIMESTAMP, true), new
Scenario(DataType.TIMESTAMP, false)};
+ }
+
+ @Test(dataProvider = "typedScenarios")
+ void typedModeMergesCountsAndHandlesNulls(Scenario scenario) {
+ String type = scenario._dataType.name();
+ String first = scenario._dataType == DataType.STRING ? "apple" :
"2026-09-03 10:11:12.123";
+ String shared = scenario._dataType == DataType.STRING ? "banana" :
"2026-09-04 10:11:12.456";
+ String last = scenario._dataType == DataType.STRING ? "cherry" :
"2026-09-05 10:11:12.789";
+ // Each instance has a different local mode; the shared value wins only
after merging full counts.
+ scenario.getDeclaringTable(true)
Review Comment:
**[Critical coverage gap] Test disabled null handling**
Every new typed execution test enables null handling; the scenario booleans
vary dictionary encoding. Please add `STRING`/`TIMESTAMP` cases with null
handling disabled, covering mixed-null, all-null, empty, and grouped inputs
where stored defaults participate in MODE. A data provider over type,
dictionary encoding, and null-handling mode would cover the relevant
combinations.
The repository's testing policy classifies missing coverage of either
null-handling mode as critical. This is a coverage finding; no disabled-mode
runtime failure was demonstrated.
##########
pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunctionTest.java:
##########
@@ -40,6 +51,58 @@ Object[] scenarios() {
};
}
+ @DataProvider
+ Object[] typedScenarios() {
+ return new Object[]{new Scenario(DataType.STRING, true), new
Scenario(DataType.STRING, false),
+ new Scenario(DataType.TIMESTAMP, true), new
Scenario(DataType.TIMESTAMP, false)};
+ }
+
+ @Test(dataProvider = "typedScenarios")
+ void typedModeMergesCountsAndHandlesNulls(Scenario scenario) {
+ String type = scenario._dataType.name();
+ String first = scenario._dataType == DataType.STRING ? "apple" :
"2026-09-03 10:11:12.123";
+ String shared = scenario._dataType == DataType.STRING ? "banana" :
"2026-09-04 10:11:12.456";
+ String last = scenario._dataType == DataType.STRING ? "cherry" :
"2026-09-05 10:11:12.789";
+ // Each instance has a different local mode; the shared value wins only
after merging full counts.
+ scenario.getDeclaringTable(true)
+ .onFirstInstance("myField", first, first, first, shared, shared,
"null")
+ .andOnSecondInstance("myField", last, last, last, shared, shared,
"null")
+ .whenQuery("select mode(myField, 'MIN', '" + type + "') as mode from
testTable")
Review Comment:
**[Major] Exercise ordinary MODE through the single-stage broker**
These execution tests manually supply the internal type argument, while the
optimizer tests check only rewritten expressions. Please add a
`CustomDataQueryClusterIntegrationTest` case using ordinary `MODE(stringCol)`,
`MODE(timestampCol)`, and scalar expressions through broker rewriting and
result reduction. Assert both values and result types for aggregation, grouped
queries, and post-aggregation expressions without supplying the third argument.
The distributed resource tests cover the separate multi-stage planning path,
so they do not exercise this single-stage gap.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java:
##########
@@ -133,4 +164,96 @@ private void maybeRewriteAggregateFunction(@Nullable
Expression expression, Sche
}
}
}
+
+ @Nullable
+ private static ColumnDataType getOperandType(Expression operand, Schema
schema) {
+ if (operand.isSetIdentifier()) {
+ FieldSpec fieldSpec =
schema.getFieldSpecFor(operand.getIdentifier().getName());
+ return fieldSpec != null
+ ? ColumnDataType.fromDataType(fieldSpec.getDataType(),
fieldSpec.isSingleValueField())
+ : null;
+ }
+ if (operand.isSetLiteral()) {
+ return
RequestUtils.getLiteralTypeAndValue(operand.getLiteral()).getLeft();
+ }
+ if (!operand.isSetFunctionCall()) {
+ return null;
+ }
+ Function function = operand.getFunctionCall();
+ List<Expression> arguments = function.getOperands();
+ String name = FunctionRegistry.canonicalize(function.getOperator());
+ switch (name) {
+ case "cast":
+ return literalType(arguments, 1);
+ case "jsonextractscalar":
+ case "jsonextractscalarfast":
+ case "jsonextractscalarfirstmatch":
+ case "jsonextractscalarfory":
+ return literalType(arguments, 2);
+ case "case":
+ // CASE stores alternating condition/result pairs followed by an
optional ELSE result.
+ ColumnDataType resultType = ColumnDataType.UNKNOWN;
+ for (int i = 1; i < arguments.size(); i += 2) {
+ resultType = commonType(resultType, getOperandType(arguments.get(i),
schema));
+ }
+ if (arguments.size() % 2 == 1) {
+ resultType = commonType(resultType,
getOperandType(arguments.get(arguments.size() - 1), schema));
+ }
+ return resultType;
+ case "datetimeconvert":
+ if (arguments.size() < 3 || !arguments.get(2).isSetLiteral()
+ || !arguments.get(2).getLiteral().isSetStringValue()) {
+ return null;
+ }
+ DateTimeFieldSpec.TimeFormat format =
+ new
DateTimeFormatSpec(arguments.get(2).getLiteral().getStringValue()).getTimeFormat();
+ return format == DateTimeFieldSpec.TimeFormat.EPOCH || format ==
DateTimeFieldSpec.TimeFormat.TIMESTAMP
+ ? ColumnDataType.LONG
+ : ColumnDataType.STRING;
+ default:
+ ColumnDataType[] argumentTypes = new ColumnDataType[arguments.size()];
+ for (int i = 0; i < arguments.size(); i++) {
+ argumentTypes[i] = getOperandType(arguments.get(i), schema);
+ if (argumentTypes[i] == null) {
+ return null;
+ }
+ }
+ FunctionInfo functionInfo = FunctionRegistry.lookupFunctionInfo(name,
argumentTypes);
+ return functionInfo != null ?
FunctionUtils.getColumnDataType(functionInfo.getMethod().getReturnType()) :
null;
+ }
+ }
+
+ @Nullable
+ private static ColumnDataType commonType(@Nullable ColumnDataType left,
@Nullable ColumnDataType right) {
+ if (left == null || right == null) {
+ return null;
+ }
+ if (left == ColumnDataType.UNKNOWN) {
+ return right;
+ }
+ return right == ColumnDataType.UNKNOWN || left == right ? left : null;
Review Comment:
**[Critical] Match CASE timestamp coercion rules**
For `MODE(CASE WHEN longCol = 0 THEN timestampCol ELSE '1970-01-01
00:00:00.000' END)`, this method rejects `TIMESTAMP` combined with a `STRING`
literal. `CaseTransformFunction` accepts that literal as `TIMESTAMP`, so MODE
silently retains `DOUBLE` output containing epoch milliseconds instead of
preserving `TIMESTAMP`, potentially losing precision beyond 2^53.
Please apply CASE's literal coercion rules during inference. A targeted
optimizer test expecting the appended `'MIN', 'TIMESTAMP'` arguments fails for
this expression.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java:
##########
@@ -719,6 +824,25 @@ public Double extractFinalResult(Double2LongOpenHashMap
intermediateResult) {
}
}
+ @Nullable
+ private Comparable<?> extractComparableFinalResult(@Nullable Map<?, Long>
counts) {
+ Comparable mode = null;
+ long maxCount = 0;
+ if (counts != null) {
+ for (Map.Entry<?, Long> entry : counts.entrySet()) {
+ Comparable value = (Comparable) entry.getKey();
+ long count = entry.getValue();
+ if (mode == null || count > maxCount || (count == maxCount
+ && (_multiModeReducerType == MultiModeReducerType.MIN
+ ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) {
Review Comment:
**[Minor] Separate multiline ternary branches**
Please move `: value.compareTo(mode) > 0` onto its own line to follow the
repository's multiline ternary convention.
--
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]
