From da4922f54fdb9a8fd8f640dd13556982ab12ea73 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 7 Sep 2026 02:22:26 -0700 Subject: [PATCH 1/8] Support string and timestamp MODE aggregations --- .../function/AggregationFunctionFactory.java | 4 + ...BaseComparableModeAggregationFunction.java | 228 ++++++++++++++++++ .../ModeStringAggregationFunction.java | 139 +++++++++++ .../ModeTimestampAggregationFunction.java | 78 ++++++ .../core/query/optimizer/QueryOptimizer.java | 3 +- ...deAggregationFunctionRewriteOptimizer.java | 186 ++++++++++++++ .../reduce/AggregationDataTableReducer.java | 4 +- .../function/ModeAggregationFunctionTest.java | 97 ++++++++ ...ModeNonNumericAggregationFunctionTest.java | 147 +++++++++++ ...gregationFunctionRewriteOptimizerTest.java | 116 +++++++++ .../apache/pinot/queries/ModeQueriesTest.java | 97 +++++++- .../PinotAggregateFunctionRewriteRule.java | 15 ++ .../calcite/sql/fun/PinotOperatorTable.java | 1 - .../query/queries/ModeSqlPlannerTest.java | 134 ++++++++++ .../resources/queries/ModeAggregates.json | 150 ++++++++++++ .../segment/spi/AggregationFunctionType.java | 25 +- 16 files changed, 1418 insertions(+), 6 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java create mode 100644 pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java create mode 100644 pinot-query-runtime/src/test/resources/queries/ModeAggregates.json diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java index fdc93df43e80..11b05e4f32a8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java @@ -245,6 +245,10 @@ public static AggregationFunction getAggregationFunction(FunctionContext functio return new AvgAggregationFunction(arguments, nullHandlingEnabled); case MODE: return new ModeAggregationFunction(arguments, nullHandlingEnabled); + case MODESTRING: + return new ModeStringAggregationFunction(arguments, nullHandlingEnabled); + case MODETIMESTAMP: + return new ModeTimestampAggregationFunction(arguments, nullHandlingEnabled); case ANYVALUE: return new AnyValueAggregationFunction(arguments, nullHandlingEnabled); case FIRSTWITHTIME: { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java new file mode 100644 index 000000000000..c8553d7e4d56 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java @@ -0,0 +1,228 @@ +/** + * 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.core.query.aggregation.function; + +import it.unimi.dsi.fastutil.ints.Int2LongOpenHashMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.ObjectAggregationResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.ObjectGroupByResultHolder; +import org.apache.pinot.segment.spi.index.reader.Dictionary; + +import static com.google.common.base.Preconditions.checkArgument; + + +/// Counts comparable values and resolves equally frequent values using MIN (default) or MAX. +/// +/// Instances are immutable and may be shared across segments. Accumulators belong to result holders, and dictionary +/// identifiers are converted to values before merging across segments. +/// Numeric MODE retains its existing implementation. +abstract class BaseComparableModeAggregationFunction> + extends BaseSingleInputAggregationFunction, T> { + private final boolean _minimum; + + protected BaseComparableModeAggregationFunction(List arguments, boolean nullHandlingEnabled, + String valueType) { + super(checkArguments(arguments), nullHandlingEnabled); + String reducer = "MIN"; + if (arguments.size() == 2) { + ExpressionContext argument = arguments.get(1); + checkArgument(argument.getType() == ExpressionContext.Type.LITERAL, + "MODE tie reducer must be a literal MIN or MAX for %s", valueType); + reducer = argument.getLiteral().getStringValue(); + } + checkArgument("MIN".equals(reducer) || "MAX".equals(reducer), + "MODE for %s supports only MIN or MAX tie reducers, got: %s", valueType, reducer); + _minimum = "MIN".equals(reducer); + } + + private static ExpressionContext checkArguments(List arguments) { + checkArgument(arguments.size() == 1 || arguments.size() == 2, + "MODE expects one or two arguments, got: %s", arguments.size()); + return arguments.get(0); + } + + protected abstract Map newValueMap(); + + protected abstract ValueCounter valueCounter(BlockValSet blockValSet); + + @FunctionalInterface + protected interface ValueCounter { + void add(Map counts, int row); + } + + protected abstract T dictionaryValue(Dictionary dictionary, int dictionaryId); + + @Override + public AggregationResultHolder createAggregationResultHolder() { + return new ObjectAggregationResultHolder(); + } + + @Override + public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int maxCapacity) { + return new ObjectGroupByResultHolder(initialCapacity, maxCapacity); + } + + @Override + public void aggregate(int length, AggregationResultHolder holder, + Map blockValSetMap) { + BlockValSet values = blockValSetMap.get(_expression); + Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; + if (dictionary != null) { + int[] ids = values.getDictionaryIdsSV(); + forEachNotNull(length, values, (from, to) -> { + DictionaryCounts counts = getValue(holder, () -> new DictionaryCounts(dictionary)); + for (int i = from; i < to; i++) { + counts._counts.addTo(ids[i], 1L); + } + }); + } else { + ValueCounter counter = valueCounter(values); + forEachNotNull(length, values, (from, to) -> { + Map counts = getValue(holder, this::newValueMap); + for (int i = from; i < to; i++) { + counter.add(counts, i); + } + }); + } + } + + @Override + public void aggregateGroupBySV(int length, int[] groupKeys, GroupByResultHolder holder, + Map blockValSetMap) { + BlockValSet values = blockValSetMap.get(_expression); + Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; + if (dictionary != null) { + int[] ids = values.getDictionaryIdsSV(); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + DictionaryCounts counts = getValue(holder, groupKeys[i], () -> new DictionaryCounts(dictionary)); + counts._counts.addTo(ids[i], 1L); + } + }); + } else { + ValueCounter counter = valueCounter(values); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + Map counts = getValue(holder, groupKeys[i], this::newValueMap); + counter.add(counts, i); + } + }); + } + } + + @Override + public void aggregateGroupByMV(int length, int[][] groupKeys, GroupByResultHolder holder, + Map blockValSetMap) { + BlockValSet values = blockValSetMap.get(_expression); + Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; + if (dictionary != null) { + int[] ids = values.getDictionaryIdsSV(); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + for (int groupKey : groupKeys[i]) { + DictionaryCounts counts = getValue(holder, groupKey, () -> new DictionaryCounts(dictionary)); + counts._counts.addTo(ids[i], 1L); + } + } + }); + } else { + ValueCounter counter = valueCounter(values); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + for (int groupKey : groupKeys[i]) { + Map counts = getValue(holder, groupKey, this::newValueMap); + counter.add(counts, i); + } + } + }); + } + } + + @Nullable + @Override + public Map extractAggregationResult(AggregationResultHolder holder) { + return extractCounts(holder.getResult()); + } + + @Nullable + @Override + public Map extractGroupByResult(GroupByResultHolder holder, int groupKey) { + return extractCounts(holder.getResult(groupKey)); + } + + @Nullable + @SuppressWarnings("unchecked") + private Map extractCounts(@Nullable Object result) { + if (result instanceof DictionaryCounts) { + DictionaryCounts dictionaryCounts = (DictionaryCounts) result; + Map counts = newValueMap(); + dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> counts.put( + dictionaryValue(dictionaryCounts._dictionary, entry.getIntKey()), entry.getLongValue())); + return counts; + } + return (Map) result; + } + + @Override + public Map merge(Map left, Map right) { + right.forEach((value, count) -> left.merge(value, count, Long::sum)); + return left; + } + + @Override + public ColumnDataType getIntermediateResultColumnType() { + return ColumnDataType.OBJECT; + } + + @Nullable + @Override + public T extractFinalResult(@Nullable Map counts) { + if (counts == null || counts.isEmpty()) { + return null; + } + T mode = null; + long maxCount = 0; + for (Map.Entry entry : counts.entrySet()) { + T value = entry.getKey(); + long count = entry.getValue(); + if (mode == null || count > maxCount || (count == maxCount + && (_minimum ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { + mode = value; + maxCount = count; + } + } + return mode; + } + + private static final class DictionaryCounts { + private final Dictionary _dictionary; + private final Int2LongOpenHashMap _counts = new Int2LongOpenHashMap(); + + private DictionaryCounts(Dictionary dictionary) { + _dictionary = dictionary; + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java new file mode 100644 index 000000000000..fc2805fb9eb8 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java @@ -0,0 +1,139 @@ +/** + * 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.core.query.aggregation.function; + +import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.CustomObject; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.ObjectSerDeUtils; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.segment.spi.index.reader.Dictionary; + + +/// String implementation of MODE, with lexicographic MIN/MAX tie resolution and a fixed STRING result type. +/// Instances are immutable; per-segment frequency maps are stored in the result holders. +public class ModeStringAggregationFunction extends BaseComparableModeAggregationFunction { + public ModeStringAggregationFunction(List arguments, boolean nullHandlingEnabled) { + super(arguments, nullHandlingEnabled, "STRING"); + } + + @Override + public AggregationFunctionType getType() { + return AggregationFunctionType.MODESTRING; + } + + @Override + public ColumnDataType getFinalResultColumnType() { + return ColumnDataType.STRING; + } + + @Override + protected Map newValueMap() { + return new StringModeCounts(); + } + + @Override + protected ValueCounter valueCounter(BlockValSet blockValSet) { + String[] values = blockValSet.getStringValuesSV(); + return (counts, row) -> ((StringModeCounts) counts).addTo(values[row], 1L); + } + + @Override + protected String dictionaryValue(Dictionary dictionary, int dictionaryId) { + return dictionary.getStringValue(dictionaryId); + } + + @Override + @SuppressWarnings({"unchecked", "rawtypes", "deprecation"}) + public SerializedIntermediateResult serializeIntermediateResult(Map counts) { + // Reuse the existing map wire encoding so generic aggregation bridges can deserialize the frequency state. + return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.Map.getValue(), + ObjectSerDeUtils.MAP_SER_DE.serialize((Map) counts)); + } + + @Override + public Map deserializeIntermediateResult(CustomObject customObject) { + return new StringModeCounts(ObjectSerDeUtils.deserialize(customObject)); + } + + /// Frequency state with an O(1) conservative estimate of the retained string-key payload. + /// Accumulation uses [#addTo] and dictionary extraction and boxed [Map#merge] use [#put]. + /// Each distinct key is charged once, assuming UTF-16 storage plus object and array overhead. + /// Instances belong to one result holder and are not thread-safe. + public static final class StringModeCounts extends Object2LongOpenHashMap { + private long _retainedStringBytes; + + public StringModeCounts() { + } + + /// Restores accounting once when a generic map is deserialized from the existing wire format. + public StringModeCounts(Map counts) { + super(counts.size()); + counts.forEach((value, count) -> put(value, count.longValue())); + } + + public long getRetainedStringBytes() { + return _retainedStringBytes; + } + + @Override + public long addTo(String value, long increment) { + int previousSize = size(); + long previousCount = super.addTo(value, increment); + if (size() != previousSize) { + _retainedStringBytes += retainedStringBytes(value); + } + return previousCount; + } + + @Override + public long put(String value, long count) { + int previousSize = size(); + long previousCount = super.put(value, count); + if (size() != previousSize) { + _retainedStringBytes += retainedStringBytes(value); + } + return previousCount; + } + + @Override + public long removeLong(Object value) { + int previousSize = size(); + long previousCount = super.removeLong(value); + if (size() != previousSize) { + _retainedStringBytes -= retainedStringBytes((String) value); + } + return previousCount; + } + + @Override + public void clear() { + super.clear(); + _retainedStringBytes = 0; + } + + private static long retainedStringBytes(String value) { + return 48 + 2L * value.length(); + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java new file mode 100644 index 000000000000..959971de51f1 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java @@ -0,0 +1,78 @@ +/** + * 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.core.query.aggregation.function; + +import it.unimi.dsi.fastutil.longs.Long2LongMap; +import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.CustomObject; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.ObjectSerDeUtils; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.segment.spi.index.reader.Dictionary; + + +/// Timestamp implementation of MODE, preserving epoch milliseconds without conversion through DOUBLE. +/// Instances are immutable; per-segment frequency maps are stored in the result holders. +public class ModeTimestampAggregationFunction extends BaseComparableModeAggregationFunction { + public ModeTimestampAggregationFunction(List arguments, boolean nullHandlingEnabled) { + super(arguments, nullHandlingEnabled, "TIMESTAMP"); + } + + @Override + public AggregationFunctionType getType() { + return AggregationFunctionType.MODETIMESTAMP; + } + + @Override + public ColumnDataType getFinalResultColumnType() { + return ColumnDataType.TIMESTAMP; + } + + @Override + protected Map newValueMap() { + return new Long2LongOpenHashMap(); + } + + @Override + protected ValueCounter valueCounter(BlockValSet blockValSet) { + long[] values = blockValSet.getLongValuesSV(); + return (counts, row) -> ((Long2LongOpenHashMap) counts).addTo(values[row], 1L); + } + + @Override + protected Long dictionaryValue(Dictionary dictionary, int dictionaryId) { + return dictionary.getLongValue(dictionaryId); + } + + @Override + public SerializedIntermediateResult serializeIntermediateResult(Map counts) { + Long2LongMap longCounts = counts instanceof Long2LongMap ? (Long2LongMap) counts : new Long2LongOpenHashMap(counts); + return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.Long2LongMap.getValue(), + ObjectSerDeUtils.LONG_2_LONG_MAP_SER_DE.serialize(longCounts)); + } + + @Override + public Map deserializeIntermediateResult(CustomObject customObject) { + return ObjectSerDeUtils.deserialize(customObject); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java index 635250800e03..68b1002e3aec 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java @@ -31,6 +31,7 @@ import org.apache.pinot.core.query.optimizer.filter.TextMatchFilterOptimizer; import org.apache.pinot.core.query.optimizer.filter.TimePredicateFilterOptimizer; import org.apache.pinot.core.query.optimizer.statement.AggregateFunctionRewriteOptimizer; +import org.apache.pinot.core.query.optimizer.statement.ModeAggregationFunctionRewriteOptimizer; import org.apache.pinot.core.query.optimizer.statement.StatementOptimizer; import org.apache.pinot.spi.data.Schema; @@ -48,7 +49,7 @@ public class QueryOptimizer { new MergeRangeFilterOptimizer(), new TextMatchFilterOptimizer()); private static final List STATEMENT_OPTIMIZERS = - List.of(new AggregateFunctionRewriteOptimizer()); + List.of(new AggregateFunctionRewriteOptimizer(), new ModeAggregationFunctionRewriteOptimizer()); /// Optimizes the given query. public void optimize(PinotQuery pinotQuery, @Nullable Schema schema) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java new file mode 100644 index 000000000000..9fb460d33ba7 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java @@ -0,0 +1,186 @@ +/** + * 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.core.query.optimizer.statement; + +import java.util.List; +import java.util.Locale; +import javax.annotation.Nullable; +import org.apache.pinot.common.function.FunctionInfo; +import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.common.function.FunctionUtils; +import org.apache.pinot.common.request.Expression; +import org.apache.pinot.common.request.Function; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.request.context.LiteralContext; +import org.apache.pinot.common.request.context.RequestContextUtils; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.spi.data.DateTimeFieldSpec; +import org.apache.pinot.spi.data.DateTimeFormatSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; + + +/// When `autoRewriteAggregationType` is enabled, resolves string and timestamp MODE expressions to implementations +/// with fixed result types before execution. +/// This also supplies the broker with the correct result type when no rows match or groups are trimmed before +/// finalization. Type inference reads schema and function metadata only: server-dependent transforms such as LOOKUP +/// must not be initialized on the broker. Expressions whose type is unknown retain the legacy MODE implementation. +public class ModeAggregationFunctionRewriteOptimizer implements StatementOptimizer { + @Override + public void optimize(PinotQuery pinotQuery, @Nullable Schema schema) { + // Keep existing timestamp MODE requests compatible with older servers during rolling upgrades. + if (schema == null || pinotQuery.getQueryOptions() == null + || !Boolean.parseBoolean(pinotQuery.getQueryOptions().get(QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE))) { + return; + } + rewriteExpressions(pinotQuery.getSelectList(), schema); + rewriteExpressions(pinotQuery.getGroupByList(), schema); + rewriteExpressions(pinotQuery.getOrderByList(), schema); + rewriteExpression(pinotQuery.getFilterExpression(), schema); + rewriteExpression(pinotQuery.getHavingExpression(), schema); + } + + private static void rewriteExpressions(@Nullable List expressions, Schema schema) { + if (expressions != null) { + for (Expression expression : expressions) { + rewriteExpression(expression, schema); + } + } + } + + private static void rewriteExpression(@Nullable Expression expression, Schema schema) { + if (expression == null || !expression.isSetFunctionCall()) { + return; + } + Function function = expression.getFunctionCall(); + List operands = function.getOperands(); + rewriteExpressions(operands, schema); + if (!AggregationFunctionType.MODE.getName().equalsIgnoreCase(function.getOperator()) || operands.isEmpty()) { + return; + } + + ColumnDataType operandType = getOperandType(operands.get(0), schema); + if (operandType == ColumnDataType.STRING) { + function.setOperator("modestring"); + } else if (operandType == ColumnDataType.TIMESTAMP) { + function.setOperator("modetimestamp"); + } + } + + @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()) { + LiteralContext literal = RequestContextUtils.getExpression(operand).getLiteral(); + return ColumnDataType.fromDataType(literal.getType(), literal.isSingleValue()); + } + if (!operand.isSetFunctionCall()) { + return null; + } + Function function = operand.getFunctionCall(); + List 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; + } + + @Nullable + private static ColumnDataType literalType(List arguments, int position) { + if (arguments.size() <= position || !arguments.get(position).isSetLiteral() + || !arguments.get(position).getLiteral().isSetStringValue()) { + return null; + } + String type = arguments.get(position).getLiteral().getStringValue().toUpperCase(Locale.ROOT); + switch (type) { + case "VARCHAR": + case "CHAR": + case "JSON": + return ColumnDataType.STRING; + case "BIGINT": + return ColumnDataType.LONG; + case "INTEGER": + return ColumnDataType.INT; + case "REAL": + return ColumnDataType.FLOAT; + case "DECIMAL": + return ColumnDataType.BIG_DECIMAL; + default: + try { + return ColumnDataType.valueOf(type); + } catch (IllegalArgumentException e) { + return null; + } + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java index 96f53690f5d6..66d84f0fe83e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java @@ -222,7 +222,9 @@ private ResultTable reduceToResultTable(DataSchema dataSchema, Object[] finalRes int numColumns = columnDataTypes.length; for (Object[] rewrittenRow : rows) { for (int j = 0; j < numColumns; j++) { - rewrittenRow[j] = columnDataTypes[j].format(rewrittenRow[j]); + if (rewrittenRow[j] != null) { + rewrittenRow[j] = columnDataTypes[j].format(rewrittenRow[j]); + } } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunctionTest.java index 78a7cdcada99..a902e5682e1a 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunctionTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.core.query.aggregation.function; +import java.sql.Timestamp; import org.apache.pinot.queries.FluentQueryTest; import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.data.FieldSpec.DataType; @@ -40,6 +41,102 @@ Object[] scenarios() { }; } + @DataProvider(name = "stringScenarios") + Object[] stringScenarios() { + return new Object[]{new Scenario(DataType.STRING, true), new Scenario(DataType.STRING, false)}; + } + + @DataProvider(name = "timestampScenarios") + Object[] timestampScenarios() { + return new Object[]{new Scenario(DataType.TIMESTAMP, true), new Scenario(DataType.TIMESTAMP, false)}; + } + + @Test(dataProvider = "stringScenarios") + void stringModeMergesValueCountsAcrossSegments(Scenario scenario) { + // Each server has a different local mode. The shared value wins after merging the counts. + scenario.getDeclaringTable(true) + .onFirstInstance("myField", "apple", "apple", "apple", "banana", "banana", "null") + .andSegment("myField", "banana") + .andOnSecondInstance("myField", "cherry", "cherry", "cherry", "banana", "banana", "null") + .whenQuery("select modeString(myField) as mode from testTable") + .thenResultTextIs("mode[STRING]\nbanana"); + } + + @Test(dataProvider = "stringScenarios") + void stringModeTieReducers(Scenario scenario) { + scenario.getDeclaringTable(true) + .onFirstInstance("myField", "zebra", "apple", "zebra", "banana", "null") + .andOnSecondInstance("myField", "apple", "apple", "zebra", "null") + .whenQuery("select modeString(myField), modeString(myField, 'MIN'), modeString(myField, 'MAX') from testTable") + .thenResultIs("STRING | STRING | STRING", "apple | apple | zebra"); + } + + @Test(dataProvider = "stringScenarios") + void stringModeWithNullAndComputedEmptyString(Scenario scenario) { + // The CSV-backed fixture treats empty cells as null. Generate an actual empty string in the query instead. + scenario.getDeclaringTable(true) + .onFirstInstance("myField", "null", "empty", "empty", "apple") + .andOnSecondInstance("myField", "null", "null", "apple") + .whenQuery("select modeString(case when myField = 'empty' then '' else myField end) as mode from testTable") + .thenResultIs(new Object[]{""}) + .whenQuery("select modeString(case when myField = 'empty' then null else myField end) as mode from testTable") + .thenResultTextIs("mode[STRING]\napple") + .whenQuery("select myField, modeString(case when myField = 'empty' then '' else myField end) from testTable " + + "group by myField order by myField") + .thenResultIs(new Object[]{"apple", "apple"}, new Object[]{"empty", ""}, new Object[]{null, null}); + } + + @Test(dataProvider = "stringScenarios") + void stringModeAllNullAndEmptyInput(Scenario scenario) { + scenario.getDeclaringTable(true) + .onFirstInstance("myField", "null", "null") + .andOnSecondInstance("myField", "null") + .whenQuery("select modeString(myField) as mode from testTable") + .thenResultIs(new Object[]{null}) + .whenQuery("select modeString(myField) as mode from testTable where myField = 'absent'") + .thenResultIs(new Object[]{null}) + .whenQuery("select 'group', modeString(myField) as mode from testTable group by 'group'") + .thenResultIs("STRING | STRING", "group | null"); + } + + @Test(dataProvider = "stringScenarios") + void stringModeWithoutNullHandling(Scenario scenario) { + // Without null handling, three null defaults must beat the two ordinary values. + scenario.getDeclaringTable(false) + .onFirstInstance("myField", "null", "null", "null", "apple") + .andOnSecondInstance("myField", "apple") + .whenQuery("select modeString(myField) as mode from testTable") + .thenResultIs(new Object[]{"null"}) + .whenQuery("select modeString(myField) as mode from testTable where myField = 'absent'") + .thenResultIs(new Object[]{null}); + } + + @Test(dataProvider = "timestampScenarios") + void timestampModePreservesTypeAndTieOrdering(Scenario scenario) { + scenario.getDeclaringTable(true) + .onFirstInstance("myField", "2026-09-03 10:11:12.123", "2026-09-04 10:11:12.456", "null") + .andOnSecondInstance("myField", "2026-09-04 10:11:12.456", "2026-09-03 10:11:12.123", "null") + .whenQuery("select modeTimestamp(myField) as mode, modeTimestamp(myField, 'MAX') as latest from testTable") + .thenResultTextIs("mode[TIMESTAMP] | latest[TIMESTAMP]\n" + + "2026-09-03 10:11:12.123 | 2026-09-04 10:11:12.456") + .whenQuery("select myField, modeTimestamp(myField) as mode from testTable group by myField order by myField") + .thenResultIs(new Object[]{"2026-09-03 10:11:12.123", "2026-09-03 10:11:12.123"}, + new Object[]{"2026-09-04 10:11:12.456", "2026-09-04 10:11:12.456"}, new Object[]{null, null}) + .whenQuery("select fromTimestamp(modeTimestamp(myField)) as epochMillis from testTable") + .thenResultIs(new Object[]{Timestamp.valueOf("2026-09-03 10:11:12.123").getTime()}); + } + + @Test(dataProvider = "timestampScenarios") + void timestampModeAllNullAndEmptyInput(Scenario scenario) { + scenario.getDeclaringTable(true) + .onFirstInstance("myField", "null", "null") + .andOnSecondInstance("myField", "null") + .whenQuery("select modeTimestamp(myField) as mode from testTable") + .thenResultIs(new Object[]{null}) + .whenQuery("select modeTimestamp(myField) as mode from testTable where myField > '2026-09-03 00:00:00'") + .thenResultIs(new Object[]{null}); + } + public class Scenario { private final DataType _dataType; private final boolean _dictionary; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java new file mode 100644 index 000000000000..2e54aa61ea66 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java @@ -0,0 +1,147 @@ +/** + * 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.core.query.aggregation.function; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.CustomObject; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.SyntheticBlockValSets; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.roaringbitmap.RoaringBitmap; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +/// Verifies serialization, exact timestamp values and null handling for the typed MODE implementations. +public class ModeNonNumericAggregationFunctionTest { + private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("value"); + + @Test + public void testStringIntermediateResultsRoundTripAndMerge() { + ModeStringAggregationFunction function = new ModeStringAggregationFunction(List.of(EXPRESSION), true); + var first = aggregateAndRoundTrip(function, + SyntheticBlockValSets.Str.create(null, new String[]{"é", "é", "é", "苹果", "苹果", ""}), 6); + var second = aggregateAndRoundTrip(function, + SyntheticBlockValSets.Str.create(null, new String[]{"zebra", "zebra", "zebra", "苹果", "苹果", "\0"}), 6); + + assertEquals(function.getFinalResultColumnType(), ColumnDataType.STRING); + assertEquals(function.extractFinalResult(first), "é"); + assertEquals(function.extractFinalResult(second), "zebra"); + assertEquals(function.extractFinalResult(function.merge(first, second)), "苹果"); + } + + @Test + public void testTimestampIntermediateResultsPreserveLongPrecision() { + // Adjacent long values above 2^53 become identical if converted through double. + long earlier = 9_007_199_254_740_992L; + long later = earlier + 1; + ModeTimestampAggregationFunction minFunction = + new ModeTimestampAggregationFunction(List.of(EXPRESSION), true); + ModeTimestampAggregationFunction maxFunction = + new ModeTimestampAggregationFunction(argumentsWithReducer("MAX"), true); + var first = aggregateAndRoundTrip(minFunction, timestampValues(earlier, later, later), 3); + var second = aggregateAndRoundTrip(minFunction, timestampValues(earlier), 1); + + assertEquals(minFunction.getFinalResultColumnType(), ColumnDataType.TIMESTAMP); + assertEquals(minFunction.extractFinalResult(first), Long.valueOf(later)); + var merged = minFunction.merge(first, second); + assertEquals(minFunction.extractFinalResult(merged), Long.valueOf(earlier)); + assertEquals(maxFunction.extractFinalResult(merged), Long.valueOf(later)); + } + + @Test + public void testStringModeSkipsNullRowsForMultiValueGroupKeys() { + ModeStringAggregationFunction function = new ModeStringAggregationFunction(List.of(EXPRESSION), true); + GroupByResultHolder holder = function.createGroupByResultHolder(3, 3); + BlockValSet values = SyntheticBlockValSets.Str.create(RoaringBitmap.bitmapOf(0, 2), + new String[]{"ignored", "alpha", "ignored", "beta", "alpha"}); + function.aggregateGroupByMV(5, new int[][]{{0, 1}, {0}, {1, 2}, {0, 1}, {0}}, holder, + Map.of(EXPRESSION, values)); + + assertEquals(function.extractFinalResult(function.extractGroupByResult(holder, 0)), "alpha"); + assertEquals(function.extractFinalResult(function.extractGroupByResult(holder, 1)), "beta"); + assertNull(function.extractFinalResult(function.extractGroupByResult(holder, 2))); + } + + @Test + public void testEmptyResultsAreNullWithEitherNullHandlingMode() { + for (boolean nullHandlingEnabled : new boolean[]{false, true}) { + ModeStringAggregationFunction stringFunction = + new ModeStringAggregationFunction(List.of(EXPRESSION), nullHandlingEnabled); + ModeTimestampAggregationFunction timestampFunction = + new ModeTimestampAggregationFunction(List.of(EXPRESSION), nullHandlingEnabled); + assertNull(stringFunction.extractFinalResult(null)); + assertNull(timestampFunction.extractFinalResult(null)); + assertNull(stringFunction.extractFinalResult( + stringFunction.extractAggregationResult(stringFunction.createAggregationResultHolder()))); + assertNull(timestampFunction.extractFinalResult( + timestampFunction.extractAggregationResult(timestampFunction.createAggregationResultHolder()))); + } + } + + @Test + public void testNonNumericModesRejectAverageReducer() { + IllegalArgumentException stringError = expectThrows(IllegalArgumentException.class, + () -> new ModeStringAggregationFunction(argumentsWithReducer("AVG"), true)); + assertTrue(stringError.getMessage().contains("AVG")); + assertTrue(stringError.getMessage().contains("STRING")); + IllegalArgumentException timestampError = expectThrows(IllegalArgumentException.class, + () -> new ModeTimestampAggregationFunction(argumentsWithReducer("AVG"), true)); + assertTrue(timestampError.getMessage().contains("AVG")); + assertTrue(timestampError.getMessage().contains("TIMESTAMP")); + } + + private static List argumentsWithReducer(String reducer) { + return List.of(EXPRESSION, ExpressionContext.forLiteral(Literal.stringValue(reducer))); + } + + private static BlockValSet timestampValues(long... values) { + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.isSingleValue()).thenReturn(true); + when(blockValSet.getValueType()).thenReturn(DataType.TIMESTAMP); + when(blockValSet.getLongValuesSV()).thenReturn(values); + return blockValSet; + } + + private static > I aggregateAndRoundTrip(AggregationFunction function, + BlockValSet values, int length) { + AggregationResultHolder holder = function.createAggregationResultHolder(); + function.aggregate(length, holder, Map.of(EXPRESSION, values)); + I intermediateResult = function.extractAggregationResult(holder); + AggregationFunction.SerializedIntermediateResult serialized = + function.serializeIntermediateResult(intermediateResult); + I deserialized = function.deserializeIntermediateResult( + new CustomObject(serialized.getType(), ByteBuffer.wrap(serialized.getBytes()))); + assertEquals(deserialized, intermediateResult); + return deserialized; + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java new file mode 100644 index 000000000000..259f31cd9b18 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java @@ -0,0 +1,116 @@ +/** + * 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.core.query.optimizer.statement; + +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.core.query.optimizer.QueryOptimizer; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +/// Verifies MODE rewriting through the single-stage optimizer, including expressions and post-aggregation clauses. +public class ModeAggregationFunctionRewriteOptimizerTest { + private static final QueryOptimizer OPTIMIZER = new QueryOptimizer(); + private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName("testTable") + .addSingleValueDimension("stringCol", DataType.STRING) + .addSingleValueDimension("intCol", DataType.INT) + .addSingleValueDimension("longCol", DataType.LONG) + .addSingleValueDimension("floatCol", DataType.FLOAT) + .addSingleValueDimension("doubleCol", DataType.DOUBLE) + .addMultiValueDimension("mvStringCol", DataType.STRING) + .addDateTime("timestampCol", DataType.TIMESTAMP, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + + @DataProvider + public Object[][] modeExpressions() { + return new Object[][]{ + {"MODE(stringCol)", "modeString(stringCol)"}, + {"MODE(timestampCol, 'MAX')", "modeTimestamp(timestampCol, 'MAX')"}, + {"MODE(CONCAT(stringCol, 'suffix'))", "modeString(CONCAT(stringCol, 'suffix'))"}, + {"MODE(JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', ''))", + "modeString(JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', ''))"}, + {"MODE(CAST(intCol AS STRING))", "modeString(CAST(intCol AS STRING))"}, + {"MODE(CAST(stringCol AS TIMESTAMP))", "modeTimestamp(CAST(stringCol AS TIMESTAMP))"}, + {"MODE(CASE WHEN intCol > 0 THEN stringCol ELSE 'other' END)", + "modeString(CASE WHEN intCol > 0 THEN stringCol ELSE 'other' END)"}, + {"MODE('literal')", "modeString('literal')"}, + {"fromTimestamp(MODE(timestampCol))", "fromTimestamp(modeTimestamp(timestampCol))"} + }; + } + + @Test(dataProvider = "modeExpressions") + public void testModeExpressions(String original, String rewritten) { + for (boolean nullHandlingEnabled : new boolean[]{false, true}) { + String prefix = + "SET autoRewriteAggregationType=true; SET enableNullHandling=" + nullHandlingEnabled + "; SELECT "; + TestHelper.assertEqualsQuery(prefix + original + " AS commonValue FROM testTable", + prefix + rewritten + " AS commonValue FROM testTable", SCHEMA); + } + } + + @Test + public void testModeInHavingAndOrderBy() { + TestHelper.assertEqualsQuery( + "SET autoRewriteAggregationType=true; " + + "SELECT intCol, MODE(stringCol) AS commonValue FROM testTable GROUP BY intCol " + + "HAVING MODE(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END) = 'value' " + + "ORDER BY MODE(timestampCol) DESC", + "SET autoRewriteAggregationType=true; " + + "SELECT intCol, modeString(stringCol) AS commonValue FROM testTable GROUP BY intCol " + + "HAVING modeString(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END) = 'value' " + + "ORDER BY modeTimestamp(timestampCol) DESC", SCHEMA); + } + + @Test + public void testNumericModeAndOptInRewritesRemainUnchanged() { + assertUnchanged("SELECT MODE(intCol), MODE(longCol), MODE(floatCol), MODE(doubleCol), " + + "MODE(CAST(stringCol AS LONG)), MODE(fromDateTime(stringCol, 'yyyy-MM-dd HH:mm:ss')), " + + "MIN(stringCol), MAX(longCol), SUM(intCol) FROM testTable", SCHEMA); + assertUnchanged("SET autoRewriteAggregationType=false; SELECT MODE(stringCol), MODE(timestampCol) FROM testTable", + SCHEMA); + assertUnchanged("SELECT MODE(stringCol), MODE(timestampCol) FROM testTable", SCHEMA); + } + + @Test + public void testServerDependentModeDoesNotInitializeOnBroker() { + assertUnchanged("SET autoRewriteAggregationType=true; " + + "SELECT MODE(LOOKUP('baseballTeams', 'teamInteger', 'teamID', stringCol)) FROM testTable", SCHEMA); + } + + @Test + public void testMissingSchemaAndColumns() { + assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(stringCol) FROM testTable", null); + assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(unknownCol) FROM testTable", SCHEMA); + assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(CONCAT(unknownCol, 'suffix')) FROM testTable", + SCHEMA); + assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(mvStringCol) FROM testTable", SCHEMA); + } + + private static void assertUnchanged(String sql, Schema schema) { + PinotQuery original = CalciteSqlParser.compileToPinotQuery(sql); + PinotQuery optimized = original.deepCopy(); + OPTIMIZER.optimize(optimized, schema); + assertEquals(optimized, original); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java index b79c9c1fbce1..0409f421d2b2 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java @@ -21,6 +21,7 @@ import com.google.common.collect.Lists; import java.io.File; import java.io.IOException; +import java.sql.Timestamp; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -31,6 +32,7 @@ import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock; import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; @@ -80,16 +82,28 @@ public class ModeQueriesTest extends BaseQueriesTest { private static final String LONG_NO_DICT_COLUMN = "longNoDictColumn"; private static final String FLOAT_NO_DICT_COLUMN = "floatNoDictColumn"; private static final String DOUBLE_NO_DICT_COLUMN = "doubleNoDictColumn"; + private static final String STRING_COLUMN = "stringColumn"; + private static final String STRING_NO_DICT_COLUMN = "stringNoDictColumn"; + private static final String JSON_COLUMN = "jsonColumn"; + private static final String TIMESTAMP_COLUMN = "timestampColumn"; + private static final String TIMESTAMP_NO_DICT_COLUMN = "timestampNoDictColumn"; + private static final long BASE_TIMESTAMP = Timestamp.valueOf("2026-09-03 10:11:12.000").getTime(); private static final Schema SCHEMA = new Schema.SchemaBuilder().addSingleValueDimension(INT_COLUMN, DataType.INT) .addMultiValueDimension(INT_MV_COLUMN, DataType.INT).addSingleValueDimension(INT_NO_DICT_COLUMN, DataType.INT) .addSingleValueDimension(LONG_COLUMN, DataType.LONG).addSingleValueDimension(LONG_NO_DICT_COLUMN, DataType.LONG) .addSingleValueDimension(FLOAT_COLUMN, DataType.FLOAT) .addSingleValueDimension(FLOAT_NO_DICT_COLUMN, DataType.FLOAT) .addSingleValueDimension(DOUBLE_COLUMN, DataType.DOUBLE) - .addSingleValueDimension(DOUBLE_NO_DICT_COLUMN, DataType.DOUBLE).build(); + .addSingleValueDimension(DOUBLE_NO_DICT_COLUMN, DataType.DOUBLE) + .addSingleValueDimension(STRING_COLUMN, DataType.STRING) + .addSingleValueDimension(STRING_NO_DICT_COLUMN, DataType.STRING) + .addSingleValueDimension(JSON_COLUMN, DataType.STRING) + .addSingleValueDimension(TIMESTAMP_COLUMN, DataType.TIMESTAMP) + .addSingleValueDimension(TIMESTAMP_NO_DICT_COLUMN, DataType.TIMESTAMP).build(); private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) .setNoDictionaryColumns( - Lists.newArrayList(INT_NO_DICT_COLUMN, LONG_NO_DICT_COLUMN, FLOAT_NO_DICT_COLUMN, DOUBLE_NO_DICT_COLUMN)) + Lists.newArrayList(INT_NO_DICT_COLUMN, LONG_NO_DICT_COLUMN, FLOAT_NO_DICT_COLUMN, DOUBLE_NO_DICT_COLUMN, + STRING_NO_DICT_COLUMN, TIMESTAMP_NO_DICT_COLUMN)) .build(); private static final double DELTA = 0.00001; @@ -136,6 +150,11 @@ public void setUp() record.putValue(FLOAT_NO_DICT_COLUMN, (float) value); record.putValue(DOUBLE_COLUMN, (double) value); record.putValue(DOUBLE_NO_DICT_COLUMN, (double) value); + record.putValue(STRING_COLUMN, Integer.toString(value)); + record.putValue(STRING_NO_DICT_COLUMN, Integer.toString(value)); + record.putValue(JSON_COLUMN, "{\"value\":\"" + value + "\"}"); + record.putValue(TIMESTAMP_COLUMN, BASE_TIMESTAMP + value); + record.putValue(TIMESTAMP_NO_DICT_COLUMN, BASE_TIMESTAMP + value); records.add(record); } long maxOccurrences = _values.values().stream().max(Long::compareTo).get(); @@ -371,6 +390,80 @@ public Object[][] testAggregationGroupByMVDataProvider() { return entries.toArray(new Object[0][]); } + @Test + public void testStringAggregationAndComputedExpression() { + long maxOccurrences = _values.values().stream().max(Long::compareTo).orElseThrow(); + String expectedMin = _values.entrySet().stream().filter(e -> e.getValue() == maxOccurrences) + .map(e -> e.getKey().toString()).min(String::compareTo).orElseThrow(); + String expectedMax = _values.entrySet().stream().filter(e -> e.getValue() == maxOccurrences) + .map(e -> e.getKey().toString()).max(String::compareTo).orElseThrow(); + BrokerResponseNative response = getBrokerResponseForOptimizedQuery( + "SET autoRewriteAggregationType=true; SELECT MODE(stringColumn), " + + "MODE(stringNoDictColumn), MODE(stringColumn, 'MAX'), MODE(CONCAT('value-', stringColumn, '')), " + + "MODE(CASE WHEN JSONEXTRACTSCALAR(jsonColumn, '$.value', 'STRING', '') = '' THEN NULL " + + "ELSE JSONEXTRACTSCALAR(jsonColumn, '$.value', 'STRING', '') END) FROM testTable", + SCHEMA); + assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); + assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, + ColumnDataType.STRING}); + assertEquals(response.getResultTable().getRows().size(), 1); + assertEquals(response.getResultTable().getRows().get(0), + new Object[]{expectedMin, expectedMin, expectedMax, "value-" + expectedMin, expectedMin}); + } + + @Test + public void testStringAggregationWithNoMatchingRows() { + BrokerResponseNative response = getBrokerResponseForOptimizedQuery( + "SET autoRewriteAggregationType=true; SELECT MODE(stringColumn), " + + "MODE(stringNoDictColumn) FROM testTable WHERE intColumn < 0", SCHEMA); + assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); + assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING}); + assertEquals(response.getResultTable().getRows().size(), 1); + assertEquals(response.getResultTable().getRows().get(0), new Object[]{null, null}); + } + + @DataProvider + public Object[][] stringGroupByColumns() { + return new Object[][]{{INT_COLUMN}, {INT_MV_COLUMN}}; + } + + @Test(dataProvider = "stringGroupByColumns") + public void testStringAggregationGroupBy(String groupByColumn) { + BrokerResponseNative response = getBrokerResponseForOptimizedQuery( + "SET autoRewriteAggregationType=true; SELECT " + groupByColumn + + ", MODE(stringColumn), MODE(stringNoDictColumn), MODE(CONCAT('value-', stringColumn, '')) " + + "FROM testTable GROUP BY " + groupByColumn + " ORDER BY " + groupByColumn, SCHEMA); + assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); + assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING}); + List rows = response.getResultTable().getRows(); + assertEquals(rows.size(), 10); + for (Object[] row : rows) { + String value = row[0].toString(); + assertEquals(row, new Object[]{row[0], value, value, "value-" + value}); + } + } + + @Test + public void testTimestampAggregationAndResultType() { + BrokerResponseNative response = getBrokerResponseForOptimizedQuery( + "SET autoRewriteAggregationType=true; SELECT MODE(timestampColumn), " + + "MODE(timestampNoDictColumn), fromTimestamp(MODE(timestampColumn)), MODE(toTimestamp(longColumn)) " + + "FROM testTable", SCHEMA); + assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); + assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.TIMESTAMP, ColumnDataType.TIMESTAMP, ColumnDataType.LONG, + ColumnDataType.TIMESTAMP}); + long expectedMillis = BASE_TIMESTAMP + _expectedResultMin.longValue(); + String expectedTimestamp = new Timestamp(expectedMillis).toString(); + assertEquals(response.getResultTable().getRows().size(), 1); + assertEquals(response.getResultTable().getRows().get(0), + new Object[]{expectedTimestamp, expectedTimestamp, expectedMillis, + new Timestamp(_expectedResultMin.longValue()).toString()}); + } + @AfterClass public void tearDown() throws IOException { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java index 6d1033f37aaf..83b372d045bb 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java @@ -44,6 +44,8 @@ /// - MAX(longType) -> MAXLONG /// - SUM(longType) -> SUMLONG /// - SUM(intType) -> SUMINT +/// - MODE(stringType) -> MODESTRING +/// - MODE(timestampType) -> MODETIMESTAMP public class PinotAggregateFunctionRewriteRule extends RelOptRule { public static final PinotAggregateFunctionRewriteRule INSTANCE = new PinotAggregateFunctionRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY, null); @@ -98,6 +100,19 @@ private static AggregateCall maybeRewriteAggCall(AggregateCall call, RelNode inp SqlAggFunction newAgg; switch (aggKind) { + case MODE: { + if (SqlTypeName.STRING_TYPES.contains(operandType)) { + newAgg = new PinotSqlAggFunction("MODESTRING", SqlKind.OTHER_FUNCTION, ReturnTypes.explicit(call.getType()), + aggFunction.getOperandTypeChecker(), SqlFunctionCategory.USER_DEFINED_FUNCTION); + } else if (operandType == SqlTypeName.TIMESTAMP) { + newAgg = new PinotSqlAggFunction("MODETIMESTAMP", SqlKind.OTHER_FUNCTION, + ReturnTypes.explicit(call.getType()), aggFunction.getOperandTypeChecker(), + SqlFunctionCategory.USER_DEFINED_FUNCTION); + } else { + return call; + } + break; + } case MIN: { if (SqlTypeName.STRING_TYPES.contains(operandType)) { newAgg = new PinotSqlAggFunction("MINSTRING", SqlKind.OTHER_FUNCTION, ReturnTypes.explicit(call.getType()), diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java index e05ee827547d..eec0041d47b2 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java @@ -192,7 +192,6 @@ public static PinotOperatorTable instance(boolean nullHandlingEnabled) { PinotMinMaxFunction.MIN, PinotMinMaxFunction.MAX, PinotAvgFunction.INSTANCE, - SqlStdOperatorTable.MODE, SqlStdOperatorTable.STDDEV_POP, SqlStdOperatorTable.COVAR_POP, SqlStdOperatorTable.COVAR_SAMP, diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java new file mode 100644 index 000000000000..b2d4a5a544c3 --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java @@ -0,0 +1,134 @@ +/** + * 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.query.queries; + +import java.util.ArrayList; +import java.util.List; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.QueryEnvironmentTestBase; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; +import org.apache.pinot.query.planner.physical.DispatchableSubPlan; +import org.apache.pinot.query.planner.plannode.AggregateNode; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/// Verifies MODE type inference and type-specific dispatch across both multi-stage planner implementations. +public class ModeSqlPlannerTest extends QueryEnvironmentTestBase { + @DataProvider + public Object[][] physicalOptimizers() { + return new Object[][]{{false}, {true}}; + } + + @Test + public void testModeReturnTypes() { + RelDataType rowType = _queryEnvironment.compile( + "SET autoRewriteAggregationType=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3), MODE(col7), " + + "MODE(CAST(col3 AS FLOAT)), MODE(CAST(col3 AS DOUBLE)), " + + "MODE(NULLIF(JSONEXTRACTSCALAR(col1, '$.user', 'STRING', ''), '')), " + + "fromTimestamp(MODE(ts_timestamp)) FROM a") + .getRelRoot().validatedRowType; + SqlTypeName[] expectedTypes = {SqlTypeName.VARCHAR, SqlTypeName.TIMESTAMP, SqlTypeName.DOUBLE, + SqlTypeName.DOUBLE, SqlTypeName.DOUBLE, SqlTypeName.DOUBLE, SqlTypeName.VARCHAR, SqlTypeName.BIGINT}; + for (int i = 0; i < expectedTypes.length; i++) { + assertEquals(rowType.getFieldList().get(i).getType().getSqlTypeName(), expectedTypes[i]); + } + } + + @Test + public void testFilteredModeNullability() { + RelDataType rowType = _queryEnvironment.compile( + "SET autoRewriteAggregationType=true; " + + "SELECT col2, MODE(col1) FILTER (WHERE col3 > 0), MODE(ts_timestamp) FILTER (WHERE col3 > 0) " + + "FROM a GROUP BY col2").getRelRoot().validatedRowType; + assertTrue(rowType.getFieldList().get(1).getType().isNullable()); + assertTrue(rowType.getFieldList().get(2).getType().isNullable()); + } + + @Test(dataProvider = "physicalOptimizers") + public void testDistributedModeTypes(boolean usePhysicalOptimizer) { + DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + + "SET autoRewriteAggregationType=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3) FROM a"); + PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); + assertEquals(root.getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP, ColumnDataType.DOUBLE}); + + List aggregates = findAggregates(plan); + assertFalse(aggregates.isEmpty()); + boolean sawIntermediate = false; + boolean sawFinal = false; + for (AggregateNode aggregate : aggregates) { + List calls = aggregate.getAggCalls(); + assertEquals(calls.stream().map(RexExpression.FunctionCall::getFunctionName).toList(), + List.of("MODESTRING", "MODETIMESTAMP", "MODE")); + if (aggregate.getAggType().isOutputIntermediateFormat() && !aggregate.isLeafReturnFinalResult()) { + sawIntermediate = true; + assertEquals(aggregate.getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.OBJECT, ColumnDataType.OBJECT, ColumnDataType.OBJECT}); + } else { + sawFinal = true; + assertEquals(aggregate.getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP, ColumnDataType.DOUBLE}); + } + } + assertTrue(sawIntermediate, "Distributed MODE must exchange frequency counts"); + assertTrue(sawFinal, "Distributed MODE must retain its final result types"); + } + + @Test(dataProvider = "physicalOptimizers") + public void testModeExpressionsAndTieBreakers(boolean usePhysicalOptimizer) { + DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + + "SET autoRewriteAggregationType=true; " + + "SELECT col2, MODE(NULLIF(col1, ''), 'MIN'), MODE(ts_timestamp, 'MAX'), MODE(col3, 'AVG') " + + "FROM a GROUP BY col2"); + PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); + assertEquals(root.getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.TIMESTAMP, + ColumnDataType.DOUBLE}); + for (AggregateNode aggregate : findAggregates(plan)) { + assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), + List.of("MODESTRING", "MODETIMESTAMP", "MODE")); + } + } + + private static List findAggregates(DispatchableSubPlan plan) { + List aggregates = new ArrayList<>(); + for (DispatchablePlanFragment fragment : plan.getQueryStageMap().values()) { + findAggregates(fragment.getPlanFragment().getFragmentRoot(), aggregates); + } + return aggregates; + } + + private static void findAggregates(PlanNode node, List aggregates) { + if (node instanceof AggregateNode) { + aggregates.add((AggregateNode) node); + } + for (PlanNode input : node.getInputs()) { + findAggregates(input, aggregates); + } + } +} diff --git a/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json new file mode 100644 index 000000000000..0976e9a830ea --- /dev/null +++ b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json @@ -0,0 +1,150 @@ +{ + "mode_string_timestamp_distributed": { + "comments": "Exercise typed MODE partial-state serialization across four partitions on two servers. The most frequent pallet and action user in group a never occur on the same input row.", + "extraProps": { + "enableColumnBasedNullHandling": true + }, + "tables": { + "items": { + "schema": [ + {"name": "lpn_id", "type": "STRING"}, + {"name": "pallet_id", "type": "STRING"}, + {"name": "created_on", "type": "TIMESTAMP"}, + {"name": "meta", "type": "STRING"}, + {"name": "first_name", "type": "STRING"}, + {"name": "last_name", "type": "STRING"}, + {"name": "manifest_id", "type": "INT"} + ], + "inputs": [ + ["a", "p2", "2026-09-03 08:00:00.001", "{\"action_user\":\"u2\"}", "Bob", "Smith", 1], + ["a", "p2", "2026-09-03 08:00:00.001", "{\"action_user\":\"u3\"}", "Cara", "Jones", 1], + ["a", "p2", "2026-09-03 09:00:00.123", "{\"action_user\":\"u2\"}", "Bob", "Smith", 2], + ["a", "p1", "2026-09-03 09:00:00.123", "{\"action_user\":\"u1\"}", "Alice", "Jones", 2], + ["a", "p1", "2026-09-03 09:00:00.123", "{\"action_user\":\"u1\"}", "Alice", "Jones", 1], + ["a", "p3", "2026-09-03 10:00:00.999", "{\"action_user\":\"u1\"}", "Alice", "Jones", 2], + ["nulls", null, null, "{}", null, null, 99], + ["nulls", null, null, "{\"action_user\":\"\"}", null, null, 99], + ["ties", "z", "1970-01-01 00:00:00.001", "{}", null, null, 99], + ["ties", "a", "1969-12-31 23:59:59.999", "{}", null, null, 99], + ["ties", "z", "1970-01-01 00:00:00.001", "{}", null, null, 99], + ["ties", "a", "1969-12-31 23:59:59.999", "{}", null, null, 99], + ["empty", "", "2026-09-03 00:00:00.001", "{}", null, null, 99], + ["empty", "z", null, "{}", null, null, 99], + ["empty", "", null, "{}", null, null, 99] + ] + }, + "manifests": { + "schema": [ + {"name": "id", "type": "INT"}, + {"name": "code", "type": "STRING"}, + {"name": "status", "type": "STRING"}, + {"name": "created_on", "type": "TIMESTAMP"} + ], + "inputs": [ + [1, "Z", "shipped", "2026-09-03 12:00:00.999"], + [2, "A", "pending", "2026-09-03 11:00:00.123"] + ] + } + }, + "queries": [ + { + "description": "Global string and timestamp modes preserve the argument types", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'a'", + "outputs": [["p2", "2026-09-03 09:00:00.123"]] + }, + { + "description": "Grouped modes ignore nulls, retain empty strings, and break ties by the smallest value", + "sql": "SET autoRewriteAggregationType=true; SELECT lpn_id, MODE(pallet_id), MODE(created_on) FROM {items} GROUP BY lpn_id", + "outputs": [ + ["a", "p2", "2026-09-03 09:00:00.123"], + ["nulls", null, null], + ["ties", "a", "1969-12-31 23:59:59.999"], + ["empty", "", "2026-09-03 00:00:00.001"] + ] + }, + { + "description": "MODE applies independently to a stored string and computed JSON string", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'a'", + "outputs": [["p2", "u1"]] + }, + { + "description": "Typed MAX tie reducers select the largest original value", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id, 'MAX'), MODE(created_on, 'MAX') FROM {items} WHERE lpn_id = 'ties'", + "outputs": [["z", "1970-01-01 00:00:00.001"]] + }, + { + "description": "Existing numeric modes retain DOUBLE results and MIN, MAX, and AVG reducers", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(manifest_id), MODE(manifest_id, 'MAX'), MODE(manifest_id, 'AVG') FROM {items} WHERE lpn_id = 'a'", + "outputs": [[1.0, 2.0, 1.5]] + }, + { + "description": "Computed null inputs do not become a string mode", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'nulls'", + "outputs": [[null]] + }, + { + "description": "String expressions and post-aggregation transforms use the string result type", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(CONCAT(first_name, last_name, ' ')), UPPER(MODE(pallet_id)) FROM {items} WHERE lpn_id = 'a'", + "outputs": [["Alice Jones", "P2"]] + }, + { + "description": "Empty aggregate inputs return null for both types", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'missing'", + "outputs": [[null, null]] + }, + { + "description": "Filtered typed modes return null while preserving their group", + "sql": "SET autoRewriteAggregationType=true; SELECT lpn_id, MODE(pallet_id) FILTER (WHERE manifest_id < 0), MODE(created_on) FILTER (WHERE manifest_id < 0) FROM {items} WHERE lpn_id = 'a' GROUP BY lpn_id", + "outputs": [["a", null, null]] + }, + { + "description": "Timestamp mode can feed a parent comparison; fixture ingestion uses Los Angeles time and the SQL literal uses UTC", + "sql": "SET autoRewriteAggregationType=true; SELECT lpn_id FROM (SELECT lpn_id, MODE(created_on) AS mode_time FROM {items} GROUP BY lpn_id) WHERE mode_time = TIMESTAMP '2026-09-03 16:00:00.123'", + "outputs": [["a"]] + }, + { + "description": "Joined grouping supports independent string, JSON, and timestamp modes, including unmatched outer rows", + "sql": "SET autoRewriteAggregationType=true; SELECT i.lpn_id, MODE(i.pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(i.meta, '$.action_user', 'STRING', ''), '')), MODE(m.code), MODE(m.status), MODE(m.created_on) FROM {items} i LEFT JOIN {manifests} m ON m.id = i.manifest_id GROUP BY i.lpn_id", + "outputs": [ + ["a", "p2", "u1", "A", "pending", "2026-09-03 11:00:00.123"], + ["nulls", null, null, null, null, null], + ["ties", "a", null, null, null, null], + ["empty", "", null, null, null, null] + ] + } + ] + }, + "mode_string_timestamp_direct": { + "comments": "Replicated input exercises AGGREGATE_DIRECT finalization when the physical optimizer colocates a global aggregate on one server.", + "extraProps": { + "enableColumnBasedNullHandling": true + }, + "tables": { + "items": { + "replicated": true, + "schema": [ + {"name": "pallet_id", "type": "STRING"}, + {"name": "created_on", "type": "TIMESTAMP"} + ], + "inputs": [ + ["p2", "2026-09-03 08:00:00.001"], + ["p2", "2026-09-03 09:00:00.123"], + ["p1", "2026-09-03 09:00:00.123"], + [null, null] + ] + } + }, + "queries": [ + { + "description": "Direct aggregate emits scalar string and timestamp values instead of frequency-map intermediates", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items}", + "outputs": [["p2", "2026-09-03 09:00:00.123"]] + }, + { + "description": "Direct aggregate with no matching values emits typed nulls", + "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id) FILTER (WHERE pallet_id = 'missing'), MODE(created_on) FILTER (WHERE pallet_id = 'missing') FROM {items}", + "outputs": [[null, null]] + } + ] + } +} diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java index 080138b3a327..ebbb4551774f 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java @@ -63,7 +63,16 @@ public enum AggregationFunctionType { SUMLONG("sumLong", ReturnTypes.AGG_SUM, OperandTypes.or(OperandTypes.INTEGER, OperandTypes.ARRAY_OF_INTEGER)), SUMPRECISION("sumPrecision", ReturnTypes.explicit(SqlTypeName.DECIMAL), OperandTypes.ANY, SqlTypeName.OTHER), AVG("avg", SqlTypeName.OTHER, SqlTypeName.DOUBLE), - MODE("mode", SqlTypeName.OTHER, SqlTypeName.DOUBLE), + MODE("mode", new ModeReturnTypeInference(), OperandTypes.or( + OperandTypes.NUMERIC, OperandTypes.CHARACTER, OperandTypes.TIMESTAMP, + OperandTypes.family(List.of(SqlTypeFamily.NUMERIC, SqlTypeFamily.CHARACTER), i -> i == 1), + OperandTypes.family(List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER), i -> i == 1), + OperandTypes.family(List.of(SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER), i -> i == 1)), + ReturnTypes.explicit(SqlTypeName.OTHER), null, SqlKind.MODE), + MODESTRING("modeString", ReturnTypes.ARG0_NULLABLE_IF_EMPTY, + OperandTypes.family(List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER), i -> i == 1), SqlTypeName.OTHER), + MODETIMESTAMP("modeTimestamp", ReturnTypes.ARG0_NULLABLE_IF_EMPTY, + OperandTypes.family(List.of(SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER), i -> i == 1), SqlTypeName.OTHER), ANYVALUE("anyValue", ReturnTypes.ARG0, OperandTypes.ANY, SqlTypeName.OTHER), FIRSTWITHTIME("firstWithTime", ReturnTypes.ARG0, OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER), SqlTypeName.OTHER), @@ -413,6 +422,20 @@ public static AggregationFunctionType getAggregationFunctionType(String function } } + /// Preserves the legacy DOUBLE result for numeric MODE while retaining the type of string and timestamp inputs. + private static class ModeReturnTypeInference implements SqlReturnTypeInference { + @Override + public RelDataType inferReturnType(SqlOperatorBinding opBinding) { + RelDataType operandType = opBinding.getOperandType(0); + if (SqlTypeName.STRING_TYPES.contains(operandType.getSqlTypeName()) + || operandType.getSqlTypeName() == SqlTypeName.TIMESTAMP) { + return ReturnTypes.ARG0_NULLABLE_IF_EMPTY.inferReturnType(opBinding); + } + return opBinding.getTypeFactory().createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(SqlTypeName.DOUBLE), true); + } + } + private static class ArrayReturnTypeInference implements SqlReturnTypeInference { final SqlTypeName _sqlTypeName; From f6e38d614565df948c1240535554cf3974b9971c Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 7 Sep 2026 02:32:19 -0700 Subject: [PATCH 2/8] Avoid boxing when reducing typed MODE frequency states --- ...BaseComparableModeAggregationFunction.java | 10 ++-- .../ModeStringAggregationFunction.java | 46 ++++++++++++++++++- .../ModeTimestampAggregationFunction.java | 46 ++++++++++++++++++- ...ModeNonNumericAggregationFunctionTest.java | 30 ++++++++++++ 4 files changed, 125 insertions(+), 7 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java index c8553d7e4d56..48f6c4bce9f2 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java @@ -73,7 +73,11 @@ protected interface ValueCounter { void add(Map counts, int row); } - protected abstract T dictionaryValue(Dictionary dictionary, int dictionaryId); + protected abstract void putDictionaryCount(Map counts, Dictionary dictionary, int dictionaryId, long count); + + protected final boolean isMinimum() { + return _minimum; + } @Override public AggregationResultHolder createAggregationResultHolder() { @@ -179,8 +183,8 @@ private Map extractCounts(@Nullable Object result) { if (result instanceof DictionaryCounts) { DictionaryCounts dictionaryCounts = (DictionaryCounts) result; Map counts = newValueMap(); - dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> counts.put( - dictionaryValue(dictionaryCounts._dictionary, entry.getIntKey()), entry.getLongValue())); + dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> putDictionaryCount( + counts, dictionaryCounts._dictionary, entry.getIntKey(), entry.getLongValue())); return counts; } return (Map) result; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java index fc2805fb9eb8..93b1a145114c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java @@ -18,9 +18,13 @@ */ package org.apache.pinot.core.query.aggregation.function; +import it.unimi.dsi.fastutil.objects.Object2LongMap; +import it.unimi.dsi.fastutil.objects.Object2LongMaps; import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap; +import it.unimi.dsi.fastutil.objects.ObjectIterator; import java.util.List; import java.util.Map; +import javax.annotation.Nullable; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; @@ -59,8 +63,46 @@ protected ValueCounter valueCounter(BlockValSet blockValSet) { } @Override - protected String dictionaryValue(Dictionary dictionary, int dictionaryId) { - return dictionary.getStringValue(dictionaryId); + protected void putDictionaryCount(Map counts, Dictionary dictionary, int dictionaryId, long count) { + ((StringModeCounts) counts).put(dictionary.getStringValue(dictionaryId), count); + } + + @Override + public Map merge(Map left, Map right) { + if (left instanceof Object2LongOpenHashMap && right instanceof Object2LongMap) { + Object2LongOpenHashMap counts = (Object2LongOpenHashMap) left; + ObjectIterator> iterator = + Object2LongMaps.fastIterator((Object2LongMap) right); + while (iterator.hasNext()) { + Object2LongMap.Entry entry = iterator.next(); + counts.addTo(entry.getKey(), entry.getLongValue()); + } + return left; + } + return super.merge(left, right); + } + + @Nullable + @Override + public String extractFinalResult(@Nullable Map counts) { + if (!(counts instanceof Object2LongMap)) { + return super.extractFinalResult(counts); + } + String mode = null; + long maxCount = 0; + ObjectIterator> iterator = + Object2LongMaps.fastIterator((Object2LongMap) counts); + while (iterator.hasNext()) { + Object2LongMap.Entry entry = iterator.next(); + String value = entry.getKey(); + long count = entry.getLongValue(); + if (mode == null || count > maxCount || (count == maxCount + && (isMinimum() ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { + mode = value; + maxCount = count; + } + } + return mode; } @Override diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java index 959971de51f1..0cd9fe996e79 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java @@ -19,9 +19,12 @@ package org.apache.pinot.core.query.aggregation.function; import it.unimi.dsi.fastutil.longs.Long2LongMap; +import it.unimi.dsi.fastutil.longs.Long2LongMaps; import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; +import it.unimi.dsi.fastutil.objects.ObjectIterator; import java.util.List; import java.util.Map; +import javax.annotation.Nullable; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; @@ -60,8 +63,47 @@ protected ValueCounter valueCounter(BlockValSet blockValSet) { } @Override - protected Long dictionaryValue(Dictionary dictionary, int dictionaryId) { - return dictionary.getLongValue(dictionaryId); + protected void putDictionaryCount(Map counts, Dictionary dictionary, int dictionaryId, long count) { + ((Long2LongOpenHashMap) counts).put(dictionary.getLongValue(dictionaryId), count); + } + + @Override + public Map merge(Map left, Map right) { + if (!(left instanceof Long2LongOpenHashMap) || !(right instanceof Long2LongMap)) { + return super.merge(left, right); + } + Long2LongOpenHashMap counts = (Long2LongOpenHashMap) left; + ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) right); + while (iterator.hasNext()) { + Long2LongMap.Entry entry = iterator.next(); + counts.addTo(entry.getLongKey(), entry.getLongValue()); + } + return counts; + } + + @Nullable + @Override + public Long extractFinalResult(@Nullable Map counts) { + if (!(counts instanceof Long2LongMap)) { + return super.extractFinalResult(counts); + } + ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) counts); + if (!iterator.hasNext()) { + return null; + } + Long2LongMap.Entry first = iterator.next(); + long mode = first.getLongKey(); + long maxCount = first.getLongValue(); + while (iterator.hasNext()) { + Long2LongMap.Entry entry = iterator.next(); + long value = entry.getLongKey(); + long count = entry.getLongValue(); + if (count > maxCount || (count == maxCount && (isMinimum() ? value < mode : value > mode))) { + mode = value; + maxCount = count; + } + } + return mode; } @Override diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java index 2e54aa61ea66..c6d7a6d27531 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java @@ -18,7 +18,9 @@ */ package org.apache.pinot.core.query.aggregation.function; +import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; import java.nio.ByteBuffer; +import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.pinot.common.CustomObject; @@ -37,6 +39,7 @@ import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; import static org.testng.Assert.expectThrows; @@ -78,6 +81,31 @@ public void testTimestampIntermediateResultsPreserveLongPrecision() { assertEquals(maxFunction.extractFinalResult(merged), Long.valueOf(later)); } + @Test + public void testTimestampMergesPrimitiveAndGenericStates() { + ModeTimestampAggregationFunction minFunction = + new ModeTimestampAggregationFunction(List.of(EXPRESSION), true); + ModeTimestampAggregationFunction maxFunction = + new ModeTimestampAggregationFunction(argumentsWithReducer("MAX"), true); + for (boolean primitiveLeft : new boolean[]{false, true}) { + for (boolean primitiveRight : new boolean[]{false, true}) { + Map left = primitiveLeft ? new Long2LongOpenHashMap() : new HashMap<>(); + Map right = primitiveRight ? new Long2LongOpenHashMap() : new HashMap<>(); + left.put(Long.MIN_VALUE, 2L); + left.put(0L, 1L); + right.put(Long.MIN_VALUE, 1L); + right.put(Long.MAX_VALUE, 3L); + + Map merged = minFunction.merge(left, right); + assertSame(merged, left); + assertEquals(merged, Map.of(Long.MIN_VALUE, 3L, 0L, 1L, Long.MAX_VALUE, 3L)); + assertEquals(minFunction.extractFinalResult(merged), Long.valueOf(Long.MIN_VALUE)); + assertEquals(maxFunction.extractFinalResult(merged), Long.valueOf(Long.MAX_VALUE)); + assertEquals(right, Map.of(Long.MIN_VALUE, 1L, Long.MAX_VALUE, 3L)); + } + } + } + @Test public void testStringModeSkipsNullRowsForMultiValueGroupKeys() { ModeStringAggregationFunction function = new ModeStringAggregationFunction(List.of(EXPRESSION), true); @@ -101,6 +129,8 @@ public void testEmptyResultsAreNullWithEitherNullHandlingMode() { new ModeTimestampAggregationFunction(List.of(EXPRESSION), nullHandlingEnabled); assertNull(stringFunction.extractFinalResult(null)); assertNull(timestampFunction.extractFinalResult(null)); + assertNull(timestampFunction.extractFinalResult(new Long2LongOpenHashMap())); + assertNull(timestampFunction.extractFinalResult(Map.of())); assertNull(stringFunction.extractFinalResult( stringFunction.extractAggregationResult(stringFunction.createAggregationResultHolder()))); assertNull(timestampFunction.extractFinalResult( From 8cb1b17fb20c08accfc666692ceb9c78712112ab Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 7 Sep 2026 02:49:49 -0700 Subject: [PATCH 3/8] Gate typed MODE rewrites behind an independent rollout option --- ...deAggregationFunctionRewriteOptimizer.java | 4 +- ...gregationFunctionRewriteOptimizerTest.java | 32 ++++--- .../apache/pinot/queries/ModeQueriesTest.java | 8 +- .../PinotAggregateFunctionRewriteRule.java | 15 ---- ...notModeAggregationFunctionRewriteRule.java | 88 +++++++++++++++++++ .../calcite/rel/rules/PinotQueryRuleSets.java | 2 + .../apache/pinot/query/QueryEnvironment.java | 17 +++- .../query/queries/ModeSqlPlannerTest.java | 86 +++++++++++++++++- .../resources/queries/ModeAggregates.json | 26 +++--- .../pinot/spi/utils/CommonConstants.java | 6 ++ 10 files changed, 233 insertions(+), 51 deletions(-) create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java index 9fb460d33ba7..9495536df477 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java @@ -38,7 +38,7 @@ import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; -/// When `autoRewriteAggregationType` is enabled, resolves string and timestamp MODE expressions to implementations +/// When `enableTypedMode` is enabled, resolves string and timestamp MODE expressions to implementations /// with fixed result types before execution. /// This also supplies the broker with the correct result type when no rows match or groups are trimmed before /// finalization. Type inference reads schema and function metadata only: server-dependent transforms such as LOOKUP @@ -48,7 +48,7 @@ public class ModeAggregationFunctionRewriteOptimizer implements StatementOptimiz public void optimize(PinotQuery pinotQuery, @Nullable Schema schema) { // Keep existing timestamp MODE requests compatible with older servers during rolling upgrades. if (schema == null || pinotQuery.getQueryOptions() == null - || !Boolean.parseBoolean(pinotQuery.getQueryOptions().get(QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE))) { + || !Boolean.parseBoolean(pinotQuery.getQueryOptions().get(QueryOptionKey.ENABLE_TYPED_MODE))) { return; } rewriteExpressions(pinotQuery.getSelectList(), schema); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java index 259f31cd9b18..8a28bbcdf7ac 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java @@ -63,7 +63,7 @@ public Object[][] modeExpressions() { public void testModeExpressions(String original, String rewritten) { for (boolean nullHandlingEnabled : new boolean[]{false, true}) { String prefix = - "SET autoRewriteAggregationType=true; SET enableNullHandling=" + nullHandlingEnabled + "; SELECT "; + "SET enableTypedMode=true; SET enableNullHandling=" + nullHandlingEnabled + "; SELECT "; TestHelper.assertEqualsQuery(prefix + original + " AS commonValue FROM testTable", prefix + rewritten + " AS commonValue FROM testTable", SCHEMA); } @@ -72,11 +72,11 @@ public void testModeExpressions(String original, String rewritten) { @Test public void testModeInHavingAndOrderBy() { TestHelper.assertEqualsQuery( - "SET autoRewriteAggregationType=true; " + "SET enableTypedMode=true; " + "SELECT intCol, MODE(stringCol) AS commonValue FROM testTable GROUP BY intCol " + "HAVING MODE(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END) = 'value' " + "ORDER BY MODE(timestampCol) DESC", - "SET autoRewriteAggregationType=true; " + "SET enableTypedMode=true; " + "SELECT intCol, modeString(stringCol) AS commonValue FROM testTable GROUP BY intCol " + "HAVING modeString(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END) = 'value' " + "ORDER BY modeTimestamp(timestampCol) DESC", SCHEMA); @@ -84,27 +84,39 @@ public void testModeInHavingAndOrderBy() { @Test public void testNumericModeAndOptInRewritesRemainUnchanged() { - assertUnchanged("SELECT MODE(intCol), MODE(longCol), MODE(floatCol), MODE(doubleCol), " + assertUnchanged("SET enableTypedMode=true; SELECT MODE(intCol), MODE(longCol), MODE(floatCol), MODE(doubleCol), " + "MODE(CAST(stringCol AS LONG)), MODE(fromDateTime(stringCol, 'yyyy-MM-dd HH:mm:ss')), " + "MIN(stringCol), MAX(longCol), SUM(intCol) FROM testTable", SCHEMA); - assertUnchanged("SET autoRewriteAggregationType=false; SELECT MODE(stringCol), MODE(timestampCol) FROM testTable", + assertUnchanged("SET enableTypedMode=false; SELECT MODE(stringCol), MODE(timestampCol) FROM testTable", SCHEMA); assertUnchanged("SELECT MODE(stringCol), MODE(timestampCol) FROM testTable", SCHEMA); } @Test - public void testServerDependentModeDoesNotInitializeOnBroker() { + public void testExistingRewriteOptionPreservesLegacyMode() { assertUnchanged("SET autoRewriteAggregationType=true; " + + "SELECT MODE(timestampCol), MODE(timestampCol, 'AVG'), MODE(stringCol) FROM testTable", SCHEMA); + assertUnchanged("SET autoRewriteAggregationType=true; SET enableTypedMode=false; " + + "SELECT MODE(timestampCol), MODE(timestampCol, 'AVG'), MODE(stringCol) FROM testTable", SCHEMA); + TestHelper.assertEqualsQuery("SET autoRewriteAggregationType=true; SET enableTypedMode=true; " + + "SELECT MODE(timestampCol), MODE(stringCol), MIN(stringCol) FROM testTable", + "SET autoRewriteAggregationType=true; SET enableTypedMode=true; " + + "SELECT MODETIMESTAMP(timestampCol), MODESTRING(stringCol), MINSTRING(stringCol) FROM testTable", SCHEMA); + } + + @Test + public void testServerDependentModeDoesNotInitializeOnBroker() { + assertUnchanged("SET enableTypedMode=true; " + "SELECT MODE(LOOKUP('baseballTeams', 'teamInteger', 'teamID', stringCol)) FROM testTable", SCHEMA); } @Test public void testMissingSchemaAndColumns() { - assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(stringCol) FROM testTable", null); - assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(unknownCol) FROM testTable", SCHEMA); - assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(CONCAT(unknownCol, 'suffix')) FROM testTable", + assertUnchanged("SET enableTypedMode=true; SELECT MODE(stringCol) FROM testTable", null); + assertUnchanged("SET enableTypedMode=true; SELECT MODE(unknownCol) FROM testTable", SCHEMA); + assertUnchanged("SET enableTypedMode=true; SELECT MODE(CONCAT(unknownCol, 'suffix')) FROM testTable", SCHEMA); - assertUnchanged("SET autoRewriteAggregationType=true; SELECT MODE(mvStringCol) FROM testTable", SCHEMA); + assertUnchanged("SET enableTypedMode=true; SELECT MODE(mvStringCol) FROM testTable", SCHEMA); } private static void assertUnchanged(String sql, Schema schema) { diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java index 0409f421d2b2..4641a9ba9945 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java @@ -398,7 +398,7 @@ public void testStringAggregationAndComputedExpression() { String expectedMax = _values.entrySet().stream().filter(e -> e.getValue() == maxOccurrences) .map(e -> e.getKey().toString()).max(String::compareTo).orElseThrow(); BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET autoRewriteAggregationType=true; SELECT MODE(stringColumn), " + "SET enableTypedMode=true; SELECT MODE(stringColumn), " + "MODE(stringNoDictColumn), MODE(stringColumn, 'MAX'), MODE(CONCAT('value-', stringColumn, '')), " + "MODE(CASE WHEN JSONEXTRACTSCALAR(jsonColumn, '$.value', 'STRING', '') = '' THEN NULL " + "ELSE JSONEXTRACTSCALAR(jsonColumn, '$.value', 'STRING', '') END) FROM testTable", @@ -415,7 +415,7 @@ public void testStringAggregationAndComputedExpression() { @Test public void testStringAggregationWithNoMatchingRows() { BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET autoRewriteAggregationType=true; SELECT MODE(stringColumn), " + "SET enableTypedMode=true; SELECT MODE(stringColumn), " + "MODE(stringNoDictColumn) FROM testTable WHERE intColumn < 0", SCHEMA); assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), @@ -432,7 +432,7 @@ public Object[][] stringGroupByColumns() { @Test(dataProvider = "stringGroupByColumns") public void testStringAggregationGroupBy(String groupByColumn) { BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET autoRewriteAggregationType=true; SELECT " + groupByColumn + "SET enableTypedMode=true; SELECT " + groupByColumn + ", MODE(stringColumn), MODE(stringNoDictColumn), MODE(CONCAT('value-', stringColumn, '')) " + "FROM testTable GROUP BY " + groupByColumn + " ORDER BY " + groupByColumn, SCHEMA); assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); @@ -449,7 +449,7 @@ public void testStringAggregationGroupBy(String groupByColumn) { @Test public void testTimestampAggregationAndResultType() { BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET autoRewriteAggregationType=true; SELECT MODE(timestampColumn), " + "SET enableTypedMode=true; SELECT MODE(timestampColumn), " + "MODE(timestampNoDictColumn), fromTimestamp(MODE(timestampColumn)), MODE(toTimestamp(longColumn)) " + "FROM testTable", SCHEMA); assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java index 83b372d045bb..6d1033f37aaf 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateFunctionRewriteRule.java @@ -44,8 +44,6 @@ /// - MAX(longType) -> MAXLONG /// - SUM(longType) -> SUMLONG /// - SUM(intType) -> SUMINT -/// - MODE(stringType) -> MODESTRING -/// - MODE(timestampType) -> MODETIMESTAMP public class PinotAggregateFunctionRewriteRule extends RelOptRule { public static final PinotAggregateFunctionRewriteRule INSTANCE = new PinotAggregateFunctionRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY, null); @@ -100,19 +98,6 @@ private static AggregateCall maybeRewriteAggCall(AggregateCall call, RelNode inp SqlAggFunction newAgg; switch (aggKind) { - case MODE: { - if (SqlTypeName.STRING_TYPES.contains(operandType)) { - newAgg = new PinotSqlAggFunction("MODESTRING", SqlKind.OTHER_FUNCTION, ReturnTypes.explicit(call.getType()), - aggFunction.getOperandTypeChecker(), SqlFunctionCategory.USER_DEFINED_FUNCTION); - } else if (operandType == SqlTypeName.TIMESTAMP) { - newAgg = new PinotSqlAggFunction("MODETIMESTAMP", SqlKind.OTHER_FUNCTION, - ReturnTypes.explicit(call.getType()), aggFunction.getOperandTypeChecker(), - SqlFunctionCategory.USER_DEFINED_FUNCTION); - } else { - return call; - } - break; - } case MIN: { if (SqlTypeName.STRING_TYPES.contains(operandType)) { newAgg = new PinotSqlAggFunction("MINSTRING", SqlKind.OTHER_FUNCTION, ReturnTypes.explicit(call.getType()), diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java new file mode 100644 index 000000000000..6886af299b8d --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java @@ -0,0 +1,88 @@ +/** + * 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.List; +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.logical.LogicalAggregate; +import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.pinot.common.function.sql.PinotSqlAggFunction; + + +/// Rewrites string and timestamp MODE calls to typed implementations after an explicit rollout opt-in. +/// Numeric MODE and the separate MIN, MAX and SUM rewrite rule are unaffected. +public class PinotModeAggregationFunctionRewriteRule extends RelOptRule { + public static PinotModeAggregationFunctionRewriteRule instanceWithDescription(String description) { + return new PinotModeAggregationFunctionRewriteRule(description); + } + + private PinotModeAggregationFunctionRewriteRule(String description) { + super(operand(LogicalAggregate.class, any()), PinotRuleUtils.PINOT_REL_FACTORY, description); + } + + @Override + public void onMatch(RelOptRuleCall call) { + Aggregate aggregate = call.rel(0); + RelNode input = aggregate.getInput(); + List originalCalls = aggregate.getAggCallList(); + List rewrittenCalls = new ArrayList<>(originalCalls.size()); + boolean changed = false; + for (AggregateCall originalCall : originalCalls) { + AggregateCall rewrittenCall = maybeRewriteAggCall(originalCall, input, aggregate.getGroupCount()); + changed |= rewrittenCall != originalCall; + rewrittenCalls.add(rewrittenCall); + } + if (changed) { + call.transformTo(aggregate.copy(aggregate.getTraitSet(), input, aggregate.getGroupSet(), aggregate.getGroupSets(), + rewrittenCalls)); + } + } + + private static AggregateCall maybeRewriteAggCall(AggregateCall call, RelNode input, int numGroups) { + SqlAggFunction aggregation = call.getAggregation(); + List arguments = call.getArgList(); + if (aggregation.getKind() != SqlKind.MODE || arguments.isEmpty()) { + return call; + } + SqlTypeName operandType = input.getRowType().getFieldList().get(arguments.get(0)).getType().getSqlTypeName(); + String functionName; + if (SqlTypeName.STRING_TYPES.contains(operandType)) { + functionName = "MODESTRING"; + } else if (operandType == SqlTypeName.TIMESTAMP) { + functionName = "MODETIMESTAMP"; + } else { + return call; + } + SqlAggFunction rewrittenAggregation = new PinotSqlAggFunction(functionName, SqlKind.OTHER_FUNCTION, + ReturnTypes.explicit(call.getType()), aggregation.getOperandTypeChecker(), + SqlFunctionCategory.USER_DEFINED_FUNCTION); + return AggregateCall.create(rewrittenAggregation, call.isDistinct(), call.isApproximate(), call.ignoreNulls(), + arguments, call.filterArg, call.distinctKeys, call.getCollation(), numGroups, input, call.getType(), + call.getName()); + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java index 38822d4dfad7..7e204b3fa076 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java @@ -174,6 +174,8 @@ private PinotQueryRuleSets() { PinotAggregateFunctionRewriteRule .instanceWithDescription(PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE), + PinotModeAggregationFunctionRewriteRule + .instanceWithDescription(PlannerRuleNames.TYPED_MODE_REWRITE), // convert CASE-style filtered aggregates into true filtered aggregates // put it after AGGREGATE_REDUCE_FUNCTIONS where SUM is converted to SUM0 diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index bfa24cac481e..61e6de8f5914 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -167,7 +167,7 @@ public QueryEnvironment(Config config, MultiClusterRoutingContext multiClusterRo rootSchema, List.of(database), _typeFactory, CONNECTION_CONFIG, config.isCaseSensitive()); _defaultDisabledPlannerRules = _envConfig.defaultDisabledPlannerRules(); // default optProgram with no skip rule options and no use rule options - _optProgram = getOptProgram(_envConfig.getRuleSet(), Set.of(), Set.of(), _defaultDisabledPlannerRules); + _optProgram = getOptProgram(_envConfig.getRuleSet(), Set.of(), Set.of(), _defaultDisabledPlannerRules, false); _multiClusterRoutingContext = multiClusterRoutingContext; } @@ -199,11 +199,16 @@ private PlannerContext getPlannerContext(SqlNodeAndOptions sqlNodeAndOptions) { if (Boolean.parseBoolean(options.get(QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE))) { useRuleSet.add(CommonConstants.Broker.PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE); } + boolean enableTypedMode = Boolean.parseBoolean(options.get(QueryOptionKey.ENABLE_TYPED_MODE)); + if (enableTypedMode) { + useRuleSet.add(CommonConstants.Broker.PlannerRuleNames.TYPED_MODE_REWRITE); + } if (MapUtils.isNotEmpty(options)) { Set skipRuleSet = QueryOptionsUtils.getSkipPlannerRules(options); if (!skipRuleSet.isEmpty() || !useRuleSet.isEmpty()) { // dynamically create optProgram according to rule options - optProgram = getOptProgram(_envConfig.getRuleSet(), skipRuleSet, useRuleSet, _defaultDisabledPlannerRules); + optProgram = getOptProgram(_envConfig.getRuleSet(), skipRuleSet, useRuleSet, _defaultDisabledPlannerRules, + enableTypedMode); } } int sortExchangeCopyLimit = QueryOptionsUtils.getSortExchangeCopyThreshold(options, @@ -546,9 +551,15 @@ private DispatchableSubPlan toDispatchableSubPlan(RelRoot relRoot, PlannerContex /// @param skipRuleSet parsed skipped rule name set from query options /// @param useRuleSet parsed use rule set from query options /// @param defaultDisabledRuleSet parsed default disabled rule set from broker config + /// @param enableTypedMode whether the query explicitly opted into typed MODE implementations /// @return HepProgram that performs logical transformations private static HepProgram getOptProgram(PinotRuleSet ruleSet, Set skipRuleSet, Set useRuleSet, - Set defaultDisabledRuleSet) { + Set defaultDisabledRuleSet, boolean enableTypedMode) { + if (!enableTypedMode) { + // Rollout safety must not depend on customized default-disabled rules or usePlannerRules overrides. + skipRuleSet = new HashSet<>(skipRuleSet); + skipRuleSet.add(CommonConstants.Broker.PlannerRuleNames.TYPED_MODE_REWRITE); + } HepProgramBuilder hepProgramBuilder = new HepProgramBuilder(); // Set the match order as DEPTH_FIRST. The default is arbitrary which works the same as DEPTH_FIRST, but it's // best to be explicit. diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java index b2d4a5a544c3..8a831dba927a 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java @@ -20,15 +20,20 @@ import java.util.ArrayList; import java.util.List; +import java.util.Set; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.routing.MockRoutingManagerFactory; +import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.QueryEnvironmentTestBase; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.routing.WorkerManager; +import org.apache.pinot.spi.utils.CommonConstants; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -47,7 +52,7 @@ public Object[][] physicalOptimizers() { @Test public void testModeReturnTypes() { RelDataType rowType = _queryEnvironment.compile( - "SET autoRewriteAggregationType=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3), MODE(col7), " + "SET enableTypedMode=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3), MODE(col7), " + "MODE(CAST(col3 AS FLOAT)), MODE(CAST(col3 AS DOUBLE)), " + "MODE(NULLIF(JSONEXTRACTSCALAR(col1, '$.user', 'STRING', ''), '')), " + "fromTimestamp(MODE(ts_timestamp)) FROM a") @@ -62,7 +67,7 @@ public void testModeReturnTypes() { @Test public void testFilteredModeNullability() { RelDataType rowType = _queryEnvironment.compile( - "SET autoRewriteAggregationType=true; " + "SET enableTypedMode=true; " + "SELECT col2, MODE(col1) FILTER (WHERE col3 > 0), MODE(ts_timestamp) FILTER (WHERE col3 > 0) " + "FROM a GROUP BY col2").getRelRoot().validatedRowType; assertTrue(rowType.getFieldList().get(1).getType().isNullable()); @@ -72,7 +77,7 @@ public void testFilteredModeNullability() { @Test(dataProvider = "physicalOptimizers") public void testDistributedModeTypes(boolean usePhysicalOptimizer) { DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + "SET autoRewriteAggregationType=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3) FROM a"); + + "SET enableTypedMode=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3) FROM a"); PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); assertEquals(root.getDataSchema().getColumnDataTypes(), new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP, ColumnDataType.DOUBLE}); @@ -102,7 +107,7 @@ public void testDistributedModeTypes(boolean usePhysicalOptimizer) { @Test(dataProvider = "physicalOptimizers") public void testModeExpressionsAndTieBreakers(boolean usePhysicalOptimizer) { DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + "SET autoRewriteAggregationType=true; " + + "SET enableTypedMode=true; " + "SELECT col2, MODE(NULLIF(col1, ''), 'MIN'), MODE(ts_timestamp, 'MAX'), MODE(col3, 'AVG') " + "FROM a GROUP BY col2"); PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); @@ -115,6 +120,79 @@ public void testModeExpressionsAndTieBreakers(boolean usePhysicalOptimizer) { } } + @Test(dataProvider = "physicalOptimizers") + public void testModeRewriteRequiresItsOwnOptIn(boolean usePhysicalOptimizer) { + for (String options : List.of("", "SET autoRewriteAggregationType=true; ", + "SET autoRewriteAggregationType=true; SET enableTypedMode=false; ")) { + DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + + options + "SELECT MODE(ts_timestamp), MODE(col3) FROM a"); + List aggregates = findAggregates(plan); + assertFalse(aggregates.isEmpty()); + for (AggregateNode aggregate : aggregates) { + assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), + List.of("MODE", "MODE"), options); + } + } + } + + @Test(dataProvider = "physicalOptimizers") + public void testTypedModeDoesNotEnableOtherAggregateRewrites(boolean usePhysicalOptimizer) { + DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + + "SET autoRewriteAggregationType=false; SET enableTypedMode=true; " + + "SELECT MODE(col1), MODE(ts_timestamp), MODE(col3), MIN(col7), MAX(col7), SUM(col7) FROM a"); + List aggregates = findAggregates(plan); + assertFalse(aggregates.isEmpty()); + for (AggregateNode aggregate : aggregates) { + List functionNames = + aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(); + assertEquals(functionNames.subList(0, 5), List.of("MODESTRING", "MODETIMESTAMP", "MODE", "MIN", "MAX")); + assertFalse(functionNames.contains("SUMLONG")); + assertFalse(functionNames.contains("SUMINT")); + } + } + + @Test(dataProvider = "physicalOptimizers") + public void testTypedModeOptInSurvivesCustomizedPlannerDefaults(boolean usePhysicalOptimizer) { + for (Set disabledRules : List.of(Set.of(), + Set.of(CommonConstants.Broker.PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE))) { + QueryEnvironment environment = buildQueryEnvironment(disabledRules); + for (String options : List.of("", "SET enableTypedMode=false; ", + "SET autoRewriteAggregationType=true; ", "SET usePlannerRules='TypedModeRewrite'; ", + "SET usePlannerRules='TypedModeRewrite'; SET enableTypedMode=false; ", + "SET enableTypedMode=true; SET skipPlannerRules='TypedModeRewrite'; ")) { + assertModeCalls(environment, usePhysicalOptimizer, options, List.of("MODE", "MODE", "MODE")); + } + assertModeCalls(environment, usePhysicalOptimizer, "SET enableTypedMode=true; ", + List.of("MODESTRING", "MODETIMESTAMP", "MODE")); + } + } + + private static QueryEnvironment buildQueryEnvironment(Set disabledRules) { + MockRoutingManagerFactory factory = new MockRoutingManagerFactory(1, 2); + TABLE_SCHEMAS.forEach((name, schema) -> factory.registerTable(schema, name)); + SERVER1_SEGMENTS.forEach((table, segments) -> segments.forEach(s -> factory.registerSegment(1, table, s))); + SERVER2_SEGMENTS.forEach((table, segments) -> segments.forEach(s -> factory.registerSegment(2, table, s))); + return new QueryEnvironment(QueryEnvironment.configBuilder() + .requestId(-1L) + .database(CommonConstants.DEFAULT_DATABASE) + .tableCache(factory.buildTableCache()) + .workerManager(new WorkerManager("Broker_localhost", "localhost", 3, factory.buildRoutingManager(null))) + .defaultDisabledPlannerRules(disabledRules) + .build()); + } + + private static void assertModeCalls(QueryEnvironment environment, boolean usePhysicalOptimizer, String options, + List expectedCalls) { + DispatchableSubPlan plan = environment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + + options + "SELECT MODE(col1), MODE(ts_timestamp), MODE(col3) FROM a"); + List aggregates = findAggregates(plan); + assertFalse(aggregates.isEmpty()); + for (AggregateNode aggregate : aggregates) { + assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), + expectedCalls, options); + } + } + private static List findAggregates(DispatchableSubPlan plan) { List aggregates = new ArrayList<>(); for (DispatchablePlanFragment fragment : plan.getQueryStageMap().values()) { diff --git a/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json index 0976e9a830ea..79cbd49ed6e8 100644 --- a/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json +++ b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json @@ -49,12 +49,12 @@ "queries": [ { "description": "Global string and timestamp modes preserve the argument types", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'a'", + "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'a'", "outputs": [["p2", "2026-09-03 09:00:00.123"]] }, { "description": "Grouped modes ignore nulls, retain empty strings, and break ties by the smallest value", - "sql": "SET autoRewriteAggregationType=true; SELECT lpn_id, MODE(pallet_id), MODE(created_on) FROM {items} GROUP BY lpn_id", + "sql": "SET enableTypedMode=true; SELECT lpn_id, MODE(pallet_id), MODE(created_on) FROM {items} GROUP BY lpn_id", "outputs": [ ["a", "p2", "2026-09-03 09:00:00.123"], ["nulls", null, null], @@ -64,47 +64,47 @@ }, { "description": "MODE applies independently to a stored string and computed JSON string", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'a'", + "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'a'", "outputs": [["p2", "u1"]] }, { "description": "Typed MAX tie reducers select the largest original value", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id, 'MAX'), MODE(created_on, 'MAX') FROM {items} WHERE lpn_id = 'ties'", + "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id, 'MAX'), MODE(created_on, 'MAX') FROM {items} WHERE lpn_id = 'ties'", "outputs": [["z", "1970-01-01 00:00:00.001"]] }, { "description": "Existing numeric modes retain DOUBLE results and MIN, MAX, and AVG reducers", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(manifest_id), MODE(manifest_id, 'MAX'), MODE(manifest_id, 'AVG') FROM {items} WHERE lpn_id = 'a'", + "sql": "SET enableTypedMode=true; SELECT MODE(manifest_id), MODE(manifest_id, 'MAX'), MODE(manifest_id, 'AVG') FROM {items} WHERE lpn_id = 'a'", "outputs": [[1.0, 2.0, 1.5]] }, { "description": "Computed null inputs do not become a string mode", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'nulls'", + "sql": "SET enableTypedMode=true; SELECT MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'nulls'", "outputs": [[null]] }, { "description": "String expressions and post-aggregation transforms use the string result type", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(CONCAT(first_name, last_name, ' ')), UPPER(MODE(pallet_id)) FROM {items} WHERE lpn_id = 'a'", + "sql": "SET enableTypedMode=true; SELECT MODE(CONCAT(first_name, last_name, ' ')), UPPER(MODE(pallet_id)) FROM {items} WHERE lpn_id = 'a'", "outputs": [["Alice Jones", "P2"]] }, { "description": "Empty aggregate inputs return null for both types", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'missing'", + "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'missing'", "outputs": [[null, null]] }, { "description": "Filtered typed modes return null while preserving their group", - "sql": "SET autoRewriteAggregationType=true; SELECT lpn_id, MODE(pallet_id) FILTER (WHERE manifest_id < 0), MODE(created_on) FILTER (WHERE manifest_id < 0) FROM {items} WHERE lpn_id = 'a' GROUP BY lpn_id", + "sql": "SET enableTypedMode=true; SELECT lpn_id, MODE(pallet_id) FILTER (WHERE manifest_id < 0), MODE(created_on) FILTER (WHERE manifest_id < 0) FROM {items} WHERE lpn_id = 'a' GROUP BY lpn_id", "outputs": [["a", null, null]] }, { "description": "Timestamp mode can feed a parent comparison; fixture ingestion uses Los Angeles time and the SQL literal uses UTC", - "sql": "SET autoRewriteAggregationType=true; SELECT lpn_id FROM (SELECT lpn_id, MODE(created_on) AS mode_time FROM {items} GROUP BY lpn_id) WHERE mode_time = TIMESTAMP '2026-09-03 16:00:00.123'", + "sql": "SET enableTypedMode=true; SELECT lpn_id FROM (SELECT lpn_id, MODE(created_on) AS mode_time FROM {items} GROUP BY lpn_id) WHERE mode_time = TIMESTAMP '2026-09-03 16:00:00.123'", "outputs": [["a"]] }, { "description": "Joined grouping supports independent string, JSON, and timestamp modes, including unmatched outer rows", - "sql": "SET autoRewriteAggregationType=true; SELECT i.lpn_id, MODE(i.pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(i.meta, '$.action_user', 'STRING', ''), '')), MODE(m.code), MODE(m.status), MODE(m.created_on) FROM {items} i LEFT JOIN {manifests} m ON m.id = i.manifest_id GROUP BY i.lpn_id", + "sql": "SET enableTypedMode=true; SELECT i.lpn_id, MODE(i.pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(i.meta, '$.action_user', 'STRING', ''), '')), MODE(m.code), MODE(m.status), MODE(m.created_on) FROM {items} i LEFT JOIN {manifests} m ON m.id = i.manifest_id GROUP BY i.lpn_id", "outputs": [ ["a", "p2", "u1", "A", "pending", "2026-09-03 11:00:00.123"], ["nulls", null, null, null, null, null], @@ -137,12 +137,12 @@ "queries": [ { "description": "Direct aggregate emits scalar string and timestamp values instead of frequency-map intermediates", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items}", + "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items}", "outputs": [["p2", "2026-09-03 09:00:00.123"]] }, { "description": "Direct aggregate with no matching values emits typed nulls", - "sql": "SET autoRewriteAggregationType=true; SELECT MODE(pallet_id) FILTER (WHERE pallet_id = 'missing'), MODE(created_on) FILTER (WHERE pallet_id = 'missing') FROM {items}", + "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id) FILTER (WHERE pallet_id = 'missing'), MODE(created_on) FILTER (WHERE pallet_id = 'missing') FROM {items}", "outputs": [[null, null]] } ] diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 50cb94b4a152..a650ef4df5f3 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1044,6 +1044,10 @@ public static class QueryOptionKey { // MAX(stringCol) -> MAXSTRING(stringCol) // SUM(intCol) -> SUMINT(intCol) public static final String AUTO_REWRITE_AGGREGATION_TYPE = "autoRewriteAggregationType"; + + /// Opts into string and TIMESTAMP MODE implementations after all query components have been upgraded. + /// Kept separate from existing aggregate rewrites to preserve timestamp MODE semantics during rolling upgrades. + public static final String ENABLE_TYPED_MODE = "enableTypedMode"; // When enabled, allows multi cluster/federated queries to be executed. public static final String ENABLE_MULTI_CLUSTER_ROUTING = "enableMultiClusterRouting"; @@ -1116,6 +1120,7 @@ public static class PlannerRuleNames { public static final String AGGREGATE_UNION_TRANSPOSE = "AggregateUnionTranspose"; public static final String AGGREGATE_REDUCE_FUNCTIONS = "AggregateReduceFunctions"; public static final String AGGREGATE_FUNCTION_REWRITE = "AggregateFunctionRewrite"; + public static final String TYPED_MODE_REWRITE = "TypedModeRewrite"; public static final String AGGREGATE_CASE_TO_FILTER = "AggregateCaseToFilter"; public static final String PROJECT_FILTER_TRANSPOSE = "ProjectFilterTranspose"; public static final String PROJECT_MERGE = "ProjectMerge"; @@ -1162,6 +1167,7 @@ public static class PlannerRuleNames { PlannerRuleNames.AGGREGATE_UNION_AGGREGATE, PlannerRuleNames.JOIN_TO_ENRICHED_JOIN, PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE, + PlannerRuleNames.TYPED_MODE_REWRITE, // Stock Calcite rule kept opt-in via usePlannerRules — see SORT_PROJECT_TRANSPOSE javadoc // above for the rationale (firing in BASIC_RULES disrupts ProjectToSemiJoinRule on // partition-hinted IN(SELECT) queries, breaking colocated broadcast semi-joins). From 5e9a75122596ce16829f7e7b786b7d4147f75049 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 7 Sep 2026 14:14:05 -0700 Subject: [PATCH 4/8] Consolidate typed MODE using inferred type arguments --- .../function/AggregationFunctionFactory.java | 4 - ...BaseComparableModeAggregationFunction.java | 232 ----------- .../function/ModeAggregationFunction.java | 381 +++++++++++++++++- .../ModeStringAggregationFunction.java | 181 --------- .../ModeTimestampAggregationFunction.java | 120 ------ ...deAggregationFunctionRewriteOptimizer.java | 20 +- .../query/reduce/BaseGapfillProcessor.java | 11 + .../core/query/reduce/GapfillProcessor.java | 5 +- .../AggregationFunctionFactoryTest.java | 11 + .../function/ModeAggregationFunctionTest.java | 39 +- ...ModeNonNumericAggregationFunctionTest.java | 131 ++++-- ...gregationFunctionRewriteOptimizerTest.java | 39 +- .../apache/pinot/queries/BaseQueriesTest.java | 20 +- .../apache/pinot/queries/ExprMinMaxTest.java | 40 ++ .../apache/pinot/queries/ModeQueriesTest.java | 43 ++ ...notModeAggregationFunctionRewriteRule.java | 104 +++-- .../query/queries/ModeSqlPlannerTest.java | 80 +++- .../segment/spi/AggregationFunctionType.java | 22 +- 18 files changed, 809 insertions(+), 674 deletions(-) delete mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java delete mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java delete mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java index 11b05e4f32a8..fdc93df43e80 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java @@ -245,10 +245,6 @@ public static AggregationFunction getAggregationFunction(FunctionContext functio return new AvgAggregationFunction(arguments, nullHandlingEnabled); case MODE: return new ModeAggregationFunction(arguments, nullHandlingEnabled); - case MODESTRING: - return new ModeStringAggregationFunction(arguments, nullHandlingEnabled); - case MODETIMESTAMP: - return new ModeTimestampAggregationFunction(arguments, nullHandlingEnabled); case ANYVALUE: return new AnyValueAggregationFunction(arguments, nullHandlingEnabled); case FIRSTWITHTIME: { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java deleted file mode 100644 index 48f6c4bce9f2..000000000000 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseComparableModeAggregationFunction.java +++ /dev/null @@ -1,232 +0,0 @@ -/** - * 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.core.query.aggregation.function; - -import it.unimi.dsi.fastutil.ints.Int2LongOpenHashMap; -import java.util.List; -import java.util.Map; -import javax.annotation.Nullable; -import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.utils.DataSchema.ColumnDataType; -import org.apache.pinot.core.common.BlockValSet; -import org.apache.pinot.core.query.aggregation.AggregationResultHolder; -import org.apache.pinot.core.query.aggregation.ObjectAggregationResultHolder; -import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; -import org.apache.pinot.core.query.aggregation.groupby.ObjectGroupByResultHolder; -import org.apache.pinot.segment.spi.index.reader.Dictionary; - -import static com.google.common.base.Preconditions.checkArgument; - - -/// Counts comparable values and resolves equally frequent values using MIN (default) or MAX. -/// -/// Instances are immutable and may be shared across segments. Accumulators belong to result holders, and dictionary -/// identifiers are converted to values before merging across segments. -/// Numeric MODE retains its existing implementation. -abstract class BaseComparableModeAggregationFunction> - extends BaseSingleInputAggregationFunction, T> { - private final boolean _minimum; - - protected BaseComparableModeAggregationFunction(List arguments, boolean nullHandlingEnabled, - String valueType) { - super(checkArguments(arguments), nullHandlingEnabled); - String reducer = "MIN"; - if (arguments.size() == 2) { - ExpressionContext argument = arguments.get(1); - checkArgument(argument.getType() == ExpressionContext.Type.LITERAL, - "MODE tie reducer must be a literal MIN or MAX for %s", valueType); - reducer = argument.getLiteral().getStringValue(); - } - checkArgument("MIN".equals(reducer) || "MAX".equals(reducer), - "MODE for %s supports only MIN or MAX tie reducers, got: %s", valueType, reducer); - _minimum = "MIN".equals(reducer); - } - - private static ExpressionContext checkArguments(List arguments) { - checkArgument(arguments.size() == 1 || arguments.size() == 2, - "MODE expects one or two arguments, got: %s", arguments.size()); - return arguments.get(0); - } - - protected abstract Map newValueMap(); - - protected abstract ValueCounter valueCounter(BlockValSet blockValSet); - - @FunctionalInterface - protected interface ValueCounter { - void add(Map counts, int row); - } - - protected abstract void putDictionaryCount(Map counts, Dictionary dictionary, int dictionaryId, long count); - - protected final boolean isMinimum() { - return _minimum; - } - - @Override - public AggregationResultHolder createAggregationResultHolder() { - return new ObjectAggregationResultHolder(); - } - - @Override - public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int maxCapacity) { - return new ObjectGroupByResultHolder(initialCapacity, maxCapacity); - } - - @Override - public void aggregate(int length, AggregationResultHolder holder, - Map blockValSetMap) { - BlockValSet values = blockValSetMap.get(_expression); - Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; - if (dictionary != null) { - int[] ids = values.getDictionaryIdsSV(); - forEachNotNull(length, values, (from, to) -> { - DictionaryCounts counts = getValue(holder, () -> new DictionaryCounts(dictionary)); - for (int i = from; i < to; i++) { - counts._counts.addTo(ids[i], 1L); - } - }); - } else { - ValueCounter counter = valueCounter(values); - forEachNotNull(length, values, (from, to) -> { - Map counts = getValue(holder, this::newValueMap); - for (int i = from; i < to; i++) { - counter.add(counts, i); - } - }); - } - } - - @Override - public void aggregateGroupBySV(int length, int[] groupKeys, GroupByResultHolder holder, - Map blockValSetMap) { - BlockValSet values = blockValSetMap.get(_expression); - Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; - if (dictionary != null) { - int[] ids = values.getDictionaryIdsSV(); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - DictionaryCounts counts = getValue(holder, groupKeys[i], () -> new DictionaryCounts(dictionary)); - counts._counts.addTo(ids[i], 1L); - } - }); - } else { - ValueCounter counter = valueCounter(values); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - Map counts = getValue(holder, groupKeys[i], this::newValueMap); - counter.add(counts, i); - } - }); - } - } - - @Override - public void aggregateGroupByMV(int length, int[][] groupKeys, GroupByResultHolder holder, - Map blockValSetMap) { - BlockValSet values = blockValSetMap.get(_expression); - Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; - if (dictionary != null) { - int[] ids = values.getDictionaryIdsSV(); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - for (int groupKey : groupKeys[i]) { - DictionaryCounts counts = getValue(holder, groupKey, () -> new DictionaryCounts(dictionary)); - counts._counts.addTo(ids[i], 1L); - } - } - }); - } else { - ValueCounter counter = valueCounter(values); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - for (int groupKey : groupKeys[i]) { - Map counts = getValue(holder, groupKey, this::newValueMap); - counter.add(counts, i); - } - } - }); - } - } - - @Nullable - @Override - public Map extractAggregationResult(AggregationResultHolder holder) { - return extractCounts(holder.getResult()); - } - - @Nullable - @Override - public Map extractGroupByResult(GroupByResultHolder holder, int groupKey) { - return extractCounts(holder.getResult(groupKey)); - } - - @Nullable - @SuppressWarnings("unchecked") - private Map extractCounts(@Nullable Object result) { - if (result instanceof DictionaryCounts) { - DictionaryCounts dictionaryCounts = (DictionaryCounts) result; - Map counts = newValueMap(); - dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> putDictionaryCount( - counts, dictionaryCounts._dictionary, entry.getIntKey(), entry.getLongValue())); - return counts; - } - return (Map) result; - } - - @Override - public Map merge(Map left, Map right) { - right.forEach((value, count) -> left.merge(value, count, Long::sum)); - return left; - } - - @Override - public ColumnDataType getIntermediateResultColumnType() { - return ColumnDataType.OBJECT; - } - - @Nullable - @Override - public T extractFinalResult(@Nullable Map counts) { - if (counts == null || counts.isEmpty()) { - return null; - } - T mode = null; - long maxCount = 0; - for (Map.Entry entry : counts.entrySet()) { - T value = entry.getKey(); - long count = entry.getValue(); - if (mode == null || count > maxCount || (count == maxCount - && (_minimum ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { - mode = value; - maxCount = count; - } - } - return mode; - } - - private static final class DictionaryCounts { - private final Dictionary _dictionary; - private final Int2LongOpenHashMap _counts = new Int2LongOpenHashMap(); - - private DictionaryCounts(Dictionary dictionary) { - _dictionary = dictionary; - } - } -} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java index b17437b1a2ff..02440f4755e8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java @@ -29,9 +29,14 @@ import it.unimi.dsi.fastutil.ints.Int2LongMap; import it.unimi.dsi.fastutil.ints.Int2LongOpenHashMap; import it.unimi.dsi.fastutil.longs.Long2LongMap; +import it.unimi.dsi.fastutil.longs.Long2LongMaps; import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; +import it.unimi.dsi.fastutil.objects.Object2LongMap; +import it.unimi.dsi.fastutil.objects.Object2LongMaps; +import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap; import it.unimi.dsi.fastutil.objects.ObjectIterator; import java.util.List; +import java.util.Locale; import java.util.Map; import javax.annotation.Nullable; import org.apache.pinot.common.CustomObject; @@ -54,25 +59,59 @@ /// /// Following arguments are supported: /// -/// - Expression: expression that contains the column to be calculated mode on, can be any Numeric column +/// - Expression: expression that contains the column to be calculated mode on /// - MultiModeReducerType (optional): the reducer to use in case of multiple modes present in data +/// +/// Numeric calls retain a DOUBLE result and support MIN, MAX and AVG tie reducers. The planner supplies an internal +/// third STRING or TIMESTAMP literal for non-numeric inputs, which retain their type and support MIN and MAX. +/// The result type is immutable so a function can be shared across segments and reconstructed on reducing stages. @SuppressWarnings({"rawtypes", "unchecked"}) -public class ModeAggregationFunction extends BaseSingleInputAggregationFunction, Double> { +public class ModeAggregationFunction extends BaseSingleInputAggregationFunction, Comparable> { private static final double DEFAULT_FINAL_RESULT = Double.NEGATIVE_INFINITY; private final MultiModeReducerType _multiModeReducerType; + private final ColumnDataType _resultType; + private final String _resultColumnName; public ModeAggregationFunction(List arguments, boolean nullHandlingEnabled) { - super(arguments.get(0), nullHandlingEnabled); + super(checkArguments(arguments), nullHandlingEnabled); int numArguments = arguments.size(); - Preconditions.checkArgument(numArguments <= 2, "Mode expects at most 2 arguments, got: %s", numArguments); + if (numArguments == 3) { + ExpressionContext typeArgument = arguments.get(2); + Preconditions.checkArgument(typeArgument.getType() == ExpressionContext.Type.LITERAL + && typeArgument.getLiteral().getType() == DataType.STRING, + "MODE result type must be a STRING or TIMESTAMP string literal"); + String resultType = typeArgument.getLiteral().getStringValue(); + Preconditions.checkArgument(resultType != null, "MODE result type must be STRING or TIMESTAMP, got: null"); + resultType = resultType.toUpperCase(Locale.ROOT); + Preconditions.checkArgument("STRING".equals(resultType) || "TIMESTAMP".equals(resultType), + "MODE result type must be STRING or TIMESTAMP, got: %s", resultType); + _resultType = ColumnDataType.valueOf(resultType); + } else { + _resultType = ColumnDataType.DOUBLE; + } if (numArguments > 1) { + Preconditions.checkArgument(arguments.get(1).getType() == ExpressionContext.Type.LITERAL, + "MODE tie reducer must be a literal"); _multiModeReducerType = MultiModeReducerType.valueOf(arguments.get(1).getLiteral().getStringValue()); } else { _multiModeReducerType = MultiModeReducerType.MIN; } + Preconditions.checkArgument( + _resultType == ColumnDataType.DOUBLE || _multiModeReducerType != MultiModeReducerType.AVG, + "MODE for %s supports only MIN or MAX tie reducers, got: %s", _resultType, _multiModeReducerType); + // Gapfill resolves aliases by matching the complete selection expression, including literal spelling. + _resultColumnName = numArguments == 3 + ? "mode(" + _expression + "," + arguments.get(1) + "," + arguments.get(2) + ")" + : super.getResultColumnName(); + } + + private static ExpressionContext checkArguments(List arguments) { + Preconditions.checkArgument(!arguments.isEmpty() && arguments.size() <= 3, + "MODE expects one to three arguments, got: %s", arguments.size()); + return arguments.get(0); } /// Helper method to create a value map for the given value type. @@ -206,7 +245,11 @@ private static Map convertToValueMap(DictIdsWrapper dict } /// Helper method to extract segment level intermediate result from the inner segment result. - private static Map extractIntermediateResult(@Nullable Object result) { + @Nullable + private Map extractIntermediateResult(@Nullable Object result) { + if (_resultType != ColumnDataType.DOUBLE) { + return extractComparableCounts(result); + } if (result == null) { // NOTE: Return an empty Int2LongOpenHashMap for empty result. return new Int2LongOpenHashMap(); @@ -226,6 +269,11 @@ public AggregationFunctionType getType() { return AggregationFunctionType.MODE; } + @Override + public String getResultColumnName() { + return _resultColumnName; + } + @Override public AggregationResultHolder createAggregationResultHolder() { return new ObjectAggregationResultHolder(); @@ -239,6 +287,10 @@ public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int ma @Override public void aggregate(int length, AggregationResultHolder aggregationResultHolder, Map blockValSetMap) { + if (_resultType != ColumnDataType.DOUBLE) { + aggregateComparable(length, aggregationResultHolder, blockValSetMap); + return; + } BlockValSet blockValSet = blockValSetMap.get(_expression); // For dictionary-encoded expression, store dictionary ids into the dictId map @@ -303,6 +355,10 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde @Override public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, Map blockValSetMap) { + if (_resultType != ColumnDataType.DOUBLE) { + aggregateComparableGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSetMap); + return; + } BlockValSet blockValSet = blockValSetMap.get(_expression); // For dictionary-encoded expression, store dictionary ids into the dictId map @@ -361,6 +417,10 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol @Override public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, Map blockValSetMap) { + if (_resultType != ColumnDataType.DOUBLE) { + aggregateComparableGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSetMap); + return; + } BlockValSet blockValSet = blockValSetMap.get(_expression); // For dictionary-encoded expression, store dictionary ids into the dictId map @@ -425,19 +485,23 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } + @Nullable @Override - public Map extractAggregationResult(AggregationResultHolder aggregationResultHolder) { + public Map extractAggregationResult(AggregationResultHolder aggregationResultHolder) { return extractIntermediateResult(aggregationResultHolder.getResult()); } + @Nullable @Override - public Map extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) { + public Map extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) { return extractIntermediateResult(groupByResultHolder.getResult(groupKey)); } @Override - public Map merge(Map intermediateResult1, - Map intermediateResult2) { + public Map merge(Map intermediateResult1, Map intermediateResult2) { + if (_resultType != ColumnDataType.DOUBLE) { + return mergeComparableCounts(intermediateResult1, intermediateResult2); + } if (intermediateResult1.isEmpty()) { return intermediateResult2; } @@ -473,7 +537,16 @@ public ColumnDataType getIntermediateResultColumnType() { } @Override - public SerializedIntermediateResult serializeIntermediateResult(Map longMap) { + @SuppressWarnings("deprecation") + public SerializedIntermediateResult serializeIntermediateResult(Map longMap) { + if (_resultType == ColumnDataType.STRING) { + // Reuse the existing map wire encoding for generic aggregation bridges. + return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.Map.getValue(), + ObjectSerDeUtils.MAP_SER_DE.serialize((Map) longMap)); + } + if (_resultType == ColumnDataType.TIMESTAMP && !(longMap instanceof Long2LongMap)) { + longMap = new Long2LongOpenHashMap((Map) longMap); + } if (longMap instanceof Int2LongMap) { return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.Int2LongMap.getValue(), ObjectSerDeUtils.INT_2_LONG_MAP_SER_DE.serialize((Int2LongMap) longMap)); @@ -494,18 +567,22 @@ public SerializedIntermediateResult serializeIntermediateResult(Map deserializeIntermediateResult(CustomObject customObject) { - return ObjectSerDeUtils.deserialize(customObject); + public Map deserializeIntermediateResult(CustomObject customObject) { + Map counts = ObjectSerDeUtils.deserialize(customObject); + return _resultType == ColumnDataType.STRING ? new StringModeCounts((Map) counts) : counts; } @Override public ColumnDataType getFinalResultColumnType() { - return ColumnDataType.DOUBLE; + return _resultType; } @Nullable @Override - public Double extractFinalResult(@Nullable Map intermediateResult) { + public Comparable extractFinalResult(@Nullable Map intermediateResult) { + if (_resultType != ColumnDataType.DOUBLE) { + return extractComparableFinalResult(intermediateResult); + } // A null intermediate result means nothing was aggregated, and the mode of nothing is NULL. An empty map is a // different thing: it is what an untouched single-stage result holder produces, and it keeps its historical // sentinel below so that path is not silently changed. @@ -519,13 +596,13 @@ public Double extractFinalResult(@Nullable Map intermedi return DEFAULT_FINAL_RESULT; } } else if (intermediateResult instanceof Int2LongOpenHashMap) { - return extractFinalResult((Int2LongOpenHashMap) intermediateResult); + return extractNumericFinalResult((Int2LongOpenHashMap) intermediateResult); } else if (intermediateResult instanceof Long2LongOpenHashMap) { - return extractFinalResult((Long2LongOpenHashMap) intermediateResult); + return extractNumericFinalResult((Long2LongOpenHashMap) intermediateResult); } else if (intermediateResult instanceof Float2LongOpenHashMap) { - return extractFinalResult((Float2LongOpenHashMap) intermediateResult); + return extractNumericFinalResult((Float2LongOpenHashMap) intermediateResult); } else if (intermediateResult instanceof Double2LongOpenHashMap) { - return extractFinalResult((Double2LongOpenHashMap) intermediateResult); + return extractNumericFinalResult((Double2LongOpenHashMap) intermediateResult); } else { throw new IllegalStateException( "Illegal data type for Intermediate Result of MODE aggregation function: " + intermediateResult.getClass() @@ -533,7 +610,7 @@ public Double extractFinalResult(@Nullable Map intermedi } } - public double extractFinalResult(Int2LongOpenHashMap intermediateResult) { + private double extractNumericFinalResult(Int2LongOpenHashMap intermediateResult) { ObjectIterator iterator = intermediateResult.int2LongEntrySet().fastIterator(); Int2LongMap.Entry first = iterator.next(); long maxFrequency = first.getLongValue(); @@ -578,7 +655,7 @@ public double extractFinalResult(Int2LongOpenHashMap intermediateResult) { } } - public double extractFinalResult(Long2LongOpenHashMap intermediateResult) { + private double extractNumericFinalResult(Long2LongOpenHashMap intermediateResult) { ObjectIterator iterator = intermediateResult.long2LongEntrySet().fastIterator(); Long2LongMap.Entry first = iterator.next(); long maxFrequency = first.getLongValue(); @@ -625,7 +702,7 @@ public double extractFinalResult(Long2LongOpenHashMap intermediateResult) { } } - public double extractFinalResult(Float2LongOpenHashMap intermediateResult) { + private double extractNumericFinalResult(Float2LongOpenHashMap intermediateResult) { ObjectIterator iterator = intermediateResult.float2LongEntrySet().fastIterator(); Float2LongMap.Entry first = iterator.next(); long maxFrequency = first.getLongValue(); @@ -672,7 +749,7 @@ public double extractFinalResult(Float2LongOpenHashMap intermediateResult) { } } - public Double extractFinalResult(Double2LongOpenHashMap intermediateResult) { + private Double extractNumericFinalResult(Double2LongOpenHashMap intermediateResult) { ObjectIterator iterator = intermediateResult.double2LongEntrySet().fastIterator(); Double2LongMap.Entry first = iterator.next(); long maxFrequency = first.getLongValue(); @@ -719,6 +796,266 @@ public Double extractFinalResult(Double2LongOpenHashMap intermediateResult) { } } + private Map newComparableValueMap() { + return _resultType == ColumnDataType.STRING ? new StringModeCounts() : new Long2LongOpenHashMap(); + } + + private ValueCounter comparableValueCounter(BlockValSet blockValSet) { + if (_resultType == ColumnDataType.STRING) { + String[] values = blockValSet.getStringValuesSV(); + return (counts, row) -> ((StringModeCounts) counts).addTo(values[row], 1L); + } + long[] values = blockValSet.getLongValuesSV(); + return (counts, row) -> ((Long2LongOpenHashMap) counts).addTo(values[row], 1L); + } + + @FunctionalInterface + private interface ValueCounter { + void add(Map counts, int row); + } + + private void aggregateComparable(int length, AggregationResultHolder holder, + Map blockValSetMap) { + BlockValSet values = blockValSetMap.get(_expression); + Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; + if (dictionary != null) { + int[] ids = values.getDictionaryIdsSV(); + forEachNotNull(length, values, (from, to) -> { + DictionaryCounts counts = getValue(holder, () -> new DictionaryCounts(dictionary)); + for (int i = from; i < to; i++) { + counts._counts.addTo(ids[i], 1L); + } + }); + } else { + ValueCounter counter = comparableValueCounter(values); + forEachNotNull(length, values, (from, to) -> { + Map counts = getValue(holder, this::newComparableValueMap); + for (int i = from; i < to; i++) { + counter.add(counts, i); + } + }); + } + } + + private void aggregateComparableGroupBySV(int length, int[] groupKeys, GroupByResultHolder holder, + Map blockValSetMap) { + BlockValSet values = blockValSetMap.get(_expression); + Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; + if (dictionary != null) { + int[] ids = values.getDictionaryIdsSV(); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + DictionaryCounts counts = getValue(holder, groupKeys[i], () -> new DictionaryCounts(dictionary)); + counts._counts.addTo(ids[i], 1L); + } + }); + } else { + ValueCounter counter = comparableValueCounter(values); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + Map counts = getValue(holder, groupKeys[i], this::newComparableValueMap); + counter.add(counts, i); + } + }); + } + } + + private void aggregateComparableGroupByMV(int length, int[][] groupKeys, GroupByResultHolder holder, + Map blockValSetMap) { + BlockValSet values = blockValSetMap.get(_expression); + Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; + if (dictionary != null) { + int[] ids = values.getDictionaryIdsSV(); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + for (int groupKey : groupKeys[i]) { + DictionaryCounts counts = getValue(holder, groupKey, () -> new DictionaryCounts(dictionary)); + counts._counts.addTo(ids[i], 1L); + } + } + }); + } else { + ValueCounter counter = comparableValueCounter(values); + forEachNotNull(length, values, (from, to) -> { + for (int i = from; i < to; i++) { + for (int groupKey : groupKeys[i]) { + Map counts = getValue(holder, groupKey, this::newComparableValueMap); + counter.add(counts, i); + } + } + }); + } + } + + @Nullable + private Map extractComparableCounts(@Nullable Object result) { + if (!(result instanceof DictionaryCounts)) { + return (Map) result; + } + DictionaryCounts dictionaryCounts = (DictionaryCounts) result; + Dictionary dictionary = dictionaryCounts._dictionary; + if (_resultType == ColumnDataType.STRING) { + StringModeCounts counts = new StringModeCounts(); + dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> + counts.put(dictionary.getStringValue(entry.getIntKey()), entry.getLongValue())); + return counts; + } + Long2LongOpenHashMap counts = new Long2LongOpenHashMap(); + dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> + counts.put(dictionary.getLongValue(entry.getIntKey()), entry.getLongValue())); + return counts; + } + + private Map mergeComparableCounts(Map left, Map right) { + if (_resultType == ColumnDataType.STRING && left instanceof Object2LongOpenHashMap + && right instanceof Object2LongMap) { + Object2LongOpenHashMap counts = (Object2LongOpenHashMap) left; + ObjectIterator> iterator = + Object2LongMaps.fastIterator((Object2LongMap) right); + while (iterator.hasNext()) { + Object2LongMap.Entry entry = iterator.next(); + counts.addTo(entry.getKey(), entry.getLongValue()); + } + return left; + } + if (_resultType == ColumnDataType.TIMESTAMP && left instanceof Long2LongOpenHashMap + && right instanceof Long2LongMap) { + Long2LongOpenHashMap counts = (Long2LongOpenHashMap) left; + ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) right); + while (iterator.hasNext()) { + Long2LongMap.Entry entry = iterator.next(); + counts.addTo(entry.getLongKey(), entry.getLongValue()); + } + return left; + } + Map counts = (Map) left; + right.forEach((value, count) -> counts.merge(value, count, Long::sum)); + return left; + } + + @Nullable + private Comparable extractComparableFinalResult(@Nullable Map counts) { + if (counts == null || counts.isEmpty()) { + return null; + } + boolean minimum = _multiModeReducerType == MultiModeReducerType.MIN; + if (_resultType == ColumnDataType.STRING && counts instanceof Object2LongMap) { + String mode = null; + long maxCount = 0; + ObjectIterator> iterator = + Object2LongMaps.fastIterator((Object2LongMap) counts); + while (iterator.hasNext()) { + Object2LongMap.Entry entry = iterator.next(); + String value = entry.getKey(); + long count = entry.getLongValue(); + if (mode == null || count > maxCount || (count == maxCount + && (minimum ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { + mode = value; + maxCount = count; + } + } + return mode; + } + if (_resultType == ColumnDataType.TIMESTAMP && counts instanceof Long2LongMap) { + ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) counts); + Long2LongMap.Entry first = iterator.next(); + long mode = first.getLongKey(); + long maxCount = first.getLongValue(); + while (iterator.hasNext()) { + Long2LongMap.Entry entry = iterator.next(); + long value = entry.getLongKey(); + long count = entry.getLongValue(); + if (count > maxCount || (count == maxCount && (minimum ? value < mode : value > mode))) { + mode = value; + maxCount = count; + } + } + return mode; + } + Comparable mode = null; + long maxCount = 0; + for (Map.Entry entry : counts.entrySet()) { + Comparable value = (Comparable) entry.getKey(); + long count = entry.getValue(); + if (mode == null || count > maxCount || (count == maxCount + && (minimum ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { + mode = value; + maxCount = count; + } + } + return mode; + } + + private static final class DictionaryCounts { + private final Dictionary _dictionary; + private final Int2LongOpenHashMap _counts = new Int2LongOpenHashMap(); + + private DictionaryCounts(Dictionary dictionary) { + _dictionary = dictionary; + } + } + + /// Frequency state with an O(1) conservative estimate of the retained string-key payload. + /// Accumulation uses [#addTo] and dictionary extraction and boxed [Map#merge] use [#put]. + /// Each distinct key is charged once, assuming UTF-16 storage plus object and array overhead. + /// Instances belong to one result holder and are not thread-safe. + public static final class StringModeCounts extends Object2LongOpenHashMap { + private long _retainedStringBytes; + + public StringModeCounts() { + } + + /// Restores accounting once when a generic map is deserialized from the existing wire format. + public StringModeCounts(Map counts) { + super(counts.size()); + counts.forEach((value, count) -> put(value, count.longValue())); + } + + public long getRetainedStringBytes() { + return _retainedStringBytes; + } + + @Override + public long addTo(String value, long increment) { + int previousSize = size(); + long previousCount = super.addTo(value, increment); + if (size() != previousSize) { + _retainedStringBytes += retainedStringBytes(value); + } + return previousCount; + } + + @Override + public long put(String value, long count) { + int previousSize = size(); + long previousCount = super.put(value, count); + if (size() != previousSize) { + _retainedStringBytes += retainedStringBytes(value); + } + return previousCount; + } + + @Override + public long removeLong(Object value) { + int previousSize = size(); + long previousCount = super.removeLong(value); + if (size() != previousSize) { + _retainedStringBytes -= retainedStringBytes((String) value); + } + return previousCount; + } + + @Override + public void clear() { + super.clear(); + _retainedStringBytes = 0; + } + + private static long retainedStringBytes(String value) { + return 48 + 2L * value.length(); + } + } + private enum MultiModeReducerType { MIN, MAX, AVG } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java deleted file mode 100644 index 93b1a145114c..000000000000 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeStringAggregationFunction.java +++ /dev/null @@ -1,181 +0,0 @@ -/** - * 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.core.query.aggregation.function; - -import it.unimi.dsi.fastutil.objects.Object2LongMap; -import it.unimi.dsi.fastutil.objects.Object2LongMaps; -import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap; -import it.unimi.dsi.fastutil.objects.ObjectIterator; -import java.util.List; -import java.util.Map; -import javax.annotation.Nullable; -import org.apache.pinot.common.CustomObject; -import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.utils.DataSchema.ColumnDataType; -import org.apache.pinot.core.common.BlockValSet; -import org.apache.pinot.core.common.ObjectSerDeUtils; -import org.apache.pinot.segment.spi.AggregationFunctionType; -import org.apache.pinot.segment.spi.index.reader.Dictionary; - - -/// String implementation of MODE, with lexicographic MIN/MAX tie resolution and a fixed STRING result type. -/// Instances are immutable; per-segment frequency maps are stored in the result holders. -public class ModeStringAggregationFunction extends BaseComparableModeAggregationFunction { - public ModeStringAggregationFunction(List arguments, boolean nullHandlingEnabled) { - super(arguments, nullHandlingEnabled, "STRING"); - } - - @Override - public AggregationFunctionType getType() { - return AggregationFunctionType.MODESTRING; - } - - @Override - public ColumnDataType getFinalResultColumnType() { - return ColumnDataType.STRING; - } - - @Override - protected Map newValueMap() { - return new StringModeCounts(); - } - - @Override - protected ValueCounter valueCounter(BlockValSet blockValSet) { - String[] values = blockValSet.getStringValuesSV(); - return (counts, row) -> ((StringModeCounts) counts).addTo(values[row], 1L); - } - - @Override - protected void putDictionaryCount(Map counts, Dictionary dictionary, int dictionaryId, long count) { - ((StringModeCounts) counts).put(dictionary.getStringValue(dictionaryId), count); - } - - @Override - public Map merge(Map left, Map right) { - if (left instanceof Object2LongOpenHashMap && right instanceof Object2LongMap) { - Object2LongOpenHashMap counts = (Object2LongOpenHashMap) left; - ObjectIterator> iterator = - Object2LongMaps.fastIterator((Object2LongMap) right); - while (iterator.hasNext()) { - Object2LongMap.Entry entry = iterator.next(); - counts.addTo(entry.getKey(), entry.getLongValue()); - } - return left; - } - return super.merge(left, right); - } - - @Nullable - @Override - public String extractFinalResult(@Nullable Map counts) { - if (!(counts instanceof Object2LongMap)) { - return super.extractFinalResult(counts); - } - String mode = null; - long maxCount = 0; - ObjectIterator> iterator = - Object2LongMaps.fastIterator((Object2LongMap) counts); - while (iterator.hasNext()) { - Object2LongMap.Entry entry = iterator.next(); - String value = entry.getKey(); - long count = entry.getLongValue(); - if (mode == null || count > maxCount || (count == maxCount - && (isMinimum() ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { - mode = value; - maxCount = count; - } - } - return mode; - } - - @Override - @SuppressWarnings({"unchecked", "rawtypes", "deprecation"}) - public SerializedIntermediateResult serializeIntermediateResult(Map counts) { - // Reuse the existing map wire encoding so generic aggregation bridges can deserialize the frequency state. - return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.Map.getValue(), - ObjectSerDeUtils.MAP_SER_DE.serialize((Map) counts)); - } - - @Override - public Map deserializeIntermediateResult(CustomObject customObject) { - return new StringModeCounts(ObjectSerDeUtils.deserialize(customObject)); - } - - /// Frequency state with an O(1) conservative estimate of the retained string-key payload. - /// Accumulation uses [#addTo] and dictionary extraction and boxed [Map#merge] use [#put]. - /// Each distinct key is charged once, assuming UTF-16 storage plus object and array overhead. - /// Instances belong to one result holder and are not thread-safe. - public static final class StringModeCounts extends Object2LongOpenHashMap { - private long _retainedStringBytes; - - public StringModeCounts() { - } - - /// Restores accounting once when a generic map is deserialized from the existing wire format. - public StringModeCounts(Map counts) { - super(counts.size()); - counts.forEach((value, count) -> put(value, count.longValue())); - } - - public long getRetainedStringBytes() { - return _retainedStringBytes; - } - - @Override - public long addTo(String value, long increment) { - int previousSize = size(); - long previousCount = super.addTo(value, increment); - if (size() != previousSize) { - _retainedStringBytes += retainedStringBytes(value); - } - return previousCount; - } - - @Override - public long put(String value, long count) { - int previousSize = size(); - long previousCount = super.put(value, count); - if (size() != previousSize) { - _retainedStringBytes += retainedStringBytes(value); - } - return previousCount; - } - - @Override - public long removeLong(Object value) { - int previousSize = size(); - long previousCount = super.removeLong(value); - if (size() != previousSize) { - _retainedStringBytes -= retainedStringBytes((String) value); - } - return previousCount; - } - - @Override - public void clear() { - super.clear(); - _retainedStringBytes = 0; - } - - private static long retainedStringBytes(String value) { - return 48 + 2L * value.length(); - } - } -} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java deleted file mode 100644 index 0cd9fe996e79..000000000000 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeTimestampAggregationFunction.java +++ /dev/null @@ -1,120 +0,0 @@ -/** - * 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.core.query.aggregation.function; - -import it.unimi.dsi.fastutil.longs.Long2LongMap; -import it.unimi.dsi.fastutil.longs.Long2LongMaps; -import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; -import it.unimi.dsi.fastutil.objects.ObjectIterator; -import java.util.List; -import java.util.Map; -import javax.annotation.Nullable; -import org.apache.pinot.common.CustomObject; -import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.utils.DataSchema.ColumnDataType; -import org.apache.pinot.core.common.BlockValSet; -import org.apache.pinot.core.common.ObjectSerDeUtils; -import org.apache.pinot.segment.spi.AggregationFunctionType; -import org.apache.pinot.segment.spi.index.reader.Dictionary; - - -/// Timestamp implementation of MODE, preserving epoch milliseconds without conversion through DOUBLE. -/// Instances are immutable; per-segment frequency maps are stored in the result holders. -public class ModeTimestampAggregationFunction extends BaseComparableModeAggregationFunction { - public ModeTimestampAggregationFunction(List arguments, boolean nullHandlingEnabled) { - super(arguments, nullHandlingEnabled, "TIMESTAMP"); - } - - @Override - public AggregationFunctionType getType() { - return AggregationFunctionType.MODETIMESTAMP; - } - - @Override - public ColumnDataType getFinalResultColumnType() { - return ColumnDataType.TIMESTAMP; - } - - @Override - protected Map newValueMap() { - return new Long2LongOpenHashMap(); - } - - @Override - protected ValueCounter valueCounter(BlockValSet blockValSet) { - long[] values = blockValSet.getLongValuesSV(); - return (counts, row) -> ((Long2LongOpenHashMap) counts).addTo(values[row], 1L); - } - - @Override - protected void putDictionaryCount(Map counts, Dictionary dictionary, int dictionaryId, long count) { - ((Long2LongOpenHashMap) counts).put(dictionary.getLongValue(dictionaryId), count); - } - - @Override - public Map merge(Map left, Map right) { - if (!(left instanceof Long2LongOpenHashMap) || !(right instanceof Long2LongMap)) { - return super.merge(left, right); - } - Long2LongOpenHashMap counts = (Long2LongOpenHashMap) left; - ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) right); - while (iterator.hasNext()) { - Long2LongMap.Entry entry = iterator.next(); - counts.addTo(entry.getLongKey(), entry.getLongValue()); - } - return counts; - } - - @Nullable - @Override - public Long extractFinalResult(@Nullable Map counts) { - if (!(counts instanceof Long2LongMap)) { - return super.extractFinalResult(counts); - } - ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) counts); - if (!iterator.hasNext()) { - return null; - } - Long2LongMap.Entry first = iterator.next(); - long mode = first.getLongKey(); - long maxCount = first.getLongValue(); - while (iterator.hasNext()) { - Long2LongMap.Entry entry = iterator.next(); - long value = entry.getLongKey(); - long count = entry.getLongValue(); - if (count > maxCount || (count == maxCount && (isMinimum() ? value < mode : value > mode))) { - mode = value; - maxCount = count; - } - } - return mode; - } - - @Override - public SerializedIntermediateResult serializeIntermediateResult(Map counts) { - Long2LongMap longCounts = counts instanceof Long2LongMap ? (Long2LongMap) counts : new Long2LongOpenHashMap(counts); - return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.Long2LongMap.getValue(), - ObjectSerDeUtils.LONG_2_LONG_MAP_SER_DE.serialize(longCounts)); - } - - @Override - public Map deserializeIntermediateResult(CustomObject customObject) { - return ObjectSerDeUtils.deserialize(customObject); - } -} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java index 9495536df477..051ac3202e30 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.core.query.optimizer.statement; +import java.util.ArrayList; import java.util.List; import java.util.Locale; import javax.annotation.Nullable; @@ -30,6 +31,7 @@ import org.apache.pinot.common.request.context.LiteralContext; import org.apache.pinot.common.request.context.RequestContextUtils; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.spi.data.DateTimeFieldSpec; import org.apache.pinot.spi.data.DateTimeFormatSpec; @@ -38,8 +40,8 @@ import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; -/// When `enableTypedMode` is enabled, resolves string and timestamp MODE expressions to implementations -/// with fixed result types before execution. +/// When `enableTypedMode` is enabled, adds an inferred type argument to string and timestamp MODE expressions +/// so the result type is fixed before execution. /// This also supplies the broker with the correct result type when no rows match or groups are trimmed before /// finalization. Type inference reads schema and function metadata only: server-dependent transforms such as LOOKUP /// must not be initialized on the broker. Expressions whose type is unknown retain the legacy MODE implementation. @@ -73,15 +75,19 @@ private static void rewriteExpression(@Nullable Expression expression, Schema sc Function function = expression.getFunctionCall(); List operands = function.getOperands(); rewriteExpressions(operands, schema); - if (!AggregationFunctionType.MODE.getName().equalsIgnoreCase(function.getOperator()) || operands.isEmpty()) { + if (!AggregationFunctionType.MODE.getName().equalsIgnoreCase(function.getOperator()) || operands.isEmpty() + || operands.size() >= 3) { return; } ColumnDataType operandType = getOperandType(operands.get(0), schema); - if (operandType == ColumnDataType.STRING) { - function.setOperator("modestring"); - } else if (operandType == ColumnDataType.TIMESTAMP) { - function.setOperator("modetimestamp"); + if (operandType == ColumnDataType.STRING || operandType == ColumnDataType.TIMESTAMP) { + List typedOperands = new ArrayList<>(operands); + if (typedOperands.size() == 1) { + typedOperands.add(RequestUtils.getLiteralExpression("MIN")); + } + typedOperands.add(RequestUtils.getLiteralExpression(operandType.name())); + function.setOperands(typedOperands); } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java index 6f815901a627..16adaf4bc3dc 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java @@ -118,6 +118,17 @@ protected void replaceColumnNameWithAlias(DataSchema dataSchema) { queryContext = _queryContext.getSubquery(); } List aliasList = queryContext.getAliasList(); + String[] columnNames = dataSchema.getColumnNames(); + if (columnNames.length == aliasList.size()) { + // Reduced results follow SELECT order. Server rewrites can change expression names (for example, inferred + // MODE type arguments), so bind aliases by position as in BaseReduceService.updateAlias. + for (int i = 0; i < columnNames.length; i++) { + if (aliasList.get(i) != null) { + columnNames[i] = aliasList.get(i); + } + } + return; + } Map columnNameToAliasMap = new HashMap<>(); for (int i = 0; i < aliasList.size(); i++) { if (aliasList.get(i) != null) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java index 0af38195b4df..a644e81a249d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java @@ -179,7 +179,10 @@ private void gapfill(long bucketTime, List bucketedResult, List '2026-09-03 00:00:00'") + .whenQuery("select mode(myField, 'MIN', 'TIMESTAMP') as mode " + + "from testTable where myField > '2026-09-03 00:00:00'") .thenResultIs(new Object[]{null}); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java index c6d7a6d27531..b4521996e045 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.core.query.aggregation.function; import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; +import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap; import java.nio.ByteBuffer; import java.util.HashMap; import java.util.List; @@ -30,6 +31,7 @@ import org.apache.pinot.core.common.BlockValSet; import org.apache.pinot.core.common.SyntheticBlockValSets; import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.function.ModeAggregationFunction.StringModeCounts; import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.roaringbitmap.RoaringBitmap; @@ -44,13 +46,13 @@ import static org.testng.Assert.expectThrows; -/// Verifies serialization, exact timestamp values and null handling for the typed MODE implementations. +/// Verifies serialization, exact timestamp values and null handling for the typed MODE calls. public class ModeNonNumericAggregationFunctionTest { private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("value"); @Test public void testStringIntermediateResultsRoundTripAndMerge() { - ModeStringAggregationFunction function = new ModeStringAggregationFunction(List.of(EXPRESSION), true); + ModeAggregationFunction function = new ModeAggregationFunction(typedArguments("MIN", "STRING"), true); var first = aggregateAndRoundTrip(function, SyntheticBlockValSets.Str.create(null, new String[]{"é", "é", "é", "苹果", "苹果", ""}), 6); var second = aggregateAndRoundTrip(function, @@ -67,10 +69,10 @@ public void testTimestampIntermediateResultsPreserveLongPrecision() { // Adjacent long values above 2^53 become identical if converted through double. long earlier = 9_007_199_254_740_992L; long later = earlier + 1; - ModeTimestampAggregationFunction minFunction = - new ModeTimestampAggregationFunction(List.of(EXPRESSION), true); - ModeTimestampAggregationFunction maxFunction = - new ModeTimestampAggregationFunction(argumentsWithReducer("MAX"), true); + ModeAggregationFunction minFunction = + new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), true); + ModeAggregationFunction maxFunction = + new ModeAggregationFunction(typedArguments("MAX", "TIMESTAMP"), true); var first = aggregateAndRoundTrip(minFunction, timestampValues(earlier, later, later), 3); var second = aggregateAndRoundTrip(minFunction, timestampValues(earlier), 1); @@ -83,10 +85,10 @@ public void testTimestampIntermediateResultsPreserveLongPrecision() { @Test public void testTimestampMergesPrimitiveAndGenericStates() { - ModeTimestampAggregationFunction minFunction = - new ModeTimestampAggregationFunction(List.of(EXPRESSION), true); - ModeTimestampAggregationFunction maxFunction = - new ModeTimestampAggregationFunction(argumentsWithReducer("MAX"), true); + ModeAggregationFunction minFunction = + new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), true); + ModeAggregationFunction maxFunction = + new ModeAggregationFunction(typedArguments("MAX", "TIMESTAMP"), true); for (boolean primitiveLeft : new boolean[]{false, true}) { for (boolean primitiveRight : new boolean[]{false, true}) { Map left = primitiveLeft ? new Long2LongOpenHashMap() : new HashMap<>(); @@ -96,7 +98,7 @@ public void testTimestampMergesPrimitiveAndGenericStates() { right.put(Long.MIN_VALUE, 1L); right.put(Long.MAX_VALUE, 3L); - Map merged = minFunction.merge(left, right); + Map merged = minFunction.merge(left, right); assertSame(merged, left); assertEquals(merged, Map.of(Long.MIN_VALUE, 3L, 0L, 1L, Long.MAX_VALUE, 3L)); assertEquals(minFunction.extractFinalResult(merged), Long.valueOf(Long.MIN_VALUE)); @@ -106,9 +108,34 @@ public void testTimestampMergesPrimitiveAndGenericStates() { } } + @Test + public void testStringMergesPrimitiveAndGenericStates() { + ModeAggregationFunction minFunction = new ModeAggregationFunction(typedArguments("MIN", "STRING"), true); + ModeAggregationFunction maxFunction = new ModeAggregationFunction(typedArguments("MAX", "STRING"), true); + for (boolean primitiveLeft : new boolean[]{false, true}) { + for (boolean primitiveRight : new boolean[]{false, true}) { + Map left = primitiveLeft ? new StringModeCounts() : new HashMap<>(); + Map right = primitiveRight ? new Object2LongOpenHashMap<>() : new HashMap<>(); + left.put("alpha", 2L); + right.put("alpha", 1L); + right.put("zebra", 3L); + + Map merged = minFunction.merge(left, right); + assertSame(merged, left); + assertEquals(merged, Map.of("alpha", 3L, "zebra", 3L)); + assertEquals(minFunction.extractFinalResult(merged), "alpha"); + assertEquals(maxFunction.extractFinalResult(merged), "zebra"); + assertEquals(right, Map.of("alpha", 1L, "zebra", 3L)); + if (primitiveLeft) { + assertEquals(((StringModeCounts) left).getRetainedStringBytes(), 2L * (48 + 2 * 5)); + } + } + } + } + @Test public void testStringModeSkipsNullRowsForMultiValueGroupKeys() { - ModeStringAggregationFunction function = new ModeStringAggregationFunction(List.of(EXPRESSION), true); + ModeAggregationFunction function = new ModeAggregationFunction(typedArguments("MIN", "STRING"), true); GroupByResultHolder holder = function.createGroupByResultHolder(3, 3); BlockValSet values = SyntheticBlockValSets.Str.create(RoaringBitmap.bitmapOf(0, 2), new String[]{"ignored", "alpha", "ignored", "beta", "alpha"}); @@ -123,10 +150,10 @@ public void testStringModeSkipsNullRowsForMultiValueGroupKeys() { @Test public void testEmptyResultsAreNullWithEitherNullHandlingMode() { for (boolean nullHandlingEnabled : new boolean[]{false, true}) { - ModeStringAggregationFunction stringFunction = - new ModeStringAggregationFunction(List.of(EXPRESSION), nullHandlingEnabled); - ModeTimestampAggregationFunction timestampFunction = - new ModeTimestampAggregationFunction(List.of(EXPRESSION), nullHandlingEnabled); + ModeAggregationFunction stringFunction = + new ModeAggregationFunction(typedArguments("MIN", "STRING"), nullHandlingEnabled); + ModeAggregationFunction timestampFunction = + new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), nullHandlingEnabled); assertNull(stringFunction.extractFinalResult(null)); assertNull(timestampFunction.extractFinalResult(null)); assertNull(timestampFunction.extractFinalResult(new Long2LongOpenHashMap())); @@ -141,17 +168,75 @@ public void testEmptyResultsAreNullWithEitherNullHandlingMode() { @Test public void testNonNumericModesRejectAverageReducer() { IllegalArgumentException stringError = expectThrows(IllegalArgumentException.class, - () -> new ModeStringAggregationFunction(argumentsWithReducer("AVG"), true)); + () -> new ModeAggregationFunction(typedArguments("AVG", "STRING"), true)); assertTrue(stringError.getMessage().contains("AVG")); assertTrue(stringError.getMessage().contains("STRING")); IllegalArgumentException timestampError = expectThrows(IllegalArgumentException.class, - () -> new ModeTimestampAggregationFunction(argumentsWithReducer("AVG"), true)); + () -> new ModeAggregationFunction(typedArguments("AVG", "TIMESTAMP"), true)); assertTrue(timestampError.getMessage().contains("AVG")); assertTrue(timestampError.getMessage().contains("TIMESTAMP")); } - private static List argumentsWithReducer(String reducer) { - return List.of(EXPRESSION, ExpressionContext.forLiteral(Literal.stringValue(reducer))); + @Test + public void testResultTypeMustBeSupportedStringLiteral() { + ExpressionContext reducer = ExpressionContext.forLiteral(Literal.stringValue("MIN")); + for (ExpressionContext type : List.of(EXPRESSION, ExpressionContext.forLiteral(Literal.intValue(1)), + ExpressionContext.forLiteral(DataType.STRING, null), + ExpressionContext.forLiteral(Literal.stringValue("LONG")), + ExpressionContext.forLiteral(Literal.stringValue("DOUBLE")), + ExpressionContext.forLiteral(Literal.stringValue("INVALID")))) { + IllegalArgumentException error = expectThrows(IllegalArgumentException.class, + () -> new ModeAggregationFunction(List.of(EXPRESSION, reducer, type), true)); + assertTrue(error.getMessage().contains("MODE result type")); + } + expectThrows(IllegalArgumentException.class, () -> new ModeAggregationFunction(List.of(), true)); + expectThrows(IllegalArgumentException.class, + () -> new ModeAggregationFunction(List.of(EXPRESSION, reducer, reducer, reducer), true)); + expectThrows(IllegalArgumentException.class, + () -> new ModeAggregationFunction(List.of(EXPRESSION, EXPRESSION, + ExpressionContext.forLiteral(Literal.stringValue("STRING"))), true)); + } + + @Test + public void testResultTypeLiteralIsCaseInsensitive() { + for (String type : List.of("string", "StRiNg", "timestamp", "TiMeStAmP")) { + ModeAggregationFunction function = new ModeAggregationFunction(typedArguments("MIN", type), true); + assertEquals(function.getFinalResultColumnType(), + type.equalsIgnoreCase("string") ? ColumnDataType.STRING : ColumnDataType.TIMESTAMP); + assertEquals(function.getResultColumnName(), "mode(value,'MIN','" + type + "')"); + } + } + + @Test + public void testLegacyNumericResultTypeAndReducersAreUnchanged() { + // The same stored longs are DOUBLE for legacy calls and exact TIMESTAMP values only with the inferred type. + long earlier = 9_007_199_254_740_992L; + long later = earlier + 1; + ModeAggregationFunction numeric = new ModeAggregationFunction(List.of(EXPRESSION), false); + ModeAggregationFunction timestamp = new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), false); + assertEquals(numeric.getFinalResultColumnType(), ColumnDataType.DOUBLE); + assertEquals(timestamp.getFinalResultColumnType(), ColumnDataType.TIMESTAMP); + Map counts = aggregateAndRoundTrip(timestamp, timestampValues(later, later, earlier), 3); + assertEquals(numeric.extractFinalResult(counts), Double.valueOf(later)); + assertEquals(timestamp.extractFinalResult(counts), Long.valueOf(later)); + assertEquals(numeric.getFinalResultColumnType(), ColumnDataType.DOUBLE); + assertEquals(timestamp.getFinalResultColumnType(), ColumnDataType.TIMESTAMP); + + for (String reducer : List.of("MIN", "MAX", "AVG")) { + ModeAggregationFunction function = new ModeAggregationFunction(List.of(EXPRESSION, + ExpressionContext.forLiteral(Literal.stringValue(reducer))), false); + assertEquals(function.getResultColumnName(), "mode(value)"); + Map tiedCounts = aggregateAndRoundTrip(function, timestampValues(2, 4), 2); + double expected = reducer.equals("MIN") ? 2D : reducer.equals("MAX") ? 4D : 3D; + assertEquals(function.extractFinalResult(tiedCounts), expected); + assertEquals(function.extractFinalResult(new Long2LongOpenHashMap()), Double.NEGATIVE_INFINITY); + assertNull(function.extractFinalResult(null)); + } + } + + private static List typedArguments(String reducer, String type) { + return List.of(EXPRESSION, ExpressionContext.forLiteral(Literal.stringValue(reducer)), + ExpressionContext.forLiteral(Literal.stringValue(type))); } private static BlockValSet timestampValues(long... values) { @@ -162,14 +247,14 @@ private static BlockValSet timestampValues(long... values) { return blockValSet; } - private static > I aggregateAndRoundTrip(AggregationFunction function, + private static Map aggregateAndRoundTrip(ModeAggregationFunction function, BlockValSet values, int length) { AggregationResultHolder holder = function.createAggregationResultHolder(); function.aggregate(length, holder, Map.of(EXPRESSION, values)); - I intermediateResult = function.extractAggregationResult(holder); + Map intermediateResult = function.extractAggregationResult(holder); AggregationFunction.SerializedIntermediateResult serialized = function.serializeIntermediateResult(intermediateResult); - I deserialized = function.deserializeIntermediateResult( + Map deserialized = function.deserializeIntermediateResult( new CustomObject(serialized.getType(), ByteBuffer.wrap(serialized.getBytes()))); assertEquals(deserialized, intermediateResult); return deserialized; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java index 8a28bbcdf7ac..04dd2d981d05 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java @@ -45,17 +45,17 @@ public class ModeAggregationFunctionRewriteOptimizerTest { @DataProvider public Object[][] modeExpressions() { return new Object[][]{ - {"MODE(stringCol)", "modeString(stringCol)"}, - {"MODE(timestampCol, 'MAX')", "modeTimestamp(timestampCol, 'MAX')"}, - {"MODE(CONCAT(stringCol, 'suffix'))", "modeString(CONCAT(stringCol, 'suffix'))"}, + {"MODE(stringCol)", "MODE(stringCol, 'MIN', 'STRING')"}, + {"MODE(timestampCol, 'MAX')", "MODE(timestampCol, 'MAX', 'TIMESTAMP')"}, + {"MODE(CONCAT(stringCol, 'suffix'))", "MODE(CONCAT(stringCol, 'suffix'), 'MIN', 'STRING')"}, {"MODE(JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', ''))", - "modeString(JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', ''))"}, - {"MODE(CAST(intCol AS STRING))", "modeString(CAST(intCol AS STRING))"}, - {"MODE(CAST(stringCol AS TIMESTAMP))", "modeTimestamp(CAST(stringCol AS TIMESTAMP))"}, + "MODE(JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', ''), 'MIN', 'STRING')"}, + {"MODE(CAST(intCol AS STRING))", "MODE(CAST(intCol AS STRING), 'MIN', 'STRING')"}, + {"MODE(CAST(stringCol AS TIMESTAMP))", "MODE(CAST(stringCol AS TIMESTAMP), 'MIN', 'TIMESTAMP')"}, {"MODE(CASE WHEN intCol > 0 THEN stringCol ELSE 'other' END)", - "modeString(CASE WHEN intCol > 0 THEN stringCol ELSE 'other' END)"}, - {"MODE('literal')", "modeString('literal')"}, - {"fromTimestamp(MODE(timestampCol))", "fromTimestamp(modeTimestamp(timestampCol))"} + "MODE(CASE WHEN intCol > 0 THEN stringCol ELSE 'other' END, 'MIN', 'STRING')"}, + {"MODE('literal')", "MODE('literal', 'MIN', 'STRING')"}, + {"fromTimestamp(MODE(timestampCol))", "fromTimestamp(MODE(timestampCol, 'MIN', 'TIMESTAMP'))"} }; } @@ -77,9 +77,9 @@ public void testModeInHavingAndOrderBy() { + "HAVING MODE(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END) = 'value' " + "ORDER BY MODE(timestampCol) DESC", "SET enableTypedMode=true; " - + "SELECT intCol, modeString(stringCol) AS commonValue FROM testTable GROUP BY intCol " - + "HAVING modeString(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END) = 'value' " - + "ORDER BY modeTimestamp(timestampCol) DESC", SCHEMA); + + "SELECT intCol, MODE(stringCol, 'MIN', 'STRING') AS commonValue FROM testTable GROUP BY intCol " + + "HAVING MODE(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END, 'MIN', 'STRING') = 'value' " + + "ORDER BY MODE(timestampCol, 'MIN', 'TIMESTAMP') DESC", SCHEMA); } @Test @@ -101,7 +101,20 @@ public void testExistingRewriteOptionPreservesLegacyMode() { TestHelper.assertEqualsQuery("SET autoRewriteAggregationType=true; SET enableTypedMode=true; " + "SELECT MODE(timestampCol), MODE(stringCol), MIN(stringCol) FROM testTable", "SET autoRewriteAggregationType=true; SET enableTypedMode=true; " - + "SELECT MODETIMESTAMP(timestampCol), MODESTRING(stringCol), MINSTRING(stringCol) FROM testTable", SCHEMA); + + "SELECT MODE(timestampCol, 'MIN', 'TIMESTAMP'), MODE(stringCol, 'MIN', 'STRING'), " + + "MINSTRING(stringCol) FROM testTable", SCHEMA); + } + + @Test + public void testExplicitTypesAndRepeatedOptimization() { + assertUnchanged("SET enableTypedMode=true; SELECT MODE(stringCol, 'MAX', 'STRING'), " + + "MODE(timestampCol, 'MIN', 'TIMESTAMP') FROM testTable", SCHEMA); + PinotQuery query = CalciteSqlParser.compileToPinotQuery( + "SET enableTypedMode=true; SELECT MODE(stringCol), MODE(timestampCol, 'MAX') FROM testTable"); + OPTIMIZER.optimize(query, SCHEMA); + PinotQuery once = query.deepCopy(); + OPTIMIZER.optimize(query, SCHEMA); + assertEquals(query, once); } @Test diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java index 3767e4c07e07..8d3e8349f78c 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java @@ -195,13 +195,22 @@ private BrokerResponseNative getBrokerResponse(@Language("sql") String query, Pl /// This can be particularly useful to test statistical aggregation functions. /// @see StatisticalQueriesTest for an example use case. private BrokerResponseNative getBrokerResponse(PinotQuery pinotQuery, PlanMaker planMaker) { + return getBrokerResponse(pinotQuery, planMaker, false, null); + } + + private BrokerResponseNative getBrokerResponse(PinotQuery pinotQuery, PlanMaker planMaker, boolean optimize, + @Nullable Schema schema) { + // Match the broker's order: retain the original gapfill query while optimizing the stripped server query. + PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); + if (optimize) { + OPTIMIZER.optimize(serverPinotQuery, schema); + } List> instances = getDistinctInstances(); if (instances.size() == 2) { - return getBrokerResponseDistinctInstances(pinotQuery, planMaker); + return getBrokerResponseDistinctInstances(pinotQuery, serverPinotQuery, planMaker); } // Server side - PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); QueryContext serverQueryContext = serverPinotQuery == pinotQuery ? queryContext : QueryContextConverterUtils.getQueryContext(serverPinotQuery); @@ -266,8 +275,7 @@ protected BrokerResponseNative reduceOnDataTable(BrokerRequest brokerRequest, Br protected BrokerResponseNative getBrokerResponseForOptimizedQuery(@Language("sql") String query, @Nullable Schema schema) { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - OPTIMIZER.optimize(pinotQuery, schema); - return getBrokerResponse(pinotQuery, PLAN_MAKER); + return getBrokerResponse(pinotQuery, PLAN_MAKER, true, schema); } /// Run query on multiple index segments with custom plan maker. @@ -280,9 +288,9 @@ protected BrokerResponseNative getBrokerResponseForOptimizedQuery(@Language("sql /// overriding getDistinctInstances. /// This can be particularly useful to test statistical aggregation functions. /// @see StatisticalQueriesTest for an example use case. - private BrokerResponseNative getBrokerResponseDistinctInstances(PinotQuery pinotQuery, PlanMaker planMaker) { + private BrokerResponseNative getBrokerResponseDistinctInstances(PinotQuery pinotQuery, PinotQuery serverPinotQuery, + PlanMaker planMaker) { // Server side - PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); QueryContext serverQueryContext = serverPinotQuery == pinotQuery ? queryContext : QueryContextConverterUtils.getQueryContext(serverPinotQuery); diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java index 82a96a5359e2..79cc27334ea0 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java @@ -38,6 +38,7 @@ import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.DateTimeFormatSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; @@ -48,6 +49,7 @@ import org.testng.Assert; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import static org.apache.pinot.spi.utils.CommonConstants.RewriterConstants.CHILD_AGGREGATION_NAME_PREFIX; @@ -612,6 +614,44 @@ public void testEmptyGroupByInterSegment() { assertEquals(rows.size(), 0); } + @DataProvider + public Object[][] gapfillParentResultTypes() { + return new Object[][]{ + // Parent aggregates expose timestamp children through their stored LONG type. + {TIMESTAMP_COLUMN, 1683138373878L}, + {BYTES_COLUMN, "31"}, + {BIG_DECIMAL_COLUMN, "1199"} + }; + } + + @Test(dataProvider = "gapfillParentResultTypes") + public void testGapfillFormatsNativeParentResults(String column, Object expectedValue) { + String format = "1:MILLISECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss.SSS"; + DateTimeFormatSpec formatter = new DateTimeFormatSpec(format); + long hourMillis = 3_600_000L; + long startMillis = 1683138373878L - 1683138373878L % hourMillis; + String start = formatter.fromMillisToFormat(startMillis); + String end = formatter.fromMillisToFormat(startMillis + 2 * hourMillis); + // Parent aggregate projections are referenced by their canonical names because AS is not supported. + String expression = "exprmin(" + column + ",intColumn)"; + String projection = "\"" + expression + "\""; + String query = "SELECT GapFill(time_col, '" + format + "', '" + start + "', '" + end + + "', '1:HOURS', FILL(" + projection + ", 'FILL_PREVIOUS_VALUE'), TIMESERIESON(groupByIntColumn)), " + + "groupByIntColumn, " + projection + " FROM (SELECT DATETIMECONVERT(fromTimestamp(timestampColumn), " + + "'1:MILLISECONDS:EPOCH', '" + format + "', '1:HOURS') AS time_col, groupByIntColumn, " + + expression + " FROM testTable WHERE intColumn = 1 GROUP BY time_col, groupByIntColumn LIMIT 1000) LIMIT 1000"; + + BrokerResponseNative response = getBrokerResponse(query); + assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); + List rows = response.getResultTable().getRows(); + // The shared fixture returns the parent projection once per server, followed by one filled row. + assertEquals(rows.size(), 3); + assertEquals(rows.get(0), new Object[]{start, 1, expectedValue}); + assertEquals(rows.get(1), new Object[]{start, 1, expectedValue}); + assertEquals(rows.get(2), + new Object[]{formatter.fromMillisToFormat(startMillis + hourMillis), 1, expectedValue}); + } + @Test public void testAlias() { // Using exprmin/exprmax with alias will fail, since the alias will not be resolved by the rewriter diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java index 4641a9ba9945..a7ecf3b3ccd9 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java @@ -48,6 +48,7 @@ import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.DateTimeFormatSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; @@ -464,6 +465,48 @@ public void testTimestampAggregationAndResultType() { new Timestamp(_expectedResultMin.longValue()).toString()}); } + @DataProvider + public Object[][] typedModeGapfillExpressions() { + return new Object[][]{ + {false, "MODE(timestampColumn, 'MIN', 'TIMESTAMP')"}, + {false, "MODE(timestampColumn, 'MIN', 'timestamp')"}, + {true, "MODE(timestampColumn)"} + }; + } + + @Test(dataProvider = "typedModeGapfillExpressions") + public void testTypedModeGapfillPreservesAliases(boolean inferType, String typedMode) { + String format = "1:MILLISECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss.SSS"; + DateTimeFormatSpec formatter = new DateTimeFormatSpec(format); + long hourMillis = 3_600_000L; + long startMillis = BASE_TIMESTAMP - BASE_TIMESTAMP % hourMillis; + String start = formatter.fromMillisToFormat(startMillis); + String end = formatter.fromMillisToFormat(startMillis + 2 * hourMillis); + String entity = Integer.toString(_expectedResultMin.intValue()); + // Keep both the legacy DOUBLE and typed TIMESTAMP result in the same aggregate subquery. + String numericMode = inferType ? "MODE(fromTimestamp(timestampColumn))" : "MODE(timestampColumn)"; + String query = "SET enableTypedMode=" + inferType + "; SELECT GapFill(time_col, '" + format + "', '" + + start + "', '" + end + "', '1:HOURS', FILL(numeric_mode, 'FILL_PREVIOUS_VALUE'), " + + "FILL(typed_mode, 'FILL_PREVIOUS_VALUE'), TIMESERIESON(stringColumn)), " + + "stringColumn, numeric_mode, typed_mode FROM (SELECT DATETIMECONVERT(fromTimestamp(timestampColumn), " + + "'1:MILLISECONDS:EPOCH', '" + format + "', '1:HOURS') AS time_col, stringColumn, " + + numericMode + " AS numeric_mode, " + typedMode + " AS typed_mode FROM testTable WHERE stringColumn = '" + + entity + "' GROUP BY time_col, stringColumn LIMIT 1000) LIMIT 1000"; + BrokerResponseNative response = getBrokerResponseForOptimizedQuery(query, SCHEMA); + + assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); + assertEquals(response.getResultTable().getDataSchema().getColumnNames(), + new String[]{"time_col", "stringColumn", "numeric_mode", "typed_mode"}); + List rows = response.getResultTable().getRows(); + assertEquals(rows.size(), 2); + long expectedMillis = BASE_TIMESTAMP + _expectedResultMin.longValue(); + String expectedTimestamp = new Timestamp(expectedMillis).toString(); + for (int i = 0; i < rows.size(); i++) { + assertEquals(rows.get(i), new Object[]{formatter.fromMillisToFormat(startMillis + i * hourMillis), entity, + (double) expectedMillis, expectedTimestamp}); + } + } + @AfterClass public void tearDown() throws IOException { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java index 6886af299b8d..8d0a83597df3 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java @@ -25,17 +25,19 @@ 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.logical.LogicalAggregate; -import org.apache.calcite.sql.SqlAggFunction; -import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlKind; -import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.sql.type.SqlTypeName; -import org.apache.pinot.common.function.sql.PinotSqlAggFunction; -/// Rewrites string and timestamp MODE calls to typed implementations after an explicit rollout opt-in. -/// Numeric MODE and the separate MIN, MAX and SUM rewrite rule are unaffected. +/// Supplies string and timestamp MODE calls with an inferred type argument after an explicit rollout opt-in. +/// This stateless rule keeps the reducer and type as projected literals so distributed stages retain both arguments. public class PinotModeAggregationFunctionRewriteRule extends RelOptRule { public static PinotModeAggregationFunctionRewriteRule instanceWithDescription(String description) { return new PinotModeAggregationFunctionRewriteRule(description); @@ -48,41 +50,75 @@ private PinotModeAggregationFunctionRewriteRule(String description) { @Override public void onMatch(RelOptRuleCall call) { Aggregate aggregate = call.rel(0); - RelNode input = aggregate.getInput(); + RelNode input = PinotRuleUtils.unboxRel(aggregate.getInput()); + List projects = new ArrayList<>(); + List names = new ArrayList<>(input.getRowType().getFieldNames()); + if (input instanceof Project) { + projects.addAll(((Project) input).getProjects()); + } else { + for (int i = 0; i < names.size(); i++) { + projects.add(RexInputRef.of(i, input.getRowType())); + } + } + RexBuilder rexBuilder = input.getCluster().getRexBuilder(); List originalCalls = aggregate.getAggCallList(); - List rewrittenCalls = new ArrayList<>(originalCalls.size()); + List> rewrittenArguments = new ArrayList<>(originalCalls.size()); boolean changed = false; for (AggregateCall originalCall : originalCalls) { - AggregateCall rewrittenCall = maybeRewriteAggCall(originalCall, input, aggregate.getGroupCount()); - changed |= rewrittenCall != originalCall; - rewrittenCalls.add(rewrittenCall); + List arguments = originalCall.getArgList(); + List rewritten = arguments; + if (originalCall.getAggregation().getKind() == SqlKind.MODE && !arguments.isEmpty() && arguments.size() < 3) { + SqlTypeName operandType = input.getRowType().getFieldList().get(arguments.get(0)).getType().getSqlTypeName(); + String type = SqlTypeName.STRING_TYPES.contains(operandType) + ? "STRING" + : operandType == SqlTypeName.TIMESTAMP ? "TIMESTAMP" : null; + if (type != null) { + rewritten = new ArrayList<>(arguments); + if (arguments.size() == 1) { + rewritten.add(addLiteral(rexBuilder, projects, names, "MIN")); + } + rewritten.add(addLiteral(rexBuilder, projects, names, type)); + changed = true; + } + } + rewrittenArguments.add(rewritten); } - if (changed) { - call.transformTo(aggregate.copy(aggregate.getTraitSet(), input, aggregate.getGroupSet(), aggregate.getGroupSets(), - rewrittenCalls)); + if (!changed) { + return; } - } - private static AggregateCall maybeRewriteAggCall(AggregateCall call, RelNode input, int numGroups) { - SqlAggFunction aggregation = call.getAggregation(); - List arguments = call.getArgList(); - if (aggregation.getKind() != SqlKind.MODE || arguments.isEmpty()) { - return call; - } - SqlTypeName operandType = input.getRowType().getFieldList().get(arguments.get(0)).getType().getSqlTypeName(); - String functionName; - if (SqlTypeName.STRING_TYPES.contains(operandType)) { - functionName = "MODESTRING"; - } else if (operandType == SqlTypeName.TIMESTAMP) { - functionName = "MODETIMESTAMP"; + RelNode rewrittenInput; + if (input instanceof Project) { + // Extend the existing projection: wrapping it in identity refs would hide the original reducer literals. + Project project = (Project) input; + RelDataTypeFactory.Builder rowType = input.getCluster().getTypeFactory().builder(); + for (int i = 0; i < projects.size(); i++) { + rowType.add(names.get(i), projects.get(i).getType()); + } + rewrittenInput = project.copy(project.getTraitSet(), project.getInput(), projects, rowType.build()); } else { - return call; + rewrittenInput = LogicalProject.create(input, List.of(), projects, names); + } + List rewrittenCalls = new ArrayList<>(originalCalls.size()); + for (int i = 0; i < originalCalls.size(); i++) { + AggregateCall original = originalCalls.get(i); + rewrittenCalls.add(AggregateCall.create(original.getAggregation(), original.isDistinct(), + original.isApproximate(), original.ignoreNulls(), rewrittenArguments.get(i), original.filterArg, + original.distinctKeys, original.getCollation(), aggregate.getGroupCount(), rewrittenInput, original.getType(), + original.getName())); + } + call.transformTo(aggregate.copy(aggregate.getTraitSet(), rewrittenInput, aggregate.getGroupSet(), + aggregate.getGroupSets(), rewrittenCalls)); + } + + private static int addLiteral(RexBuilder rexBuilder, List projects, List names, String value) { + RexNode literal = rexBuilder.makeLiteral(value); + int index = projects.indexOf(literal); + if (index < 0) { + index = projects.size(); + projects.add(literal); + names.add("$mode$" + index); } - SqlAggFunction rewrittenAggregation = new PinotSqlAggFunction(functionName, SqlKind.OTHER_FUNCTION, - ReturnTypes.explicit(call.getType()), aggregation.getOperandTypeChecker(), - SqlFunctionCategory.USER_DEFINED_FUNCTION); - return AggregateCall.create(rewrittenAggregation, call.isDistinct(), call.isApproximate(), call.ignoreNulls(), - arguments, call.filterArg, call.distinctKeys, call.getCollation(), numGroups, input, call.getType(), - call.getName()); + return index; } } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java index 8a831dba927a..116c2006daa0 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java @@ -33,6 +33,7 @@ import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.routing.WorkerManager; +import org.apache.pinot.spi.exception.QueryException; import org.apache.pinot.spi.utils.CommonConstants; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -40,9 +41,10 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; -/// Verifies MODE type inference and type-specific dispatch across both multi-stage planner implementations. +/// Verifies MODE type inference and inferred type arguments across both multi-stage planner implementations. public class ModeSqlPlannerTest extends QueryEnvironmentTestBase { @DataProvider public Object[][] physicalOptimizers() { @@ -89,7 +91,10 @@ public void testDistributedModeTypes(boolean usePhysicalOptimizer) { for (AggregateNode aggregate : aggregates) { List calls = aggregate.getAggCalls(); assertEquals(calls.stream().map(RexExpression.FunctionCall::getFunctionName).toList(), - List.of("MODESTRING", "MODETIMESTAMP", "MODE")); + List.of("MODE", "MODE", "MODE")); + assertTypedCall(calls.get(0), "MIN", "STRING"); + assertTypedCall(calls.get(1), "MIN", "TIMESTAMP"); + assertEquals(calls.get(2).getFunctionOperands().size(), 1); if (aggregate.getAggType().isOutputIntermediateFormat() && !aggregate.isLeafReturnFinalResult()) { sawIntermediate = true; assertEquals(aggregate.getDataSchema().getColumnDataTypes(), @@ -116,7 +121,10 @@ public void testModeExpressionsAndTieBreakers(boolean usePhysicalOptimizer) { ColumnDataType.DOUBLE}); for (AggregateNode aggregate : findAggregates(plan)) { assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), - List.of("MODESTRING", "MODETIMESTAMP", "MODE")); + List.of("MODE", "MODE", "MODE")); + assertTypedCall(aggregate.getAggCalls().get(0), "MIN", "STRING"); + assertTypedCall(aggregate.getAggCalls().get(1), "MAX", "TIMESTAMP"); + assertEquals(aggregate.getAggCalls().get(2).getFunctionOperands().size(), 2); } } @@ -131,6 +139,9 @@ public void testModeRewriteRequiresItsOwnOptIn(boolean usePhysicalOptimizer) { for (AggregateNode aggregate : aggregates) { assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), List.of("MODE", "MODE"), options); + for (RexExpression.FunctionCall mode : aggregate.getAggCalls()) { + assertEquals(mode.getFunctionOperands().size(), 1, options); + } } } } @@ -145,9 +156,11 @@ public void testTypedModeDoesNotEnableOtherAggregateRewrites(boolean usePhysical for (AggregateNode aggregate : aggregates) { List functionNames = aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(); - assertEquals(functionNames.subList(0, 5), List.of("MODESTRING", "MODETIMESTAMP", "MODE", "MIN", "MAX")); + assertEquals(functionNames.subList(0, 5), List.of("MODE", "MODE", "MODE", "MIN", "MAX")); assertFalse(functionNames.contains("SUMLONG")); assertFalse(functionNames.contains("SUMINT")); + assertTypedCall(aggregate.getAggCalls().get(0), "MIN", "STRING"); + assertTypedCall(aggregate.getAggCalls().get(1), "MIN", "TIMESTAMP"); } } @@ -160,10 +173,9 @@ public void testTypedModeOptInSurvivesCustomizedPlannerDefaults(boolean usePhysi "SET autoRewriteAggregationType=true; ", "SET usePlannerRules='TypedModeRewrite'; ", "SET usePlannerRules='TypedModeRewrite'; SET enableTypedMode=false; ", "SET enableTypedMode=true; SET skipPlannerRules='TypedModeRewrite'; ")) { - assertModeCalls(environment, usePhysicalOptimizer, options, List.of("MODE", "MODE", "MODE")); + assertModeCalls(environment, usePhysicalOptimizer, options, false); } - assertModeCalls(environment, usePhysicalOptimizer, "SET enableTypedMode=true; ", - List.of("MODESTRING", "MODETIMESTAMP", "MODE")); + assertModeCalls(environment, usePhysicalOptimizer, "SET enableTypedMode=true; ", true); } } @@ -182,17 +194,67 @@ private static QueryEnvironment buildQueryEnvironment(Set disabledRules) } private static void assertModeCalls(QueryEnvironment environment, boolean usePhysicalOptimizer, String options, - List expectedCalls) { + boolean typed) { DispatchableSubPlan plan = environment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + options + "SELECT MODE(col1), MODE(ts_timestamp), MODE(col3) FROM a"); List aggregates = findAggregates(plan); assertFalse(aggregates.isEmpty()); for (AggregateNode aggregate : aggregates) { assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), - expectedCalls, options); + List.of("MODE", "MODE", "MODE"), options); + if (typed) { + assertTypedCall(aggregate.getAggCalls().get(0), "MIN", "STRING"); + assertTypedCall(aggregate.getAggCalls().get(1), "MIN", "TIMESTAMP"); + } else { + for (RexExpression.FunctionCall mode : aggregate.getAggCalls()) { + assertEquals(mode.getFunctionOperands().size(), 1, options); + } + } } } + @Test(dataProvider = "physicalOptimizers") + public void testExplicitTypeArguments(boolean usePhysicalOptimizer) { + DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + + "SELECT MODE(col1, 'MAX', 'STRING'), MODE(ts_timestamp, 'MIN', 'TIMESTAMP') FROM a"); + for (AggregateNode aggregate : findAggregates(plan)) { + assertTypedCall(aggregate.getAggCalls().get(0), "MAX", "STRING"); + assertTypedCall(aggregate.getAggCalls().get(1), "MIN", "TIMESTAMP"); + } + PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); + assertEquals(root.getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP}); + } + + @Test + public void testInvalidTypeAnnotations() { + for (String expression : List.of("MODE(col1, 'MIN', 'TIMESTAMP')", "MODE(ts_timestamp, 'MIN', 'STRING')", + "MODE(col1, 'MIN', 'INVALID')", "MODE(col1, 'MIN', col1)")) { + QueryException error = expectThrows(QueryException.class, + () -> _queryEnvironment.compile("SELECT " + expression + " FROM a")); + assertTrue(error.getMessage().contains("MODE type argument"), error.getMessage()); + } + } + + @Test + public void testTypeAnnotationsAreCaseInsensitive() { + RelDataType rowType = _queryEnvironment.compile( + "SELECT MODE(col1, 'MIN', 'string'), MODE(ts_timestamp, 'MAX', 'TimeStamp') FROM a") + .getRelRoot().validatedRowType; + assertEquals(rowType.getFieldList().get(0).getType().getSqlTypeName(), SqlTypeName.VARCHAR); + assertEquals(rowType.getFieldList().get(1).getType().getSqlTypeName(), SqlTypeName.TIMESTAMP); + } + + private static void assertTypedCall(RexExpression.FunctionCall call, String reducer, String type) { + assertEquals(call.getFunctionName(), "MODE"); + List arguments = call.getFunctionOperands(); + assertEquals(arguments.size(), 3); + assertTrue(arguments.get(1) instanceof RexExpression.Literal); + assertTrue(arguments.get(2) instanceof RexExpression.Literal); + assertEquals(((RexExpression.Literal) arguments.get(1)).getValue(), reducer); + assertEquals(((RexExpression.Literal) arguments.get(2)).getValue(), type); + } + private static List findAggregates(DispatchableSubPlan plan) { List aggregates = new ArrayList<>(); for (DispatchablePlanFragment fragment : plan.getQueryStageMap().values()) { diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java index ebbb4551774f..8acee5ad7105 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java @@ -20,12 +20,14 @@ import java.util.Arrays; import java.util.List; +import java.util.Locale; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nullable; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.SqlCallBinding; import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.SqlOperatorBinding; import org.apache.calcite.sql.type.OperandTypes; @@ -38,6 +40,8 @@ import org.apache.commons.lang3.Strings; import org.apache.pinot.spi.utils.CommonConstants; +import static com.google.common.base.Preconditions.checkArgument; + /// NOTES: /// - No underscore is allowed in the enum name. @@ -67,12 +71,9 @@ public enum AggregationFunctionType { OperandTypes.NUMERIC, OperandTypes.CHARACTER, OperandTypes.TIMESTAMP, OperandTypes.family(List.of(SqlTypeFamily.NUMERIC, SqlTypeFamily.CHARACTER), i -> i == 1), OperandTypes.family(List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER), i -> i == 1), - OperandTypes.family(List.of(SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER), i -> i == 1)), + OperandTypes.family(List.of(SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER), i -> i == 1), + OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER)), ReturnTypes.explicit(SqlTypeName.OTHER), null, SqlKind.MODE), - MODESTRING("modeString", ReturnTypes.ARG0_NULLABLE_IF_EMPTY, - OperandTypes.family(List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER), i -> i == 1), SqlTypeName.OTHER), - MODETIMESTAMP("modeTimestamp", ReturnTypes.ARG0_NULLABLE_IF_EMPTY, - OperandTypes.family(List.of(SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER), i -> i == 1), SqlTypeName.OTHER), ANYVALUE("anyValue", ReturnTypes.ARG0, OperandTypes.ANY, SqlTypeName.OTHER), FIRSTWITHTIME("firstWithTime", ReturnTypes.ARG0, OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER), SqlTypeName.OTHER), @@ -427,6 +428,17 @@ private static class ModeReturnTypeInference implements SqlReturnTypeInference { @Override public RelDataType inferReturnType(SqlOperatorBinding opBinding) { RelDataType operandType = opBinding.getOperandType(0); + // Aggregate bindings contain operand types only. Validate the annotation while SQL literals are available; + // later planning derives the same result type from the original input and carries it through merge stages. + if (opBinding.getOperandCount() == 3 && opBinding instanceof SqlCallBinding) { + checkArgument(opBinding.isOperandLiteral(2, false), + "MODE type argument must be a STRING or TIMESTAMP literal"); + String type = opBinding.getOperandLiteralValue(2, String.class); + type = type != null ? type.toUpperCase(Locale.ROOT) : null; + boolean matches = "STRING".equals(type) && SqlTypeName.STRING_TYPES.contains(operandType.getSqlTypeName()) + || "TIMESTAMP".equals(type) && operandType.getSqlTypeName() == SqlTypeName.TIMESTAMP; + checkArgument(matches, "MODE type argument must match the STRING or TIMESTAMP input, got: %s", type); + } if (SqlTypeName.STRING_TYPES.contains(operandType.getSqlTypeName()) || operandType.getSqlTypeName() == SqlTypeName.TIMESTAMP) { return ReturnTypes.ARG0_NULLABLE_IF_EMPTY.inferReturnType(opBinding); From 1dc9259a4acf250a5b2245500d206d5f15bce2b8 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 7 Sep 2026 14:57:47 -0700 Subject: [PATCH 5/8] Keep typed MODE focused on core aggregation and inference --- .../function/ModeAggregationFunction.java | 347 ++++-------------- ...deAggregationFunctionRewriteOptimizer.java | 33 +- .../query/reduce/BaseGapfillProcessor.java | 11 - .../core/query/reduce/GapfillProcessor.java | 5 +- .../AggregationFunctionFactoryTest.java | 11 - .../function/ModeAggregationFunctionTest.java | 141 +++---- ...ModeNonNumericAggregationFunctionTest.java | 262 ------------- ...gregationFunctionRewriteOptimizerTest.java | 88 +---- .../apache/pinot/queries/BaseQueriesTest.java | 20 +- .../apache/pinot/queries/ExprMinMaxTest.java | 40 -- .../apache/pinot/queries/ModeQueriesTest.java | 140 +------ ...notModeAggregationFunctionRewriteRule.java | 12 +- .../query/queries/ModeSqlPlannerTest.java | 130 +------ .../resources/queries/ModeAggregates.json | 130 +------ 14 files changed, 186 insertions(+), 1184 deletions(-) delete mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java index 02440f4755e8..91e10fbb3833 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunction.java @@ -29,10 +29,7 @@ import it.unimi.dsi.fastutil.ints.Int2LongMap; import it.unimi.dsi.fastutil.ints.Int2LongOpenHashMap; import it.unimi.dsi.fastutil.longs.Long2LongMap; -import it.unimi.dsi.fastutil.longs.Long2LongMaps; import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; -import it.unimi.dsi.fastutil.objects.Object2LongMap; -import it.unimi.dsi.fastutil.objects.Object2LongMaps; import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap; import it.unimi.dsi.fastutil.objects.ObjectIterator; import java.util.List; @@ -102,7 +99,7 @@ public ModeAggregationFunction(List arguments, boolean nullHa Preconditions.checkArgument( _resultType == ColumnDataType.DOUBLE || _multiModeReducerType != MultiModeReducerType.AVG, "MODE for %s supports only MIN or MAX tie reducers, got: %s", _resultType, _multiModeReducerType); - // Gapfill resolves aliases by matching the complete selection expression, including literal spelling. + // Include the inferred type in the result identity while retaining legacy numeric names. _resultColumnName = numArguments == 3 ? "mode(" + _expression + "," + arguments.get(1) + "," + arguments.get(2) + ")" : super.getResultColumnName(); @@ -115,7 +112,7 @@ private static ExpressionContext checkArguments(List argument } /// Helper method to create a value map for the given value type. - private static Map getValueMap(DataType valueType) { + private static Map getValueMap(DataType valueType) { switch (valueType) { case INT: return new Int2LongOpenHashMap(); @@ -125,15 +122,17 @@ private static Map getValueMap(DataType valueType) { return new Float2LongOpenHashMap(); case DOUBLE: return new Double2LongOpenHashMap(); + case STRING: + return new Object2LongOpenHashMap(); default: throw new IllegalStateException("Illegal data type for MODE aggregation function: " + valueType); } } /// Returns the value map from the result holder or creates a new one if it does not exist. - private static Map getValueMap(AggregationResultHolder aggregationResultHolder, + private static Map getValueMap(AggregationResultHolder aggregationResultHolder, DataType valueType) { - Map valueMap = aggregationResultHolder.getResult(); + Map valueMap = aggregationResultHolder.getResult(); if (valueMap == null) { valueMap = getValueMap(valueType); aggregationResultHolder.setValue(valueMap); @@ -181,6 +180,15 @@ private static void setValueForGroupKeys(GroupByResultHolder groupByResultHolder valueMap.merge(value, 1, Long::sum); } + private static void setValueForGroupKeys(GroupByResultHolder holder, int groupKey, String value) { + Object2LongOpenHashMap counts = holder.getResult(groupKey); + if (counts == null) { + counts = new Object2LongOpenHashMap<>(); + holder.setValueForKey(groupKey, counts); + } + counts.addTo(value, 1L); + } + /// Returns the dictionary id count map from the result holder or creates a new one if it does not exist. protected static Int2IntOpenHashMap getDictIdCountMap(AggregationResultHolder aggregationResultHolder, Dictionary dictionary) { @@ -204,7 +212,7 @@ protected static Int2IntOpenHashMap getDictIdCountMap(GroupByResultHolder groupB } /// Helper method to read dictionary and convert dictionary ids to values for dictionary-encoded expression. - private static Map convertToValueMap(DictIdsWrapper dictIdsWrapper) { + private static Map convertToValueMap(DictIdsWrapper dictIdsWrapper) { Dictionary dictionary = dictIdsWrapper._dictionary; Int2IntOpenHashMap dictIdCountMap = dictIdsWrapper._dictIdCountMap; int numValues = dictIdCountMap.size(); @@ -239,6 +247,13 @@ private static Map convertToValueMap(DictIdsWrapper dict doubleValueMap.put(dictionary.getDoubleValue(next.getIntKey()), next.getIntValue()); } return doubleValueMap; + case STRING: + Object2LongOpenHashMap stringValueMap = new Object2LongOpenHashMap<>(numValues); + while (iterator.hasNext()) { + Int2IntMap.Entry next = iterator.next(); + stringValueMap.put(dictionary.getStringValue(next.getIntKey()), next.getIntValue()); + } + return stringValueMap; default: throw new IllegalStateException("Illegal data type for MODE aggregation function: " + storedType); } @@ -247,12 +262,9 @@ private static Map convertToValueMap(DictIdsWrapper dict /// Helper method to extract segment level intermediate result from the inner segment result. @Nullable private Map extractIntermediateResult(@Nullable Object result) { - if (_resultType != ColumnDataType.DOUBLE) { - return extractComparableCounts(result); - } if (result == null) { - // NOTE: Return an empty Int2LongOpenHashMap for empty result. - return new Int2LongOpenHashMap(); + // Preserve the legacy numeric empty-result sentinel. + return _resultType == ColumnDataType.DOUBLE ? new Int2LongOpenHashMap() : null; } if (result instanceof DictIdsWrapper) { @@ -287,10 +299,6 @@ public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int ma @Override public void aggregate(int length, AggregationResultHolder aggregationResultHolder, Map blockValSetMap) { - if (_resultType != ColumnDataType.DOUBLE) { - aggregateComparable(length, aggregationResultHolder, blockValSetMap); - return; - } BlockValSet blockValSet = blockValSetMap.get(_expression); // For dictionary-encoded expression, store dictionary ids into the dictId map @@ -309,7 +317,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde // For non-dictionary-encoded expression, store values into the value map DataType storedType = blockValSet.getValueType().getStoredType(); - Map valueMap = getValueMap(aggregationResultHolder, storedType); + Map valueMap = getValueMap(aggregationResultHolder, storedType); switch (storedType) { case INT: Int2LongOpenHashMap intMap = (Int2LongOpenHashMap) valueMap; @@ -347,6 +355,15 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } }); break; + case STRING: + Object2LongOpenHashMap stringMap = (Object2LongOpenHashMap) valueMap; + String[] stringValues = blockValSet.getStringValuesSV(); + forEachNotNull(length, blockValSet, (from, to) -> { + for (int i = from; i < to; i++) { + stringMap.addTo(stringValues[i], 1L); + } + }); + break; default: throw new IllegalStateException("Illegal data type for MODE aggregation function: " + storedType); } @@ -355,10 +372,6 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde @Override public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, Map blockValSetMap) { - if (_resultType != ColumnDataType.DOUBLE) { - aggregateComparableGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSetMap); - return; - } BlockValSet blockValSet = blockValSetMap.get(_expression); // For dictionary-encoded expression, store dictionary ids into the dictId map @@ -409,6 +422,14 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } }); break; + case STRING: + String[] stringValues = blockValSet.getStringValuesSV(); + forEachNotNull(length, blockValSet, (from, to) -> { + for (int i = from; i < to; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeyArray[i], stringValues[i]); + } + }); + break; default: throw new IllegalStateException("Illegal data type for MODE aggregation function: " + storedType); } @@ -417,10 +438,6 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol @Override public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, Map blockValSetMap) { - if (_resultType != ColumnDataType.DOUBLE) { - aggregateComparableGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSetMap); - return; - } BlockValSet blockValSet = blockValSetMap.get(_expression); // For dictionary-encoded expression, store dictionary ids into the dictId map @@ -480,6 +497,16 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } }); break; + case STRING: + String[] stringValues = blockValSet.getStringValuesSV(); + forEachNotNull(length, blockValSet, (from, to) -> { + for (int i = from; i < to; i++) { + for (int groupKey : groupKeysArray[i]) { + setValueForGroupKeys(groupByResultHolder, groupKey, stringValues[i]); + } + } + }); + break; default: throw new IllegalStateException("Illegal data type for MODE aggregation function: " + storedType); } @@ -499,15 +526,17 @@ public Map extractGroupByResult(GroupByResultHolder groupByResultHolder @Override public Map merge(Map intermediateResult1, Map intermediateResult2) { - if (_resultType != ColumnDataType.DOUBLE) { - return mergeComparableCounts(intermediateResult1, intermediateResult2); - } if (intermediateResult1.isEmpty()) { return intermediateResult2; } if (intermediateResult2.isEmpty()) { return intermediateResult1; } + if (_resultType != ColumnDataType.DOUBLE) { + Map counts = (Map) intermediateResult1; + intermediateResult2.forEach((value, count) -> counts.merge(value, count, Long::sum)); + return counts; + } if (intermediateResult1 instanceof Int2LongOpenHashMap && intermediateResult2 instanceof Int2LongOpenHashMap) { ((Int2LongOpenHashMap) intermediateResult2).int2LongEntrySet().fastForEach( e -> ((Int2LongOpenHashMap) intermediateResult1).merge(e.getIntKey(), e.getLongValue(), Long::sum)); @@ -568,8 +597,7 @@ public SerializedIntermediateResult serializeIntermediateResult(Map lon @Override public Map deserializeIntermediateResult(CustomObject customObject) { - Map counts = ObjectSerDeUtils.deserialize(customObject); - return _resultType == ColumnDataType.STRING ? new StringModeCounts((Map) counts) : counts; + return ObjectSerDeUtils.deserialize(customObject); } @Override @@ -796,266 +824,25 @@ private Double extractNumericFinalResult(Double2LongOpenHashMap intermediateResu } } - private Map newComparableValueMap() { - return _resultType == ColumnDataType.STRING ? new StringModeCounts() : new Long2LongOpenHashMap(); - } - - private ValueCounter comparableValueCounter(BlockValSet blockValSet) { - if (_resultType == ColumnDataType.STRING) { - String[] values = blockValSet.getStringValuesSV(); - return (counts, row) -> ((StringModeCounts) counts).addTo(values[row], 1L); - } - long[] values = blockValSet.getLongValuesSV(); - return (counts, row) -> ((Long2LongOpenHashMap) counts).addTo(values[row], 1L); - } - - @FunctionalInterface - private interface ValueCounter { - void add(Map counts, int row); - } - - private void aggregateComparable(int length, AggregationResultHolder holder, - Map blockValSetMap) { - BlockValSet values = blockValSetMap.get(_expression); - Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; - if (dictionary != null) { - int[] ids = values.getDictionaryIdsSV(); - forEachNotNull(length, values, (from, to) -> { - DictionaryCounts counts = getValue(holder, () -> new DictionaryCounts(dictionary)); - for (int i = from; i < to; i++) { - counts._counts.addTo(ids[i], 1L); - } - }); - } else { - ValueCounter counter = comparableValueCounter(values); - forEachNotNull(length, values, (from, to) -> { - Map counts = getValue(holder, this::newComparableValueMap); - for (int i = from; i < to; i++) { - counter.add(counts, i); - } - }); - } - } - - private void aggregateComparableGroupBySV(int length, int[] groupKeys, GroupByResultHolder holder, - Map blockValSetMap) { - BlockValSet values = blockValSetMap.get(_expression); - Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; - if (dictionary != null) { - int[] ids = values.getDictionaryIdsSV(); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - DictionaryCounts counts = getValue(holder, groupKeys[i], () -> new DictionaryCounts(dictionary)); - counts._counts.addTo(ids[i], 1L); - } - }); - } else { - ValueCounter counter = comparableValueCounter(values); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - Map counts = getValue(holder, groupKeys[i], this::newComparableValueMap); - counter.add(counts, i); - } - }); - } - } - - private void aggregateComparableGroupByMV(int length, int[][] groupKeys, GroupByResultHolder holder, - Map blockValSetMap) { - BlockValSet values = blockValSetMap.get(_expression); - Dictionary dictionary = values.isDictionaryEncoded() ? values.getDictionary() : null; - if (dictionary != null) { - int[] ids = values.getDictionaryIdsSV(); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - for (int groupKey : groupKeys[i]) { - DictionaryCounts counts = getValue(holder, groupKey, () -> new DictionaryCounts(dictionary)); - counts._counts.addTo(ids[i], 1L); - } - } - }); - } else { - ValueCounter counter = comparableValueCounter(values); - forEachNotNull(length, values, (from, to) -> { - for (int i = from; i < to; i++) { - for (int groupKey : groupKeys[i]) { - Map counts = getValue(holder, groupKey, this::newComparableValueMap); - counter.add(counts, i); - } - } - }); - } - } - - @Nullable - private Map extractComparableCounts(@Nullable Object result) { - if (!(result instanceof DictionaryCounts)) { - return (Map) result; - } - DictionaryCounts dictionaryCounts = (DictionaryCounts) result; - Dictionary dictionary = dictionaryCounts._dictionary; - if (_resultType == ColumnDataType.STRING) { - StringModeCounts counts = new StringModeCounts(); - dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> - counts.put(dictionary.getStringValue(entry.getIntKey()), entry.getLongValue())); - return counts; - } - Long2LongOpenHashMap counts = new Long2LongOpenHashMap(); - dictionaryCounts._counts.int2LongEntrySet().fastForEach(entry -> - counts.put(dictionary.getLongValue(entry.getIntKey()), entry.getLongValue())); - return counts; - } - - private Map mergeComparableCounts(Map left, Map right) { - if (_resultType == ColumnDataType.STRING && left instanceof Object2LongOpenHashMap - && right instanceof Object2LongMap) { - Object2LongOpenHashMap counts = (Object2LongOpenHashMap) left; - ObjectIterator> iterator = - Object2LongMaps.fastIterator((Object2LongMap) right); - while (iterator.hasNext()) { - Object2LongMap.Entry entry = iterator.next(); - counts.addTo(entry.getKey(), entry.getLongValue()); - } - return left; - } - if (_resultType == ColumnDataType.TIMESTAMP && left instanceof Long2LongOpenHashMap - && right instanceof Long2LongMap) { - Long2LongOpenHashMap counts = (Long2LongOpenHashMap) left; - ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) right); - while (iterator.hasNext()) { - Long2LongMap.Entry entry = iterator.next(); - counts.addTo(entry.getLongKey(), entry.getLongValue()); - } - return left; - } - Map counts = (Map) left; - right.forEach((value, count) -> counts.merge(value, count, Long::sum)); - return left; - } - @Nullable private Comparable extractComparableFinalResult(@Nullable Map counts) { - if (counts == null || counts.isEmpty()) { - return null; - } - boolean minimum = _multiModeReducerType == MultiModeReducerType.MIN; - if (_resultType == ColumnDataType.STRING && counts instanceof Object2LongMap) { - String mode = null; - long maxCount = 0; - ObjectIterator> iterator = - Object2LongMaps.fastIterator((Object2LongMap) counts); - while (iterator.hasNext()) { - Object2LongMap.Entry entry = iterator.next(); - String value = entry.getKey(); - long count = entry.getLongValue(); + Comparable mode = null; + long maxCount = 0; + if (counts != null) { + for (Map.Entry entry : counts.entrySet()) { + Comparable value = (Comparable) entry.getKey(); + long count = entry.getValue(); if (mode == null || count > maxCount || (count == maxCount - && (minimum ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { - mode = value; - maxCount = count; - } - } - return mode; - } - if (_resultType == ColumnDataType.TIMESTAMP && counts instanceof Long2LongMap) { - ObjectIterator iterator = Long2LongMaps.fastIterator((Long2LongMap) counts); - Long2LongMap.Entry first = iterator.next(); - long mode = first.getLongKey(); - long maxCount = first.getLongValue(); - while (iterator.hasNext()) { - Long2LongMap.Entry entry = iterator.next(); - long value = entry.getLongKey(); - long count = entry.getLongValue(); - if (count > maxCount || (count == maxCount && (minimum ? value < mode : value > mode))) { + && (_multiModeReducerType == MultiModeReducerType.MIN + ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { mode = value; maxCount = count; } } - return mode; - } - Comparable mode = null; - long maxCount = 0; - for (Map.Entry entry : counts.entrySet()) { - Comparable value = (Comparable) entry.getKey(); - long count = entry.getValue(); - if (mode == null || count > maxCount || (count == maxCount - && (minimum ? value.compareTo(mode) < 0 : value.compareTo(mode) > 0))) { - mode = value; - maxCount = count; - } } return mode; } - private static final class DictionaryCounts { - private final Dictionary _dictionary; - private final Int2LongOpenHashMap _counts = new Int2LongOpenHashMap(); - - private DictionaryCounts(Dictionary dictionary) { - _dictionary = dictionary; - } - } - - /// Frequency state with an O(1) conservative estimate of the retained string-key payload. - /// Accumulation uses [#addTo] and dictionary extraction and boxed [Map#merge] use [#put]. - /// Each distinct key is charged once, assuming UTF-16 storage plus object and array overhead. - /// Instances belong to one result holder and are not thread-safe. - public static final class StringModeCounts extends Object2LongOpenHashMap { - private long _retainedStringBytes; - - public StringModeCounts() { - } - - /// Restores accounting once when a generic map is deserialized from the existing wire format. - public StringModeCounts(Map counts) { - super(counts.size()); - counts.forEach((value, count) -> put(value, count.longValue())); - } - - public long getRetainedStringBytes() { - return _retainedStringBytes; - } - - @Override - public long addTo(String value, long increment) { - int previousSize = size(); - long previousCount = super.addTo(value, increment); - if (size() != previousSize) { - _retainedStringBytes += retainedStringBytes(value); - } - return previousCount; - } - - @Override - public long put(String value, long count) { - int previousSize = size(); - long previousCount = super.put(value, count); - if (size() != previousSize) { - _retainedStringBytes += retainedStringBytes(value); - } - return previousCount; - } - - @Override - public long removeLong(Object value) { - int previousSize = size(); - long previousCount = super.removeLong(value); - if (size() != previousSize) { - _retainedStringBytes -= retainedStringBytes((String) value); - } - return previousCount; - } - - @Override - public void clear() { - super.clear(); - _retainedStringBytes = 0; - } - - private static long retainedStringBytes(String value) { - return 48 + 2L * value.length(); - } - } - private enum MultiModeReducerType { MIN, MAX, AVG } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java index 051ac3202e30..da6aad85d38b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java @@ -28,8 +28,6 @@ import org.apache.pinot.common.request.Expression; import org.apache.pinot.common.request.Function; import org.apache.pinot.common.request.PinotQuery; -import org.apache.pinot.common.request.context.LiteralContext; -import org.apache.pinot.common.request.context.RequestContextUtils; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; @@ -100,8 +98,7 @@ private static ColumnDataType getOperandType(Expression operand, Schema schema) : null; } if (operand.isSetLiteral()) { - LiteralContext literal = RequestContextUtils.getExpression(operand).getLiteral(); - return ColumnDataType.fromDataType(literal.getType(), literal.isSingleValue()); + return RequestUtils.getLiteralTypeAndValue(operand.getLiteral()).getLeft(); } if (!operand.isSetFunctionCall()) { return null; @@ -168,25 +165,19 @@ private static ColumnDataType literalType(List arguments, int positi return null; } String type = arguments.get(position).getLiteral().getStringValue().toUpperCase(Locale.ROOT); - switch (type) { - case "VARCHAR": - case "CHAR": - case "JSON": - return ColumnDataType.STRING; - case "BIGINT": - return ColumnDataType.LONG; - case "INTEGER": - return ColumnDataType.INT; - case "REAL": - return ColumnDataType.FLOAT; - case "DECIMAL": - return ColumnDataType.BIG_DECIMAL; - default: + return switch (type) { + case "VARCHAR", "CHAR", "JSON" -> ColumnDataType.STRING; + case "BIGINT" -> ColumnDataType.LONG; + case "INTEGER" -> ColumnDataType.INT; + case "REAL" -> ColumnDataType.FLOAT; + case "DECIMAL" -> ColumnDataType.BIG_DECIMAL; + default -> { try { - return ColumnDataType.valueOf(type); + yield ColumnDataType.valueOf(type); } catch (IllegalArgumentException e) { - return null; + yield null; } - } + } + }; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java index 16adaf4bc3dc..6f815901a627 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java @@ -118,17 +118,6 @@ protected void replaceColumnNameWithAlias(DataSchema dataSchema) { queryContext = _queryContext.getSubquery(); } List aliasList = queryContext.getAliasList(); - String[] columnNames = dataSchema.getColumnNames(); - if (columnNames.length == aliasList.size()) { - // Reduced results follow SELECT order. Server rewrites can change expression names (for example, inferred - // MODE type arguments), so bind aliases by position as in BaseReduceService.updateAlias. - for (int i = 0; i < columnNames.length; i++) { - if (aliasList.get(i) != null) { - columnNames[i] = aliasList.get(i); - } - } - return; - } Map columnNameToAliasMap = new HashMap<>(); for (int i = 0; i < aliasList.size(); i++) { if (aliasList.get(i) != null) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java index a644e81a249d..0af38195b4df 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java @@ -179,10 +179,7 @@ private void gapfill(long bucketTime, List bucketedResult, List counts = min.deserializeIntermediateResult( + new CustomObject(serialized.getType(), ByteBuffer.wrap(serialized.getBytes()))); + counts = min.merge(counts, Map.of(larger, 1L)); + assertEquals(min.getFinalResultColumnType(), ColumnDataType.valueOf(type)); + assertEquals(min.extractFinalResult(counts), smaller); + assertEquals(typedMode(type, "MAX").extractFinalResult(counts), larger); + assertNull(min.extractFinalResult(null)); + expectThrows(IllegalArgumentException.class, () -> typedMode(type, "AVG")); } - @Test(dataProvider = "timestampScenarios") - void timestampModeAllNullAndEmptyInput(Scenario scenario) { - scenario.getDeclaringTable(true) - .onFirstInstance("myField", "null", "null") - .andOnSecondInstance("myField", "null") - .whenQuery("select mode(myField, 'MIN', 'TIMESTAMP') as mode from testTable") - .thenResultIs(new Object[]{null}) - .whenQuery("select mode(myField, 'MIN', 'TIMESTAMP') as mode " - + "from testTable where myField > '2026-09-03 00:00:00'") - .thenResultIs(new Object[]{null}); + private static ModeAggregationFunction typedMode(String type, String reducer) { + return new ModeAggregationFunction(List.of(ExpressionContext.forIdentifier("value"), + ExpressionContext.forLiteral(Literal.stringValue(reducer)), + ExpressionContext.forLiteral(Literal.stringValue(type))), true); } public class Scenario { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java deleted file mode 100644 index b4521996e045..000000000000 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeNonNumericAggregationFunctionTest.java +++ /dev/null @@ -1,262 +0,0 @@ -/** - * 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.core.query.aggregation.function; - -import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap; -import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap; -import java.nio.ByteBuffer; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import org.apache.pinot.common.CustomObject; -import org.apache.pinot.common.request.Literal; -import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.utils.DataSchema.ColumnDataType; -import org.apache.pinot.core.common.BlockValSet; -import org.apache.pinot.core.common.SyntheticBlockValSets; -import org.apache.pinot.core.query.aggregation.AggregationResultHolder; -import org.apache.pinot.core.query.aggregation.function.ModeAggregationFunction.StringModeCounts; -import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; -import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.roaringbitmap.RoaringBitmap; -import org.testng.annotations.Test; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNull; -import static org.testng.Assert.assertSame; -import static org.testng.Assert.assertTrue; -import static org.testng.Assert.expectThrows; - - -/// Verifies serialization, exact timestamp values and null handling for the typed MODE calls. -public class ModeNonNumericAggregationFunctionTest { - private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("value"); - - @Test - public void testStringIntermediateResultsRoundTripAndMerge() { - ModeAggregationFunction function = new ModeAggregationFunction(typedArguments("MIN", "STRING"), true); - var first = aggregateAndRoundTrip(function, - SyntheticBlockValSets.Str.create(null, new String[]{"é", "é", "é", "苹果", "苹果", ""}), 6); - var second = aggregateAndRoundTrip(function, - SyntheticBlockValSets.Str.create(null, new String[]{"zebra", "zebra", "zebra", "苹果", "苹果", "\0"}), 6); - - assertEquals(function.getFinalResultColumnType(), ColumnDataType.STRING); - assertEquals(function.extractFinalResult(first), "é"); - assertEquals(function.extractFinalResult(second), "zebra"); - assertEquals(function.extractFinalResult(function.merge(first, second)), "苹果"); - } - - @Test - public void testTimestampIntermediateResultsPreserveLongPrecision() { - // Adjacent long values above 2^53 become identical if converted through double. - long earlier = 9_007_199_254_740_992L; - long later = earlier + 1; - ModeAggregationFunction minFunction = - new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), true); - ModeAggregationFunction maxFunction = - new ModeAggregationFunction(typedArguments("MAX", "TIMESTAMP"), true); - var first = aggregateAndRoundTrip(minFunction, timestampValues(earlier, later, later), 3); - var second = aggregateAndRoundTrip(minFunction, timestampValues(earlier), 1); - - assertEquals(minFunction.getFinalResultColumnType(), ColumnDataType.TIMESTAMP); - assertEquals(minFunction.extractFinalResult(first), Long.valueOf(later)); - var merged = minFunction.merge(first, second); - assertEquals(minFunction.extractFinalResult(merged), Long.valueOf(earlier)); - assertEquals(maxFunction.extractFinalResult(merged), Long.valueOf(later)); - } - - @Test - public void testTimestampMergesPrimitiveAndGenericStates() { - ModeAggregationFunction minFunction = - new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), true); - ModeAggregationFunction maxFunction = - new ModeAggregationFunction(typedArguments("MAX", "TIMESTAMP"), true); - for (boolean primitiveLeft : new boolean[]{false, true}) { - for (boolean primitiveRight : new boolean[]{false, true}) { - Map left = primitiveLeft ? new Long2LongOpenHashMap() : new HashMap<>(); - Map right = primitiveRight ? new Long2LongOpenHashMap() : new HashMap<>(); - left.put(Long.MIN_VALUE, 2L); - left.put(0L, 1L); - right.put(Long.MIN_VALUE, 1L); - right.put(Long.MAX_VALUE, 3L); - - Map merged = minFunction.merge(left, right); - assertSame(merged, left); - assertEquals(merged, Map.of(Long.MIN_VALUE, 3L, 0L, 1L, Long.MAX_VALUE, 3L)); - assertEquals(minFunction.extractFinalResult(merged), Long.valueOf(Long.MIN_VALUE)); - assertEquals(maxFunction.extractFinalResult(merged), Long.valueOf(Long.MAX_VALUE)); - assertEquals(right, Map.of(Long.MIN_VALUE, 1L, Long.MAX_VALUE, 3L)); - } - } - } - - @Test - public void testStringMergesPrimitiveAndGenericStates() { - ModeAggregationFunction minFunction = new ModeAggregationFunction(typedArguments("MIN", "STRING"), true); - ModeAggregationFunction maxFunction = new ModeAggregationFunction(typedArguments("MAX", "STRING"), true); - for (boolean primitiveLeft : new boolean[]{false, true}) { - for (boolean primitiveRight : new boolean[]{false, true}) { - Map left = primitiveLeft ? new StringModeCounts() : new HashMap<>(); - Map right = primitiveRight ? new Object2LongOpenHashMap<>() : new HashMap<>(); - left.put("alpha", 2L); - right.put("alpha", 1L); - right.put("zebra", 3L); - - Map merged = minFunction.merge(left, right); - assertSame(merged, left); - assertEquals(merged, Map.of("alpha", 3L, "zebra", 3L)); - assertEquals(minFunction.extractFinalResult(merged), "alpha"); - assertEquals(maxFunction.extractFinalResult(merged), "zebra"); - assertEquals(right, Map.of("alpha", 1L, "zebra", 3L)); - if (primitiveLeft) { - assertEquals(((StringModeCounts) left).getRetainedStringBytes(), 2L * (48 + 2 * 5)); - } - } - } - } - - @Test - public void testStringModeSkipsNullRowsForMultiValueGroupKeys() { - ModeAggregationFunction function = new ModeAggregationFunction(typedArguments("MIN", "STRING"), true); - GroupByResultHolder holder = function.createGroupByResultHolder(3, 3); - BlockValSet values = SyntheticBlockValSets.Str.create(RoaringBitmap.bitmapOf(0, 2), - new String[]{"ignored", "alpha", "ignored", "beta", "alpha"}); - function.aggregateGroupByMV(5, new int[][]{{0, 1}, {0}, {1, 2}, {0, 1}, {0}}, holder, - Map.of(EXPRESSION, values)); - - assertEquals(function.extractFinalResult(function.extractGroupByResult(holder, 0)), "alpha"); - assertEquals(function.extractFinalResult(function.extractGroupByResult(holder, 1)), "beta"); - assertNull(function.extractFinalResult(function.extractGroupByResult(holder, 2))); - } - - @Test - public void testEmptyResultsAreNullWithEitherNullHandlingMode() { - for (boolean nullHandlingEnabled : new boolean[]{false, true}) { - ModeAggregationFunction stringFunction = - new ModeAggregationFunction(typedArguments("MIN", "STRING"), nullHandlingEnabled); - ModeAggregationFunction timestampFunction = - new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), nullHandlingEnabled); - assertNull(stringFunction.extractFinalResult(null)); - assertNull(timestampFunction.extractFinalResult(null)); - assertNull(timestampFunction.extractFinalResult(new Long2LongOpenHashMap())); - assertNull(timestampFunction.extractFinalResult(Map.of())); - assertNull(stringFunction.extractFinalResult( - stringFunction.extractAggregationResult(stringFunction.createAggregationResultHolder()))); - assertNull(timestampFunction.extractFinalResult( - timestampFunction.extractAggregationResult(timestampFunction.createAggregationResultHolder()))); - } - } - - @Test - public void testNonNumericModesRejectAverageReducer() { - IllegalArgumentException stringError = expectThrows(IllegalArgumentException.class, - () -> new ModeAggregationFunction(typedArguments("AVG", "STRING"), true)); - assertTrue(stringError.getMessage().contains("AVG")); - assertTrue(stringError.getMessage().contains("STRING")); - IllegalArgumentException timestampError = expectThrows(IllegalArgumentException.class, - () -> new ModeAggregationFunction(typedArguments("AVG", "TIMESTAMP"), true)); - assertTrue(timestampError.getMessage().contains("AVG")); - assertTrue(timestampError.getMessage().contains("TIMESTAMP")); - } - - @Test - public void testResultTypeMustBeSupportedStringLiteral() { - ExpressionContext reducer = ExpressionContext.forLiteral(Literal.stringValue("MIN")); - for (ExpressionContext type : List.of(EXPRESSION, ExpressionContext.forLiteral(Literal.intValue(1)), - ExpressionContext.forLiteral(DataType.STRING, null), - ExpressionContext.forLiteral(Literal.stringValue("LONG")), - ExpressionContext.forLiteral(Literal.stringValue("DOUBLE")), - ExpressionContext.forLiteral(Literal.stringValue("INVALID")))) { - IllegalArgumentException error = expectThrows(IllegalArgumentException.class, - () -> new ModeAggregationFunction(List.of(EXPRESSION, reducer, type), true)); - assertTrue(error.getMessage().contains("MODE result type")); - } - expectThrows(IllegalArgumentException.class, () -> new ModeAggregationFunction(List.of(), true)); - expectThrows(IllegalArgumentException.class, - () -> new ModeAggregationFunction(List.of(EXPRESSION, reducer, reducer, reducer), true)); - expectThrows(IllegalArgumentException.class, - () -> new ModeAggregationFunction(List.of(EXPRESSION, EXPRESSION, - ExpressionContext.forLiteral(Literal.stringValue("STRING"))), true)); - } - - @Test - public void testResultTypeLiteralIsCaseInsensitive() { - for (String type : List.of("string", "StRiNg", "timestamp", "TiMeStAmP")) { - ModeAggregationFunction function = new ModeAggregationFunction(typedArguments("MIN", type), true); - assertEquals(function.getFinalResultColumnType(), - type.equalsIgnoreCase("string") ? ColumnDataType.STRING : ColumnDataType.TIMESTAMP); - assertEquals(function.getResultColumnName(), "mode(value,'MIN','" + type + "')"); - } - } - - @Test - public void testLegacyNumericResultTypeAndReducersAreUnchanged() { - // The same stored longs are DOUBLE for legacy calls and exact TIMESTAMP values only with the inferred type. - long earlier = 9_007_199_254_740_992L; - long later = earlier + 1; - ModeAggregationFunction numeric = new ModeAggregationFunction(List.of(EXPRESSION), false); - ModeAggregationFunction timestamp = new ModeAggregationFunction(typedArguments("MIN", "TIMESTAMP"), false); - assertEquals(numeric.getFinalResultColumnType(), ColumnDataType.DOUBLE); - assertEquals(timestamp.getFinalResultColumnType(), ColumnDataType.TIMESTAMP); - Map counts = aggregateAndRoundTrip(timestamp, timestampValues(later, later, earlier), 3); - assertEquals(numeric.extractFinalResult(counts), Double.valueOf(later)); - assertEquals(timestamp.extractFinalResult(counts), Long.valueOf(later)); - assertEquals(numeric.getFinalResultColumnType(), ColumnDataType.DOUBLE); - assertEquals(timestamp.getFinalResultColumnType(), ColumnDataType.TIMESTAMP); - - for (String reducer : List.of("MIN", "MAX", "AVG")) { - ModeAggregationFunction function = new ModeAggregationFunction(List.of(EXPRESSION, - ExpressionContext.forLiteral(Literal.stringValue(reducer))), false); - assertEquals(function.getResultColumnName(), "mode(value)"); - Map tiedCounts = aggregateAndRoundTrip(function, timestampValues(2, 4), 2); - double expected = reducer.equals("MIN") ? 2D : reducer.equals("MAX") ? 4D : 3D; - assertEquals(function.extractFinalResult(tiedCounts), expected); - assertEquals(function.extractFinalResult(new Long2LongOpenHashMap()), Double.NEGATIVE_INFINITY); - assertNull(function.extractFinalResult(null)); - } - } - - private static List typedArguments(String reducer, String type) { - return List.of(EXPRESSION, ExpressionContext.forLiteral(Literal.stringValue(reducer)), - ExpressionContext.forLiteral(Literal.stringValue(type))); - } - - private static BlockValSet timestampValues(long... values) { - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.isSingleValue()).thenReturn(true); - when(blockValSet.getValueType()).thenReturn(DataType.TIMESTAMP); - when(blockValSet.getLongValuesSV()).thenReturn(values); - return blockValSet; - } - - private static Map aggregateAndRoundTrip(ModeAggregationFunction function, - BlockValSet values, int length) { - AggregationResultHolder holder = function.createAggregationResultHolder(); - function.aggregate(length, holder, Map.of(EXPRESSION, values)); - Map intermediateResult = function.extractAggregationResult(holder); - AggregationFunction.SerializedIntermediateResult serialized = - function.serializeIntermediateResult(intermediateResult); - Map deserialized = function.deserializeIntermediateResult( - new CustomObject(serialized.getType(), ByteBuffer.wrap(serialized.getBytes()))); - assertEquals(deserialized, intermediateResult); - return deserialized; - } -} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java index 04dd2d981d05..552f7ba6ff4f 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java @@ -29,15 +29,12 @@ import static org.testng.Assert.assertEquals; -/// Verifies MODE rewriting through the single-stage optimizer, including expressions and post-aggregation clauses. +/// Covers inferred MODE arguments and the independent rollout opt-in. public class ModeAggregationFunctionRewriteOptimizerTest { private static final QueryOptimizer OPTIMIZER = new QueryOptimizer(); private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName("testTable") .addSingleValueDimension("stringCol", DataType.STRING) - .addSingleValueDimension("intCol", DataType.INT) .addSingleValueDimension("longCol", DataType.LONG) - .addSingleValueDimension("floatCol", DataType.FLOAT) - .addSingleValueDimension("doubleCol", DataType.DOUBLE) .addMultiValueDimension("mvStringCol", DataType.STRING) .addDateTime("timestampCol", DataType.TIMESTAMP, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") .build(); @@ -48,88 +45,35 @@ public Object[][] modeExpressions() { {"MODE(stringCol)", "MODE(stringCol, 'MIN', 'STRING')"}, {"MODE(timestampCol, 'MAX')", "MODE(timestampCol, 'MAX', 'TIMESTAMP')"}, {"MODE(CONCAT(stringCol, 'suffix'))", "MODE(CONCAT(stringCol, 'suffix'), 'MIN', 'STRING')"}, - {"MODE(JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', ''))", - "MODE(JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', ''), 'MIN', 'STRING')"}, - {"MODE(CAST(intCol AS STRING))", "MODE(CAST(intCol AS STRING), 'MIN', 'STRING')"}, + {"MODE(CASE WHEN stringCol = '' THEN NULL ELSE JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', '') END)", + "MODE(CASE WHEN stringCol = '' THEN NULL ELSE JSONEXTRACTSCALAR(stringCol, '$.user', 'STRING', '') END, " + + "'MIN', 'STRING')"}, {"MODE(CAST(stringCol AS TIMESTAMP))", "MODE(CAST(stringCol AS TIMESTAMP), 'MIN', 'TIMESTAMP')"}, - {"MODE(CASE WHEN intCol > 0 THEN stringCol ELSE 'other' END)", - "MODE(CASE WHEN intCol > 0 THEN stringCol ELSE 'other' END, 'MIN', 'STRING')"}, - {"MODE('literal')", "MODE('literal', 'MIN', 'STRING')"}, {"fromTimestamp(MODE(timestampCol))", "fromTimestamp(MODE(timestampCol, 'MIN', 'TIMESTAMP'))"} }; } @Test(dataProvider = "modeExpressions") public void testModeExpressions(String original, String rewritten) { - for (boolean nullHandlingEnabled : new boolean[]{false, true}) { - String prefix = - "SET enableTypedMode=true; SET enableNullHandling=" + nullHandlingEnabled + "; SELECT "; - TestHelper.assertEqualsQuery(prefix + original + " AS commonValue FROM testTable", - prefix + rewritten + " AS commonValue FROM testTable", SCHEMA); - } - } - - @Test - public void testModeInHavingAndOrderBy() { - TestHelper.assertEqualsQuery( - "SET enableTypedMode=true; " - + "SELECT intCol, MODE(stringCol) AS commonValue FROM testTable GROUP BY intCol " - + "HAVING MODE(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END) = 'value' " - + "ORDER BY MODE(timestampCol) DESC", - "SET enableTypedMode=true; " - + "SELECT intCol, MODE(stringCol, 'MIN', 'STRING') AS commonValue FROM testTable GROUP BY intCol " - + "HAVING MODE(CASE WHEN stringCol = '' THEN NULL ELSE stringCol END, 'MIN', 'STRING') = 'value' " - + "ORDER BY MODE(timestampCol, 'MIN', 'TIMESTAMP') DESC", SCHEMA); - } - - @Test - public void testNumericModeAndOptInRewritesRemainUnchanged() { - assertUnchanged("SET enableTypedMode=true; SELECT MODE(intCol), MODE(longCol), MODE(floatCol), MODE(doubleCol), " - + "MODE(CAST(stringCol AS LONG)), MODE(fromDateTime(stringCol, 'yyyy-MM-dd HH:mm:ss')), " - + "MIN(stringCol), MAX(longCol), SUM(intCol) FROM testTable", SCHEMA); - assertUnchanged("SET enableTypedMode=false; SELECT MODE(stringCol), MODE(timestampCol) FROM testTable", - SCHEMA); - assertUnchanged("SELECT MODE(stringCol), MODE(timestampCol) FROM testTable", SCHEMA); + String prefix = "SET enableTypedMode=true; SELECT "; + TestHelper.assertEqualsQuery(prefix + original + " FROM testTable", prefix + rewritten + " FROM testTable", SCHEMA); } @Test - public void testExistingRewriteOptionPreservesLegacyMode() { - assertUnchanged("SET autoRewriteAggregationType=true; " - + "SELECT MODE(timestampCol), MODE(timestampCol, 'AVG'), MODE(stringCol) FROM testTable", SCHEMA); - assertUnchanged("SET autoRewriteAggregationType=true; SET enableTypedMode=false; " - + "SELECT MODE(timestampCol), MODE(timestampCol, 'AVG'), MODE(stringCol) FROM testTable", SCHEMA); - TestHelper.assertEqualsQuery("SET autoRewriteAggregationType=true; SET enableTypedMode=true; " - + "SELECT MODE(timestampCol), MODE(stringCol), MIN(stringCol) FROM testTable", - "SET autoRewriteAggregationType=true; SET enableTypedMode=true; " - + "SELECT MODE(timestampCol, 'MIN', 'TIMESTAMP'), MODE(stringCol, 'MIN', 'STRING'), " - + "MINSTRING(stringCol) FROM testTable", SCHEMA); - } - - @Test - public void testExplicitTypesAndRepeatedOptimization() { - assertUnchanged("SET enableTypedMode=true; SELECT MODE(stringCol, 'MAX', 'STRING'), " - + "MODE(timestampCol, 'MIN', 'TIMESTAMP') FROM testTable", SCHEMA); - PinotQuery query = CalciteSqlParser.compileToPinotQuery( - "SET enableTypedMode=true; SELECT MODE(stringCol), MODE(timestampCol, 'MAX') FROM testTable"); - OPTIMIZER.optimize(query, SCHEMA); - PinotQuery once = query.deepCopy(); - OPTIMIZER.optimize(query, SCHEMA); - assertEquals(query, once); - } - - @Test - public void testServerDependentModeDoesNotInitializeOnBroker() { - assertUnchanged("SET enableTypedMode=true; " - + "SELECT MODE(LOOKUP('baseballTeams', 'teamInteger', 'teamID', stringCol)) FROM testTable", SCHEMA); + public void testLegacyCallsAndOtherRewritesAreUnchanged() { + for (String options : new String[]{"", "SET enableTypedMode=false; ", "SET autoRewriteAggregationType=true; "}) { + assertUnchanged(options + "SELECT MODE(stringCol), MODE(timestampCol), MODE(longCol, 'AVG') FROM testTable", + SCHEMA); + } + assertUnchanged("SET enableTypedMode=true; SELECT MODE(longCol, 'AVG'), MIN(stringCol), SUM(longCol), " + + "MODE(stringCol, 'MAX', 'STRING'), MODE(timestampCol, 'MIN', 'TIMESTAMP') FROM testTable", SCHEMA); } @Test - public void testMissingSchemaAndColumns() { + public void testUnknownTypesDoNotInitializeServerTransforms() { + assertUnchanged("SET enableTypedMode=true; SELECT MODE(unknownCol), MODE(mvStringCol), " + + "MODE(LOOKUP('baseballTeams', 'teamInteger', 'teamID', stringCol)) FROM testTable", SCHEMA); assertUnchanged("SET enableTypedMode=true; SELECT MODE(stringCol) FROM testTable", null); - assertUnchanged("SET enableTypedMode=true; SELECT MODE(unknownCol) FROM testTable", SCHEMA); - assertUnchanged("SET enableTypedMode=true; SELECT MODE(CONCAT(unknownCol, 'suffix')) FROM testTable", - SCHEMA); - assertUnchanged("SET enableTypedMode=true; SELECT MODE(mvStringCol) FROM testTable", SCHEMA); } private static void assertUnchanged(String sql, Schema schema) { diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java index 8d3e8349f78c..3767e4c07e07 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java @@ -195,22 +195,13 @@ private BrokerResponseNative getBrokerResponse(@Language("sql") String query, Pl /// This can be particularly useful to test statistical aggregation functions. /// @see StatisticalQueriesTest for an example use case. private BrokerResponseNative getBrokerResponse(PinotQuery pinotQuery, PlanMaker planMaker) { - return getBrokerResponse(pinotQuery, planMaker, false, null); - } - - private BrokerResponseNative getBrokerResponse(PinotQuery pinotQuery, PlanMaker planMaker, boolean optimize, - @Nullable Schema schema) { - // Match the broker's order: retain the original gapfill query while optimizing the stripped server query. - PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); - if (optimize) { - OPTIMIZER.optimize(serverPinotQuery, schema); - } List> instances = getDistinctInstances(); if (instances.size() == 2) { - return getBrokerResponseDistinctInstances(pinotQuery, serverPinotQuery, planMaker); + return getBrokerResponseDistinctInstances(pinotQuery, planMaker); } // Server side + PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); QueryContext serverQueryContext = serverPinotQuery == pinotQuery ? queryContext : QueryContextConverterUtils.getQueryContext(serverPinotQuery); @@ -275,7 +266,8 @@ protected BrokerResponseNative reduceOnDataTable(BrokerRequest brokerRequest, Br protected BrokerResponseNative getBrokerResponseForOptimizedQuery(@Language("sql") String query, @Nullable Schema schema) { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - return getBrokerResponse(pinotQuery, PLAN_MAKER, true, schema); + OPTIMIZER.optimize(pinotQuery, schema); + return getBrokerResponse(pinotQuery, PLAN_MAKER); } /// Run query on multiple index segments with custom plan maker. @@ -288,9 +280,9 @@ protected BrokerResponseNative getBrokerResponseForOptimizedQuery(@Language("sql /// overriding getDistinctInstances. /// This can be particularly useful to test statistical aggregation functions. /// @see StatisticalQueriesTest for an example use case. - private BrokerResponseNative getBrokerResponseDistinctInstances(PinotQuery pinotQuery, PinotQuery serverPinotQuery, - PlanMaker planMaker) { + private BrokerResponseNative getBrokerResponseDistinctInstances(PinotQuery pinotQuery, PlanMaker planMaker) { // Server side + PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); QueryContext serverQueryContext = serverPinotQuery == pinotQuery ? queryContext : QueryContextConverterUtils.getQueryContext(serverPinotQuery); diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java index 79cc27334ea0..82a96a5359e2 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ExprMinMaxTest.java @@ -38,7 +38,6 @@ import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; -import org.apache.pinot.spi.data.DateTimeFormatSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; @@ -49,7 +48,6 @@ import org.testng.Assert; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; -import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import static org.apache.pinot.spi.utils.CommonConstants.RewriterConstants.CHILD_AGGREGATION_NAME_PREFIX; @@ -614,44 +612,6 @@ public void testEmptyGroupByInterSegment() { assertEquals(rows.size(), 0); } - @DataProvider - public Object[][] gapfillParentResultTypes() { - return new Object[][]{ - // Parent aggregates expose timestamp children through their stored LONG type. - {TIMESTAMP_COLUMN, 1683138373878L}, - {BYTES_COLUMN, "31"}, - {BIG_DECIMAL_COLUMN, "1199"} - }; - } - - @Test(dataProvider = "gapfillParentResultTypes") - public void testGapfillFormatsNativeParentResults(String column, Object expectedValue) { - String format = "1:MILLISECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss.SSS"; - DateTimeFormatSpec formatter = new DateTimeFormatSpec(format); - long hourMillis = 3_600_000L; - long startMillis = 1683138373878L - 1683138373878L % hourMillis; - String start = formatter.fromMillisToFormat(startMillis); - String end = formatter.fromMillisToFormat(startMillis + 2 * hourMillis); - // Parent aggregate projections are referenced by their canonical names because AS is not supported. - String expression = "exprmin(" + column + ",intColumn)"; - String projection = "\"" + expression + "\""; - String query = "SELECT GapFill(time_col, '" + format + "', '" + start + "', '" + end - + "', '1:HOURS', FILL(" + projection + ", 'FILL_PREVIOUS_VALUE'), TIMESERIESON(groupByIntColumn)), " - + "groupByIntColumn, " + projection + " FROM (SELECT DATETIMECONVERT(fromTimestamp(timestampColumn), " - + "'1:MILLISECONDS:EPOCH', '" + format + "', '1:HOURS') AS time_col, groupByIntColumn, " - + expression + " FROM testTable WHERE intColumn = 1 GROUP BY time_col, groupByIntColumn LIMIT 1000) LIMIT 1000"; - - BrokerResponseNative response = getBrokerResponse(query); - assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); - List rows = response.getResultTable().getRows(); - // The shared fixture returns the parent projection once per server, followed by one filled row. - assertEquals(rows.size(), 3); - assertEquals(rows.get(0), new Object[]{start, 1, expectedValue}); - assertEquals(rows.get(1), new Object[]{start, 1, expectedValue}); - assertEquals(rows.get(2), - new Object[]{formatter.fromMillisToFormat(startMillis + hourMillis), 1, expectedValue}); - } - @Test public void testAlias() { // Using exprmin/exprmax with alias will fail, since the alias will not be resolved by the rewriter diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java index a7ecf3b3ccd9..b79c9c1fbce1 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/ModeQueriesTest.java @@ -21,7 +21,6 @@ import com.google.common.collect.Lists; import java.io.File; import java.io.IOException; -import java.sql.Timestamp; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -32,7 +31,6 @@ import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; import org.apache.pinot.common.response.broker.BrokerResponseNative; -import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock; import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; @@ -48,7 +46,6 @@ import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; -import org.apache.pinot.spi.data.DateTimeFormatSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; @@ -83,28 +80,16 @@ public class ModeQueriesTest extends BaseQueriesTest { private static final String LONG_NO_DICT_COLUMN = "longNoDictColumn"; private static final String FLOAT_NO_DICT_COLUMN = "floatNoDictColumn"; private static final String DOUBLE_NO_DICT_COLUMN = "doubleNoDictColumn"; - private static final String STRING_COLUMN = "stringColumn"; - private static final String STRING_NO_DICT_COLUMN = "stringNoDictColumn"; - private static final String JSON_COLUMN = "jsonColumn"; - private static final String TIMESTAMP_COLUMN = "timestampColumn"; - private static final String TIMESTAMP_NO_DICT_COLUMN = "timestampNoDictColumn"; - private static final long BASE_TIMESTAMP = Timestamp.valueOf("2026-09-03 10:11:12.000").getTime(); private static final Schema SCHEMA = new Schema.SchemaBuilder().addSingleValueDimension(INT_COLUMN, DataType.INT) .addMultiValueDimension(INT_MV_COLUMN, DataType.INT).addSingleValueDimension(INT_NO_DICT_COLUMN, DataType.INT) .addSingleValueDimension(LONG_COLUMN, DataType.LONG).addSingleValueDimension(LONG_NO_DICT_COLUMN, DataType.LONG) .addSingleValueDimension(FLOAT_COLUMN, DataType.FLOAT) .addSingleValueDimension(FLOAT_NO_DICT_COLUMN, DataType.FLOAT) .addSingleValueDimension(DOUBLE_COLUMN, DataType.DOUBLE) - .addSingleValueDimension(DOUBLE_NO_DICT_COLUMN, DataType.DOUBLE) - .addSingleValueDimension(STRING_COLUMN, DataType.STRING) - .addSingleValueDimension(STRING_NO_DICT_COLUMN, DataType.STRING) - .addSingleValueDimension(JSON_COLUMN, DataType.STRING) - .addSingleValueDimension(TIMESTAMP_COLUMN, DataType.TIMESTAMP) - .addSingleValueDimension(TIMESTAMP_NO_DICT_COLUMN, DataType.TIMESTAMP).build(); + .addSingleValueDimension(DOUBLE_NO_DICT_COLUMN, DataType.DOUBLE).build(); private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) .setNoDictionaryColumns( - Lists.newArrayList(INT_NO_DICT_COLUMN, LONG_NO_DICT_COLUMN, FLOAT_NO_DICT_COLUMN, DOUBLE_NO_DICT_COLUMN, - STRING_NO_DICT_COLUMN, TIMESTAMP_NO_DICT_COLUMN)) + Lists.newArrayList(INT_NO_DICT_COLUMN, LONG_NO_DICT_COLUMN, FLOAT_NO_DICT_COLUMN, DOUBLE_NO_DICT_COLUMN)) .build(); private static final double DELTA = 0.00001; @@ -151,11 +136,6 @@ public void setUp() record.putValue(FLOAT_NO_DICT_COLUMN, (float) value); record.putValue(DOUBLE_COLUMN, (double) value); record.putValue(DOUBLE_NO_DICT_COLUMN, (double) value); - record.putValue(STRING_COLUMN, Integer.toString(value)); - record.putValue(STRING_NO_DICT_COLUMN, Integer.toString(value)); - record.putValue(JSON_COLUMN, "{\"value\":\"" + value + "\"}"); - record.putValue(TIMESTAMP_COLUMN, BASE_TIMESTAMP + value); - record.putValue(TIMESTAMP_NO_DICT_COLUMN, BASE_TIMESTAMP + value); records.add(record); } long maxOccurrences = _values.values().stream().max(Long::compareTo).get(); @@ -391,122 +371,6 @@ public Object[][] testAggregationGroupByMVDataProvider() { return entries.toArray(new Object[0][]); } - @Test - public void testStringAggregationAndComputedExpression() { - long maxOccurrences = _values.values().stream().max(Long::compareTo).orElseThrow(); - String expectedMin = _values.entrySet().stream().filter(e -> e.getValue() == maxOccurrences) - .map(e -> e.getKey().toString()).min(String::compareTo).orElseThrow(); - String expectedMax = _values.entrySet().stream().filter(e -> e.getValue() == maxOccurrences) - .map(e -> e.getKey().toString()).max(String::compareTo).orElseThrow(); - BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET enableTypedMode=true; SELECT MODE(stringColumn), " - + "MODE(stringNoDictColumn), MODE(stringColumn, 'MAX'), MODE(CONCAT('value-', stringColumn, '')), " - + "MODE(CASE WHEN JSONEXTRACTSCALAR(jsonColumn, '$.value', 'STRING', '') = '' THEN NULL " - + "ELSE JSONEXTRACTSCALAR(jsonColumn, '$.value', 'STRING', '') END) FROM testTable", - SCHEMA); - assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); - assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, - ColumnDataType.STRING}); - assertEquals(response.getResultTable().getRows().size(), 1); - assertEquals(response.getResultTable().getRows().get(0), - new Object[]{expectedMin, expectedMin, expectedMax, "value-" + expectedMin, expectedMin}); - } - - @Test - public void testStringAggregationWithNoMatchingRows() { - BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET enableTypedMode=true; SELECT MODE(stringColumn), " - + "MODE(stringNoDictColumn) FROM testTable WHERE intColumn < 0", SCHEMA); - assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); - assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING}); - assertEquals(response.getResultTable().getRows().size(), 1); - assertEquals(response.getResultTable().getRows().get(0), new Object[]{null, null}); - } - - @DataProvider - public Object[][] stringGroupByColumns() { - return new Object[][]{{INT_COLUMN}, {INT_MV_COLUMN}}; - } - - @Test(dataProvider = "stringGroupByColumns") - public void testStringAggregationGroupBy(String groupByColumn) { - BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET enableTypedMode=true; SELECT " + groupByColumn - + ", MODE(stringColumn), MODE(stringNoDictColumn), MODE(CONCAT('value-', stringColumn, '')) " - + "FROM testTable GROUP BY " + groupByColumn + " ORDER BY " + groupByColumn, SCHEMA); - assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); - assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING}); - List rows = response.getResultTable().getRows(); - assertEquals(rows.size(), 10); - for (Object[] row : rows) { - String value = row[0].toString(); - assertEquals(row, new Object[]{row[0], value, value, "value-" + value}); - } - } - - @Test - public void testTimestampAggregationAndResultType() { - BrokerResponseNative response = getBrokerResponseForOptimizedQuery( - "SET enableTypedMode=true; SELECT MODE(timestampColumn), " - + "MODE(timestampNoDictColumn), fromTimestamp(MODE(timestampColumn)), MODE(toTimestamp(longColumn)) " - + "FROM testTable", SCHEMA); - assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); - assertEquals(response.getResultTable().getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.TIMESTAMP, ColumnDataType.TIMESTAMP, ColumnDataType.LONG, - ColumnDataType.TIMESTAMP}); - long expectedMillis = BASE_TIMESTAMP + _expectedResultMin.longValue(); - String expectedTimestamp = new Timestamp(expectedMillis).toString(); - assertEquals(response.getResultTable().getRows().size(), 1); - assertEquals(response.getResultTable().getRows().get(0), - new Object[]{expectedTimestamp, expectedTimestamp, expectedMillis, - new Timestamp(_expectedResultMin.longValue()).toString()}); - } - - @DataProvider - public Object[][] typedModeGapfillExpressions() { - return new Object[][]{ - {false, "MODE(timestampColumn, 'MIN', 'TIMESTAMP')"}, - {false, "MODE(timestampColumn, 'MIN', 'timestamp')"}, - {true, "MODE(timestampColumn)"} - }; - } - - @Test(dataProvider = "typedModeGapfillExpressions") - public void testTypedModeGapfillPreservesAliases(boolean inferType, String typedMode) { - String format = "1:MILLISECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss.SSS"; - DateTimeFormatSpec formatter = new DateTimeFormatSpec(format); - long hourMillis = 3_600_000L; - long startMillis = BASE_TIMESTAMP - BASE_TIMESTAMP % hourMillis; - String start = formatter.fromMillisToFormat(startMillis); - String end = formatter.fromMillisToFormat(startMillis + 2 * hourMillis); - String entity = Integer.toString(_expectedResultMin.intValue()); - // Keep both the legacy DOUBLE and typed TIMESTAMP result in the same aggregate subquery. - String numericMode = inferType ? "MODE(fromTimestamp(timestampColumn))" : "MODE(timestampColumn)"; - String query = "SET enableTypedMode=" + inferType + "; SELECT GapFill(time_col, '" + format + "', '" - + start + "', '" + end + "', '1:HOURS', FILL(numeric_mode, 'FILL_PREVIOUS_VALUE'), " - + "FILL(typed_mode, 'FILL_PREVIOUS_VALUE'), TIMESERIESON(stringColumn)), " - + "stringColumn, numeric_mode, typed_mode FROM (SELECT DATETIMECONVERT(fromTimestamp(timestampColumn), " - + "'1:MILLISECONDS:EPOCH', '" + format + "', '1:HOURS') AS time_col, stringColumn, " - + numericMode + " AS numeric_mode, " + typedMode + " AS typed_mode FROM testTable WHERE stringColumn = '" - + entity + "' GROUP BY time_col, stringColumn LIMIT 1000) LIMIT 1000"; - BrokerResponseNative response = getBrokerResponseForOptimizedQuery(query, SCHEMA); - - assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); - assertEquals(response.getResultTable().getDataSchema().getColumnNames(), - new String[]{"time_col", "stringColumn", "numeric_mode", "typed_mode"}); - List rows = response.getResultTable().getRows(); - assertEquals(rows.size(), 2); - long expectedMillis = BASE_TIMESTAMP + _expectedResultMin.longValue(); - String expectedTimestamp = new Timestamp(expectedMillis).toString(); - for (int i = 0; i < rows.size(); i++) { - assertEquals(rows.get(i), new Object[]{formatter.fromMillisToFormat(startMillis + i * hourMillis), entity, - (double) expectedMillis, expectedTimestamp}); - } - } - @AfterClass public void tearDown() throws IOException { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java index 8d0a83597df3..d3e3746cb0c4 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java @@ -62,7 +62,7 @@ public void onMatch(RelOptRuleCall call) { } RexBuilder rexBuilder = input.getCluster().getRexBuilder(); List originalCalls = aggregate.getAggCallList(); - List> rewrittenArguments = new ArrayList<>(originalCalls.size()); + List rewrittenCalls = new ArrayList<>(originalCalls.size()); boolean changed = false; for (AggregateCall originalCall : originalCalls) { List arguments = originalCall.getArgList(); @@ -81,7 +81,7 @@ public void onMatch(RelOptRuleCall call) { changed = true; } } - rewrittenArguments.add(rewritten); + rewrittenCalls.add(originalCall.withArgList(rewritten)); } if (!changed) { return; @@ -99,14 +99,6 @@ public void onMatch(RelOptRuleCall call) { } else { rewrittenInput = LogicalProject.create(input, List.of(), projects, names); } - List rewrittenCalls = new ArrayList<>(originalCalls.size()); - for (int i = 0; i < originalCalls.size(); i++) { - AggregateCall original = originalCalls.get(i); - rewrittenCalls.add(AggregateCall.create(original.getAggregation(), original.isDistinct(), - original.isApproximate(), original.ignoreNulls(), rewrittenArguments.get(i), original.filterArg, - original.distinctKeys, original.getCollation(), aggregate.getGroupCount(), rewrittenInput, original.getType(), - original.getName())); - } call.transformTo(aggregate.copy(aggregate.getTraitSet(), rewrittenInput, aggregate.getGroupSet(), aggregate.getGroupSets(), rewrittenCalls)); } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java index 116c2006daa0..2cf552b9e81a 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java @@ -21,8 +21,6 @@ import java.util.ArrayList; import java.util.List; import java.util.Set; -import org.apache.calcite.rel.type.RelDataType; -import org.apache.calcite.sql.type.SqlTypeName; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.core.routing.MockRoutingManagerFactory; import org.apache.pinot.query.QueryEnvironment; @@ -44,42 +42,18 @@ import static org.testng.Assert.expectThrows; -/// Verifies MODE type inference and inferred type arguments across both multi-stage planner implementations. +/// Verifies distributed MODE arguments and the independent rollout opt-in with both physical planners. public class ModeSqlPlannerTest extends QueryEnvironmentTestBase { @DataProvider public Object[][] physicalOptimizers() { return new Object[][]{{false}, {true}}; } - @Test - public void testModeReturnTypes() { - RelDataType rowType = _queryEnvironment.compile( - "SET enableTypedMode=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3), MODE(col7), " - + "MODE(CAST(col3 AS FLOAT)), MODE(CAST(col3 AS DOUBLE)), " - + "MODE(NULLIF(JSONEXTRACTSCALAR(col1, '$.user', 'STRING', ''), '')), " - + "fromTimestamp(MODE(ts_timestamp)) FROM a") - .getRelRoot().validatedRowType; - SqlTypeName[] expectedTypes = {SqlTypeName.VARCHAR, SqlTypeName.TIMESTAMP, SqlTypeName.DOUBLE, - SqlTypeName.DOUBLE, SqlTypeName.DOUBLE, SqlTypeName.DOUBLE, SqlTypeName.VARCHAR, SqlTypeName.BIGINT}; - for (int i = 0; i < expectedTypes.length; i++) { - assertEquals(rowType.getFieldList().get(i).getType().getSqlTypeName(), expectedTypes[i]); - } - } - - @Test - public void testFilteredModeNullability() { - RelDataType rowType = _queryEnvironment.compile( - "SET enableTypedMode=true; " - + "SELECT col2, MODE(col1) FILTER (WHERE col3 > 0), MODE(ts_timestamp) FILTER (WHERE col3 > 0) " - + "FROM a GROUP BY col2").getRelRoot().validatedRowType; - assertTrue(rowType.getFieldList().get(1).getType().isNullable()); - assertTrue(rowType.getFieldList().get(2).getType().isNullable()); - } - @Test(dataProvider = "physicalOptimizers") public void testDistributedModeTypes(boolean usePhysicalOptimizer) { DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + "SET enableTypedMode=true; SELECT MODE(col1), MODE(ts_timestamp), MODE(col3) FROM a"); + + "SET enableTypedMode=true; " + + "SELECT MODE(NULLIF(col1, '')), MODE(ts_timestamp, 'MAX'), MODE(col3, 'AVG') FROM a"); PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); assertEquals(root.getDataSchema().getColumnDataTypes(), new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP, ColumnDataType.DOUBLE}); @@ -93,8 +67,8 @@ public void testDistributedModeTypes(boolean usePhysicalOptimizer) { assertEquals(calls.stream().map(RexExpression.FunctionCall::getFunctionName).toList(), List.of("MODE", "MODE", "MODE")); assertTypedCall(calls.get(0), "MIN", "STRING"); - assertTypedCall(calls.get(1), "MIN", "TIMESTAMP"); - assertEquals(calls.get(2).getFunctionOperands().size(), 1); + assertTypedCall(calls.get(1), "MAX", "TIMESTAMP"); + assertEquals(calls.get(2).getFunctionOperands().size(), 2); if (aggregate.getAggType().isOutputIntermediateFormat() && !aggregate.isLeafReturnFinalResult()) { sawIntermediate = true; assertEquals(aggregate.getDataSchema().getColumnDataTypes(), @@ -110,73 +84,15 @@ public void testDistributedModeTypes(boolean usePhysicalOptimizer) { } @Test(dataProvider = "physicalOptimizers") - public void testModeExpressionsAndTieBreakers(boolean usePhysicalOptimizer) { - DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + "SET enableTypedMode=true; " - + "SELECT col2, MODE(NULLIF(col1, ''), 'MIN'), MODE(ts_timestamp, 'MAX'), MODE(col3, 'AVG') " - + "FROM a GROUP BY col2"); - PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); - assertEquals(root.getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.TIMESTAMP, - ColumnDataType.DOUBLE}); - for (AggregateNode aggregate : findAggregates(plan)) { - assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), - List.of("MODE", "MODE", "MODE")); - assertTypedCall(aggregate.getAggCalls().get(0), "MIN", "STRING"); - assertTypedCall(aggregate.getAggCalls().get(1), "MAX", "TIMESTAMP"); - assertEquals(aggregate.getAggCalls().get(2).getFunctionOperands().size(), 2); - } - } - - @Test(dataProvider = "physicalOptimizers") - public void testModeRewriteRequiresItsOwnOptIn(boolean usePhysicalOptimizer) { + public void testTypedModeRequiresOptIn(boolean usePhysicalOptimizer) { + // Even clearing the broker's disabled-rule defaults or explicitly selecting the rule cannot bypass the opt-in. + QueryEnvironment environment = buildQueryEnvironment(Set.of()); for (String options : List.of("", "SET autoRewriteAggregationType=true; ", - "SET autoRewriteAggregationType=true; SET enableTypedMode=false; ")) { - DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + options + "SELECT MODE(ts_timestamp), MODE(col3) FROM a"); - List aggregates = findAggregates(plan); - assertFalse(aggregates.isEmpty()); - for (AggregateNode aggregate : aggregates) { - assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), - List.of("MODE", "MODE"), options); - for (RexExpression.FunctionCall mode : aggregate.getAggCalls()) { - assertEquals(mode.getFunctionOperands().size(), 1, options); - } - } - } - } - - @Test(dataProvider = "physicalOptimizers") - public void testTypedModeDoesNotEnableOtherAggregateRewrites(boolean usePhysicalOptimizer) { - DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + "SET autoRewriteAggregationType=false; SET enableTypedMode=true; " - + "SELECT MODE(col1), MODE(ts_timestamp), MODE(col3), MIN(col7), MAX(col7), SUM(col7) FROM a"); - List aggregates = findAggregates(plan); - assertFalse(aggregates.isEmpty()); - for (AggregateNode aggregate : aggregates) { - List functionNames = - aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(); - assertEquals(functionNames.subList(0, 5), List.of("MODE", "MODE", "MODE", "MIN", "MAX")); - assertFalse(functionNames.contains("SUMLONG")); - assertFalse(functionNames.contains("SUMINT")); - assertTypedCall(aggregate.getAggCalls().get(0), "MIN", "STRING"); - assertTypedCall(aggregate.getAggCalls().get(1), "MIN", "TIMESTAMP"); - } - } - - @Test(dataProvider = "physicalOptimizers") - public void testTypedModeOptInSurvivesCustomizedPlannerDefaults(boolean usePhysicalOptimizer) { - for (Set disabledRules : List.of(Set.of(), - Set.of(CommonConstants.Broker.PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE))) { - QueryEnvironment environment = buildQueryEnvironment(disabledRules); - for (String options : List.of("", "SET enableTypedMode=false; ", - "SET autoRewriteAggregationType=true; ", "SET usePlannerRules='TypedModeRewrite'; ", - "SET usePlannerRules='TypedModeRewrite'; SET enableTypedMode=false; ", - "SET enableTypedMode=true; SET skipPlannerRules='TypedModeRewrite'; ")) { - assertModeCalls(environment, usePhysicalOptimizer, options, false); - } - assertModeCalls(environment, usePhysicalOptimizer, "SET enableTypedMode=true; ", true); + "SET enableTypedMode=false; SET usePlannerRules='TypedModeRewrite'; ", + "SET enableTypedMode=true; SET skipPlannerRules='TypedModeRewrite'; ")) { + assertModeCalls(environment, usePhysicalOptimizer, options, false); } + assertModeCalls(environment, usePhysicalOptimizer, "SET enableTypedMode=true; ", true); } private static QueryEnvironment buildQueryEnvironment(Set disabledRules) { @@ -213,19 +129,6 @@ private static void assertModeCalls(QueryEnvironment environment, boolean usePhy } } - @Test(dataProvider = "physicalOptimizers") - public void testExplicitTypeArguments(boolean usePhysicalOptimizer) { - DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + "SELECT MODE(col1, 'MAX', 'STRING'), MODE(ts_timestamp, 'MIN', 'TIMESTAMP') FROM a"); - for (AggregateNode aggregate : findAggregates(plan)) { - assertTypedCall(aggregate.getAggCalls().get(0), "MAX", "STRING"); - assertTypedCall(aggregate.getAggCalls().get(1), "MIN", "TIMESTAMP"); - } - PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); - assertEquals(root.getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP}); - } - @Test public void testInvalidTypeAnnotations() { for (String expression : List.of("MODE(col1, 'MIN', 'TIMESTAMP')", "MODE(ts_timestamp, 'MIN', 'STRING')", @@ -236,15 +139,6 @@ public void testInvalidTypeAnnotations() { } } - @Test - public void testTypeAnnotationsAreCaseInsensitive() { - RelDataType rowType = _queryEnvironment.compile( - "SELECT MODE(col1, 'MIN', 'string'), MODE(ts_timestamp, 'MAX', 'TimeStamp') FROM a") - .getRelRoot().validatedRowType; - assertEquals(rowType.getFieldList().get(0).getType().getSqlTypeName(), SqlTypeName.VARCHAR); - assertEquals(rowType.getFieldList().get(1).getType().getSqlTypeName(), SqlTypeName.TIMESTAMP); - } - private static void assertTypedCall(RexExpression.FunctionCall call, String reducer, String type) { assertEquals(call.getFunctionName(), "MODE"); List arguments = call.getFunctionOperands(); diff --git a/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json index 79cbd49ed6e8..a0c0757c730c 100644 --- a/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json +++ b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json @@ -1,57 +1,34 @@ { "mode_string_timestamp_distributed": { - "comments": "Exercise typed MODE partial-state serialization across four partitions on two servers. The most frequent pallet and action user in group a never occur on the same input row.", - "extraProps": { - "enableColumnBasedNullHandling": true - }, + "comments": "Merge STRING and TIMESTAMP frequency states across four partitions on two servers.", + "extraProps": {"enableColumnBasedNullHandling": true}, "tables": { "items": { "schema": [ {"name": "lpn_id", "type": "STRING"}, {"name": "pallet_id", "type": "STRING"}, - {"name": "created_on", "type": "TIMESTAMP"}, - {"name": "meta", "type": "STRING"}, - {"name": "first_name", "type": "STRING"}, - {"name": "last_name", "type": "STRING"}, - {"name": "manifest_id", "type": "INT"} - ], - "inputs": [ - ["a", "p2", "2026-09-03 08:00:00.001", "{\"action_user\":\"u2\"}", "Bob", "Smith", 1], - ["a", "p2", "2026-09-03 08:00:00.001", "{\"action_user\":\"u3\"}", "Cara", "Jones", 1], - ["a", "p2", "2026-09-03 09:00:00.123", "{\"action_user\":\"u2\"}", "Bob", "Smith", 2], - ["a", "p1", "2026-09-03 09:00:00.123", "{\"action_user\":\"u1\"}", "Alice", "Jones", 2], - ["a", "p1", "2026-09-03 09:00:00.123", "{\"action_user\":\"u1\"}", "Alice", "Jones", 1], - ["a", "p3", "2026-09-03 10:00:00.999", "{\"action_user\":\"u1\"}", "Alice", "Jones", 2], - ["nulls", null, null, "{}", null, null, 99], - ["nulls", null, null, "{\"action_user\":\"\"}", null, null, 99], - ["ties", "z", "1970-01-01 00:00:00.001", "{}", null, null, 99], - ["ties", "a", "1969-12-31 23:59:59.999", "{}", null, null, 99], - ["ties", "z", "1970-01-01 00:00:00.001", "{}", null, null, 99], - ["ties", "a", "1969-12-31 23:59:59.999", "{}", null, null, 99], - ["empty", "", "2026-09-03 00:00:00.001", "{}", null, null, 99], - ["empty", "z", null, "{}", null, null, 99], - ["empty", "", null, "{}", null, null, 99] - ] - }, - "manifests": { - "schema": [ - {"name": "id", "type": "INT"}, - {"name": "code", "type": "STRING"}, - {"name": "status", "type": "STRING"}, {"name": "created_on", "type": "TIMESTAMP"} ], "inputs": [ - [1, "Z", "shipped", "2026-09-03 12:00:00.999"], - [2, "A", "pending", "2026-09-03 11:00:00.123"] + ["a", "p2", "2026-09-03 08:00:00.001"], + ["a", "p2", "2026-09-03 08:00:00.001"], + ["a", "p2", "2026-09-03 09:00:00.123"], + ["a", "p1", "2026-09-03 09:00:00.123"], + ["a", "p1", "2026-09-03 09:00:00.123"], + ["a", "p3", "2026-09-03 10:00:00.999"], + ["nulls", null, null], + ["nulls", null, null], + ["ties", "z", "1970-01-01 00:00:00.001"], + ["ties", "a", "1969-12-31 23:59:59.999"], + ["ties", "z", "1970-01-01 00:00:00.001"], + ["ties", "a", "1969-12-31 23:59:59.999"], + ["empty", "", "2026-09-03 00:00:00.001"], + ["empty", "z", null], + ["empty", "", null] ] } }, "queries": [ - { - "description": "Global string and timestamp modes preserve the argument types", - "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'a'", - "outputs": [["p2", "2026-09-03 09:00:00.123"]] - }, { "description": "Grouped modes ignore nulls, retain empty strings, and break ties by the smallest value", "sql": "SET enableTypedMode=true; SELECT lpn_id, MODE(pallet_id), MODE(created_on) FROM {items} GROUP BY lpn_id", @@ -62,88 +39,15 @@ ["empty", "", "2026-09-03 00:00:00.001"] ] }, - { - "description": "MODE applies independently to a stored string and computed JSON string", - "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'a'", - "outputs": [["p2", "u1"]] - }, { "description": "Typed MAX tie reducers select the largest original value", "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id, 'MAX'), MODE(created_on, 'MAX') FROM {items} WHERE lpn_id = 'ties'", "outputs": [["z", "1970-01-01 00:00:00.001"]] }, - { - "description": "Existing numeric modes retain DOUBLE results and MIN, MAX, and AVG reducers", - "sql": "SET enableTypedMode=true; SELECT MODE(manifest_id), MODE(manifest_id, 'MAX'), MODE(manifest_id, 'AVG') FROM {items} WHERE lpn_id = 'a'", - "outputs": [[1.0, 2.0, 1.5]] - }, - { - "description": "Computed null inputs do not become a string mode", - "sql": "SET enableTypedMode=true; SELECT MODE(NULLIF(JSONEXTRACTSCALAR(meta, '$.action_user', 'STRING', ''), '')) FROM {items} WHERE lpn_id = 'nulls'", - "outputs": [[null]] - }, - { - "description": "String expressions and post-aggregation transforms use the string result type", - "sql": "SET enableTypedMode=true; SELECT MODE(CONCAT(first_name, last_name, ' ')), UPPER(MODE(pallet_id)) FROM {items} WHERE lpn_id = 'a'", - "outputs": [["Alice Jones", "P2"]] - }, { "description": "Empty aggregate inputs return null for both types", "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'missing'", "outputs": [[null, null]] - }, - { - "description": "Filtered typed modes return null while preserving their group", - "sql": "SET enableTypedMode=true; SELECT lpn_id, MODE(pallet_id) FILTER (WHERE manifest_id < 0), MODE(created_on) FILTER (WHERE manifest_id < 0) FROM {items} WHERE lpn_id = 'a' GROUP BY lpn_id", - "outputs": [["a", null, null]] - }, - { - "description": "Timestamp mode can feed a parent comparison; fixture ingestion uses Los Angeles time and the SQL literal uses UTC", - "sql": "SET enableTypedMode=true; SELECT lpn_id FROM (SELECT lpn_id, MODE(created_on) AS mode_time FROM {items} GROUP BY lpn_id) WHERE mode_time = TIMESTAMP '2026-09-03 16:00:00.123'", - "outputs": [["a"]] - }, - { - "description": "Joined grouping supports independent string, JSON, and timestamp modes, including unmatched outer rows", - "sql": "SET enableTypedMode=true; SELECT i.lpn_id, MODE(i.pallet_id), MODE(NULLIF(JSONEXTRACTSCALAR(i.meta, '$.action_user', 'STRING', ''), '')), MODE(m.code), MODE(m.status), MODE(m.created_on) FROM {items} i LEFT JOIN {manifests} m ON m.id = i.manifest_id GROUP BY i.lpn_id", - "outputs": [ - ["a", "p2", "u1", "A", "pending", "2026-09-03 11:00:00.123"], - ["nulls", null, null, null, null, null], - ["ties", "a", null, null, null, null], - ["empty", "", null, null, null, null] - ] - } - ] - }, - "mode_string_timestamp_direct": { - "comments": "Replicated input exercises AGGREGATE_DIRECT finalization when the physical optimizer colocates a global aggregate on one server.", - "extraProps": { - "enableColumnBasedNullHandling": true - }, - "tables": { - "items": { - "replicated": true, - "schema": [ - {"name": "pallet_id", "type": "STRING"}, - {"name": "created_on", "type": "TIMESTAMP"} - ], - "inputs": [ - ["p2", "2026-09-03 08:00:00.001"], - ["p2", "2026-09-03 09:00:00.123"], - ["p1", "2026-09-03 09:00:00.123"], - [null, null] - ] - } - }, - "queries": [ - { - "description": "Direct aggregate emits scalar string and timestamp values instead of frequency-map intermediates", - "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items}", - "outputs": [["p2", "2026-09-03 09:00:00.123"]] - }, - { - "description": "Direct aggregate with no matching values emits typed nulls", - "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id) FILTER (WHERE pallet_id = 'missing'), MODE(created_on) FILTER (WHERE pallet_id = 'missing') FROM {items}", - "outputs": [[null, null]] } ] } From df8fcc232224c930cb9838568ae3f170a3f2d4c7 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 7 Sep 2026 16:57:09 -0700 Subject: [PATCH 6/8] Enable MODE type inference by default --- ...deAggregationFunctionRewriteOptimizer.java | 7 +-- ...gregationFunctionRewriteOptimizerTest.java | 14 ++--- ...notModeAggregationFunctionRewriteRule.java | 2 +- .../apache/pinot/query/QueryEnvironment.java | 17 ++---- .../query/queries/ModeSqlPlannerTest.java | 54 +------------------ .../resources/queries/ModeAggregates.json | 6 +-- .../pinot/spi/utils/CommonConstants.java | 5 -- 7 files changed, 15 insertions(+), 90 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java index da6aad85d38b..ff1996ff4898 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java @@ -35,10 +35,9 @@ import org.apache.pinot.spi.data.DateTimeFormatSpec; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; -import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; -/// When `enableTypedMode` is enabled, adds an inferred type argument to string and timestamp MODE expressions +/// Adds an inferred type argument to string and timestamp MODE expressions /// so the result type is fixed before execution. /// This also supplies the broker with the correct result type when no rows match or groups are trimmed before /// finalization. Type inference reads schema and function metadata only: server-dependent transforms such as LOOKUP @@ -46,9 +45,7 @@ public class ModeAggregationFunctionRewriteOptimizer implements StatementOptimizer { @Override public void optimize(PinotQuery pinotQuery, @Nullable Schema schema) { - // Keep existing timestamp MODE requests compatible with older servers during rolling upgrades. - if (schema == null || pinotQuery.getQueryOptions() == null - || !Boolean.parseBoolean(pinotQuery.getQueryOptions().get(QueryOptionKey.ENABLE_TYPED_MODE))) { + if (schema == null) { return; } rewriteExpressions(pinotQuery.getSelectList(), schema); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java index 552f7ba6ff4f..304a16bc5742 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java @@ -29,7 +29,7 @@ import static org.testng.Assert.assertEquals; -/// Covers inferred MODE arguments and the independent rollout opt-in. +/// Covers automatic MODE type inference and unchanged numeric calls. public class ModeAggregationFunctionRewriteOptimizerTest { private static final QueryOptimizer OPTIMIZER = new QueryOptimizer(); private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName("testTable") @@ -55,25 +55,21 @@ public Object[][] modeExpressions() { @Test(dataProvider = "modeExpressions") public void testModeExpressions(String original, String rewritten) { - String prefix = "SET enableTypedMode=true; SELECT "; + String prefix = "SELECT "; TestHelper.assertEqualsQuery(prefix + original + " FROM testTable", prefix + rewritten + " FROM testTable", SCHEMA); } @Test public void testLegacyCallsAndOtherRewritesAreUnchanged() { - for (String options : new String[]{"", "SET enableTypedMode=false; ", "SET autoRewriteAggregationType=true; "}) { - assertUnchanged(options + "SELECT MODE(stringCol), MODE(timestampCol), MODE(longCol, 'AVG') FROM testTable", - SCHEMA); - } - assertUnchanged("SET enableTypedMode=true; SELECT MODE(longCol, 'AVG'), MIN(stringCol), SUM(longCol), " + assertUnchanged("SELECT MODE(longCol, 'AVG'), MIN(stringCol), SUM(longCol), " + "MODE(stringCol, 'MAX', 'STRING'), MODE(timestampCol, 'MIN', 'TIMESTAMP') FROM testTable", SCHEMA); } @Test public void testUnknownTypesDoNotInitializeServerTransforms() { - assertUnchanged("SET enableTypedMode=true; SELECT MODE(unknownCol), MODE(mvStringCol), " + assertUnchanged("SELECT MODE(unknownCol), MODE(mvStringCol), " + "MODE(LOOKUP('baseballTeams', 'teamInteger', 'teamID', stringCol)) FROM testTable", SCHEMA); - assertUnchanged("SET enableTypedMode=true; SELECT MODE(stringCol) FROM testTable", null); + assertUnchanged("SELECT MODE(stringCol) FROM testTable", null); } private static void assertUnchanged(String sql, Schema schema) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java index d3e3746cb0c4..44ccca23fad9 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java @@ -36,7 +36,7 @@ import org.apache.calcite.sql.type.SqlTypeName; -/// Supplies string and timestamp MODE calls with an inferred type argument after an explicit rollout opt-in. +/// Supplies string and timestamp MODE calls with an inferred type argument. /// This stateless rule keeps the reducer and type as projected literals so distributed stages retain both arguments. public class PinotModeAggregationFunctionRewriteRule extends RelOptRule { public static PinotModeAggregationFunctionRewriteRule instanceWithDescription(String description) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index 61e6de8f5914..bfa24cac481e 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -167,7 +167,7 @@ public QueryEnvironment(Config config, MultiClusterRoutingContext multiClusterRo rootSchema, List.of(database), _typeFactory, CONNECTION_CONFIG, config.isCaseSensitive()); _defaultDisabledPlannerRules = _envConfig.defaultDisabledPlannerRules(); // default optProgram with no skip rule options and no use rule options - _optProgram = getOptProgram(_envConfig.getRuleSet(), Set.of(), Set.of(), _defaultDisabledPlannerRules, false); + _optProgram = getOptProgram(_envConfig.getRuleSet(), Set.of(), Set.of(), _defaultDisabledPlannerRules); _multiClusterRoutingContext = multiClusterRoutingContext; } @@ -199,16 +199,11 @@ private PlannerContext getPlannerContext(SqlNodeAndOptions sqlNodeAndOptions) { if (Boolean.parseBoolean(options.get(QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE))) { useRuleSet.add(CommonConstants.Broker.PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE); } - boolean enableTypedMode = Boolean.parseBoolean(options.get(QueryOptionKey.ENABLE_TYPED_MODE)); - if (enableTypedMode) { - useRuleSet.add(CommonConstants.Broker.PlannerRuleNames.TYPED_MODE_REWRITE); - } if (MapUtils.isNotEmpty(options)) { Set skipRuleSet = QueryOptionsUtils.getSkipPlannerRules(options); if (!skipRuleSet.isEmpty() || !useRuleSet.isEmpty()) { // dynamically create optProgram according to rule options - optProgram = getOptProgram(_envConfig.getRuleSet(), skipRuleSet, useRuleSet, _defaultDisabledPlannerRules, - enableTypedMode); + optProgram = getOptProgram(_envConfig.getRuleSet(), skipRuleSet, useRuleSet, _defaultDisabledPlannerRules); } } int sortExchangeCopyLimit = QueryOptionsUtils.getSortExchangeCopyThreshold(options, @@ -551,15 +546,9 @@ private DispatchableSubPlan toDispatchableSubPlan(RelRoot relRoot, PlannerContex /// @param skipRuleSet parsed skipped rule name set from query options /// @param useRuleSet parsed use rule set from query options /// @param defaultDisabledRuleSet parsed default disabled rule set from broker config - /// @param enableTypedMode whether the query explicitly opted into typed MODE implementations /// @return HepProgram that performs logical transformations private static HepProgram getOptProgram(PinotRuleSet ruleSet, Set skipRuleSet, Set useRuleSet, - Set defaultDisabledRuleSet, boolean enableTypedMode) { - if (!enableTypedMode) { - // Rollout safety must not depend on customized default-disabled rules or usePlannerRules overrides. - skipRuleSet = new HashSet<>(skipRuleSet); - skipRuleSet.add(CommonConstants.Broker.PlannerRuleNames.TYPED_MODE_REWRITE); - } + Set defaultDisabledRuleSet) { HepProgramBuilder hepProgramBuilder = new HepProgramBuilder(); // Set the match order as DEPTH_FIRST. The default is arbitrary which works the same as DEPTH_FIRST, but it's // best to be explicit. diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java index 2cf552b9e81a..fe4f187160a8 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/ModeSqlPlannerTest.java @@ -20,19 +20,14 @@ import java.util.ArrayList; import java.util.List; -import java.util.Set; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; -import org.apache.pinot.core.routing.MockRoutingManagerFactory; -import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.QueryEnvironmentTestBase; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.PlanNode; -import org.apache.pinot.query.routing.WorkerManager; import org.apache.pinot.spi.exception.QueryException; -import org.apache.pinot.spi.utils.CommonConstants; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -42,7 +37,7 @@ import static org.testng.Assert.expectThrows; -/// Verifies distributed MODE arguments and the independent rollout opt-in with both physical planners. +/// Verifies automatic distributed MODE type inference with both physical planners. public class ModeSqlPlannerTest extends QueryEnvironmentTestBase { @DataProvider public Object[][] physicalOptimizers() { @@ -52,7 +47,6 @@ public Object[][] physicalOptimizers() { @Test(dataProvider = "physicalOptimizers") public void testDistributedModeTypes(boolean usePhysicalOptimizer) { DispatchableSubPlan plan = _queryEnvironment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + "SET enableTypedMode=true; " + "SELECT MODE(NULLIF(col1, '')), MODE(ts_timestamp, 'MAX'), MODE(col3, 'AVG') FROM a"); PlanNode root = plan.getQueryStageMap().get(0).getPlanFragment().getFragmentRoot(); assertEquals(root.getDataSchema().getColumnDataTypes(), @@ -83,52 +77,6 @@ public void testDistributedModeTypes(boolean usePhysicalOptimizer) { assertTrue(sawFinal, "Distributed MODE must retain its final result types"); } - @Test(dataProvider = "physicalOptimizers") - public void testTypedModeRequiresOptIn(boolean usePhysicalOptimizer) { - // Even clearing the broker's disabled-rule defaults or explicitly selecting the rule cannot bypass the opt-in. - QueryEnvironment environment = buildQueryEnvironment(Set.of()); - for (String options : List.of("", "SET autoRewriteAggregationType=true; ", - "SET enableTypedMode=false; SET usePlannerRules='TypedModeRewrite'; ", - "SET enableTypedMode=true; SET skipPlannerRules='TypedModeRewrite'; ")) { - assertModeCalls(environment, usePhysicalOptimizer, options, false); - } - assertModeCalls(environment, usePhysicalOptimizer, "SET enableTypedMode=true; ", true); - } - - private static QueryEnvironment buildQueryEnvironment(Set disabledRules) { - MockRoutingManagerFactory factory = new MockRoutingManagerFactory(1, 2); - TABLE_SCHEMAS.forEach((name, schema) -> factory.registerTable(schema, name)); - SERVER1_SEGMENTS.forEach((table, segments) -> segments.forEach(s -> factory.registerSegment(1, table, s))); - SERVER2_SEGMENTS.forEach((table, segments) -> segments.forEach(s -> factory.registerSegment(2, table, s))); - return new QueryEnvironment(QueryEnvironment.configBuilder() - .requestId(-1L) - .database(CommonConstants.DEFAULT_DATABASE) - .tableCache(factory.buildTableCache()) - .workerManager(new WorkerManager("Broker_localhost", "localhost", 3, factory.buildRoutingManager(null))) - .defaultDisabledPlannerRules(disabledRules) - .build()); - } - - private static void assertModeCalls(QueryEnvironment environment, boolean usePhysicalOptimizer, String options, - boolean typed) { - DispatchableSubPlan plan = environment.planQuery("SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " - + options + "SELECT MODE(col1), MODE(ts_timestamp), MODE(col3) FROM a"); - List aggregates = findAggregates(plan); - assertFalse(aggregates.isEmpty()); - for (AggregateNode aggregate : aggregates) { - assertEquals(aggregate.getAggCalls().stream().map(RexExpression.FunctionCall::getFunctionName).toList(), - List.of("MODE", "MODE", "MODE"), options); - if (typed) { - assertTypedCall(aggregate.getAggCalls().get(0), "MIN", "STRING"); - assertTypedCall(aggregate.getAggCalls().get(1), "MIN", "TIMESTAMP"); - } else { - for (RexExpression.FunctionCall mode : aggregate.getAggCalls()) { - assertEquals(mode.getFunctionOperands().size(), 1, options); - } - } - } - } - @Test public void testInvalidTypeAnnotations() { for (String expression : List.of("MODE(col1, 'MIN', 'TIMESTAMP')", "MODE(ts_timestamp, 'MIN', 'STRING')", diff --git a/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json index a0c0757c730c..de72a8941234 100644 --- a/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json +++ b/pinot-query-runtime/src/test/resources/queries/ModeAggregates.json @@ -31,7 +31,7 @@ "queries": [ { "description": "Grouped modes ignore nulls, retain empty strings, and break ties by the smallest value", - "sql": "SET enableTypedMode=true; SELECT lpn_id, MODE(pallet_id), MODE(created_on) FROM {items} GROUP BY lpn_id", + "sql": "SELECT lpn_id, MODE(pallet_id), MODE(created_on) FROM {items} GROUP BY lpn_id", "outputs": [ ["a", "p2", "2026-09-03 09:00:00.123"], ["nulls", null, null], @@ -41,12 +41,12 @@ }, { "description": "Typed MAX tie reducers select the largest original value", - "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id, 'MAX'), MODE(created_on, 'MAX') FROM {items} WHERE lpn_id = 'ties'", + "sql": "SELECT MODE(pallet_id, 'MAX'), MODE(created_on, 'MAX') FROM {items} WHERE lpn_id = 'ties'", "outputs": [["z", "1970-01-01 00:00:00.001"]] }, { "description": "Empty aggregate inputs return null for both types", - "sql": "SET enableTypedMode=true; SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'missing'", + "sql": "SELECT MODE(pallet_id), MODE(created_on) FROM {items} WHERE lpn_id = 'missing'", "outputs": [[null, null]] } ] diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index a650ef4df5f3..bdb005ca0f3f 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1044,10 +1044,6 @@ public static class QueryOptionKey { // MAX(stringCol) -> MAXSTRING(stringCol) // SUM(intCol) -> SUMINT(intCol) public static final String AUTO_REWRITE_AGGREGATION_TYPE = "autoRewriteAggregationType"; - - /// Opts into string and TIMESTAMP MODE implementations after all query components have been upgraded. - /// Kept separate from existing aggregate rewrites to preserve timestamp MODE semantics during rolling upgrades. - public static final String ENABLE_TYPED_MODE = "enableTypedMode"; // When enabled, allows multi cluster/federated queries to be executed. public static final String ENABLE_MULTI_CLUSTER_ROUTING = "enableMultiClusterRouting"; @@ -1167,7 +1163,6 @@ public static class PlannerRuleNames { PlannerRuleNames.AGGREGATE_UNION_AGGREGATE, PlannerRuleNames.JOIN_TO_ENRICHED_JOIN, PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE, - PlannerRuleNames.TYPED_MODE_REWRITE, // Stock Calcite rule kept opt-in via usePlannerRules — see SORT_PROJECT_TRANSPOSE javadoc // above for the rationale (firing in BASIC_RULES disrupts ProjectToSemiJoinRule on // partition-hinted IN(SELECT) queries, breaking colocated broadcast semi-joins). From 287647cbf4d162ec4e76f3235c9a6be73060eeed Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 7 Sep 2026 18:57:30 -0700 Subject: [PATCH 7/8] Update explain plan expectation for typed MODE rule --- .../pinot/integration/tests/OfflineClusterIntegrationTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java index e59cf4ebe85f..bb10e3a6fa7a 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java @@ -3664,6 +3664,7 @@ public void testExplainPlanQueryV2() + " PinotLogicalAggregate(group=[{23}], agg#0=[COUNT()], aggType=[LEAF])\n" + " PinotLogicalTableScan(table=[[default, mytable]])\n"); assertEquals(response1Json.get("rows").get(0).get(2).asText(), "Rule Execution Times\n" + + "Rule: TypedModeRewrite -> Time:*\n" + "Rule: SortRemove -> Time:*\n" + "Rule: AggregateProjectMerge -> Time:*\n" + "Rule: AggregateProjectPullUpConstants -> Time:*\n" From 0904ef96b74f1c3b0f37ab05ef54b40a6ad78ba4 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 8 Sep 2026 00:43:36 -0700 Subject: [PATCH 8/8] Infer aggregate arguments through shared planning paths --- .../core/query/optimizer/QueryOptimizer.java | 3 +- .../AggregateFunctionRewriteOptimizer.java | 147 ++++++++++++-- ...deAggregationFunctionRewriteOptimizer.java | 180 ------------------ ...ggregateFunctionRewriteOptimizerTest.java} | 11 +- .../tests/OfflineClusterIntegrationTest.java | 1 - .../PinotAggregateExchangeNodeInsertRule.java | 2 + ...notModeAggregationFunctionRewriteRule.java | 116 ----------- .../calcite/rel/rules/PinotQueryRuleSets.java | 2 - .../calcite/rel/rules/PinotRuleUtils.java | 36 ++++ .../v2/opt/rules/AggregatePushdownRule.java | 3 + .../query/queries/ModeSqlPlannerTest.java | 45 ++++- .../segment/spi/AggregationFunctionType.java | 28 ++- .../pinot/spi/utils/CommonConstants.java | 1 - 13 files changed, 249 insertions(+), 326 deletions(-) delete mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java rename pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/{ModeAggregationFunctionRewriteOptimizerTest.java => AggregateFunctionRewriteOptimizerTest.java} (87%) delete mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotModeAggregationFunctionRewriteRule.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java index 68b1002e3aec..635250800e03 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java @@ -31,7 +31,6 @@ import org.apache.pinot.core.query.optimizer.filter.TextMatchFilterOptimizer; import org.apache.pinot.core.query.optimizer.filter.TimePredicateFilterOptimizer; import org.apache.pinot.core.query.optimizer.statement.AggregateFunctionRewriteOptimizer; -import org.apache.pinot.core.query.optimizer.statement.ModeAggregationFunctionRewriteOptimizer; import org.apache.pinot.core.query.optimizer.statement.StatementOptimizer; import org.apache.pinot.spi.data.Schema; @@ -49,7 +48,7 @@ public class QueryOptimizer { new MergeRangeFilterOptimizer(), new TextMatchFilterOptimizer()); private static final List STATEMENT_OPTIMIZERS = - List.of(new AggregateFunctionRewriteOptimizer(), new ModeAggregationFunctionRewriteOptimizer()); + List.of(new AggregateFunctionRewriteOptimizer()); /// Optimizes the given query. public void optimize(PinotQuery pinotQuery, @Nullable Schema schema) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java index f88c1cf372de..0451cefc1498 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java @@ -18,18 +18,28 @@ */ package org.apache.pinot.core.query.optimizer.statement; +import java.util.ArrayList; import java.util.List; +import java.util.Locale; import javax.annotation.Nullable; +import org.apache.pinot.common.function.FunctionInfo; +import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.common.function.FunctionUtils; import org.apache.pinot.common.request.Expression; import org.apache.pinot.common.request.Function; import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.spi.data.DateTimeFieldSpec; +import org.apache.pinot.spi.data.DateTimeFormatSpec; import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.CommonConstants; -/// Rewrites certain aggregation functions based on operand types to support polymorphic aggregations. +/// Binds inferred arguments declared by aggregation functions and optionally rewrites type-specific variants. /// /// Currently supported rewrites: /// - MIN(stringType) -> MINSTRING @@ -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 selectList = pinotQuery.getSelectList(); if (selectList != null) { for (Expression expression : selectList) { - maybeRewriteAggregateFunction(expression, schema); + maybeRewriteAggregateFunction(expression, schema, autoRewrite); } } List groupByList = pinotQuery.getGroupByList(); if (groupByList != null) { for (Expression expression : groupByList) { - maybeRewriteAggregateFunction(expression, schema); + maybeRewriteAggregateFunction(expression, schema, autoRewrite); } } List 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 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 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(); + }); + if (!inferredArguments.isEmpty()) { + List typedOperands = new ArrayList<>(operands); + for (String argument : inferredArguments) { + typedOperands.add(RequestUtils.getLiteralExpression(argument)); + } + function.setOperands(typedOperands); + } + if (!autoRewrite || operands.isEmpty()) { + return; + } + Expression operand = function.getOperands().get(0); FieldSpec.DataType operandType; // TODO: Handle more complex expressions (e.g. MIN(trim(stringCol)) ) @@ -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 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; + } + + @Nullable + private static ColumnDataType literalType(List arguments, int position) { + if (arguments.size() <= position || !arguments.get(position).isSetLiteral() + || !arguments.get(position).getLiteral().isSetStringValue()) { + return null; + } + String type = arguments.get(position).getLiteral().getStringValue().toUpperCase(Locale.ROOT); + return switch (type) { + case "VARCHAR", "CHAR", "JSON" -> ColumnDataType.STRING; + case "BIGINT" -> ColumnDataType.LONG; + case "INTEGER" -> ColumnDataType.INT; + case "REAL" -> ColumnDataType.FLOAT; + case "DECIMAL" -> ColumnDataType.BIG_DECIMAL; + default -> { + try { + yield ColumnDataType.valueOf(type); + } catch (IllegalArgumentException e) { + yield null; + } + } + }; + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java deleted file mode 100644 index ff1996ff4898..000000000000 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizer.java +++ /dev/null @@ -1,180 +0,0 @@ -/** - * 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.core.query.optimizer.statement; - -import java.util.ArrayList; -import java.util.List; -import java.util.Locale; -import javax.annotation.Nullable; -import org.apache.pinot.common.function.FunctionInfo; -import org.apache.pinot.common.function.FunctionRegistry; -import org.apache.pinot.common.function.FunctionUtils; -import org.apache.pinot.common.request.Expression; -import org.apache.pinot.common.request.Function; -import org.apache.pinot.common.request.PinotQuery; -import org.apache.pinot.common.utils.DataSchema.ColumnDataType; -import org.apache.pinot.common.utils.request.RequestUtils; -import org.apache.pinot.segment.spi.AggregationFunctionType; -import org.apache.pinot.spi.data.DateTimeFieldSpec; -import org.apache.pinot.spi.data.DateTimeFormatSpec; -import org.apache.pinot.spi.data.FieldSpec; -import org.apache.pinot.spi.data.Schema; - - -/// Adds an inferred type argument to string and timestamp MODE expressions -/// so the result type is fixed before execution. -/// This also supplies the broker with the correct result type when no rows match or groups are trimmed before -/// finalization. Type inference reads schema and function metadata only: server-dependent transforms such as LOOKUP -/// must not be initialized on the broker. Expressions whose type is unknown retain the legacy MODE implementation. -public class ModeAggregationFunctionRewriteOptimizer implements StatementOptimizer { - @Override - public void optimize(PinotQuery pinotQuery, @Nullable Schema schema) { - if (schema == null) { - return; - } - rewriteExpressions(pinotQuery.getSelectList(), schema); - rewriteExpressions(pinotQuery.getGroupByList(), schema); - rewriteExpressions(pinotQuery.getOrderByList(), schema); - rewriteExpression(pinotQuery.getFilterExpression(), schema); - rewriteExpression(pinotQuery.getHavingExpression(), schema); - } - - private static void rewriteExpressions(@Nullable List expressions, Schema schema) { - if (expressions != null) { - for (Expression expression : expressions) { - rewriteExpression(expression, schema); - } - } - } - - private static void rewriteExpression(@Nullable Expression expression, Schema schema) { - if (expression == null || !expression.isSetFunctionCall()) { - return; - } - Function function = expression.getFunctionCall(); - List operands = function.getOperands(); - rewriteExpressions(operands, schema); - if (!AggregationFunctionType.MODE.getName().equalsIgnoreCase(function.getOperator()) || operands.isEmpty() - || operands.size() >= 3) { - return; - } - - ColumnDataType operandType = getOperandType(operands.get(0), schema); - if (operandType == ColumnDataType.STRING || operandType == ColumnDataType.TIMESTAMP) { - List typedOperands = new ArrayList<>(operands); - if (typedOperands.size() == 1) { - typedOperands.add(RequestUtils.getLiteralExpression("MIN")); - } - typedOperands.add(RequestUtils.getLiteralExpression(operandType.name())); - function.setOperands(typedOperands); - } - } - - @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 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; - } - - @Nullable - private static ColumnDataType literalType(List arguments, int position) { - if (arguments.size() <= position || !arguments.get(position).isSetLiteral() - || !arguments.get(position).getLiteral().isSetStringValue()) { - return null; - } - String type = arguments.get(position).getLiteral().getStringValue().toUpperCase(Locale.ROOT); - return switch (type) { - case "VARCHAR", "CHAR", "JSON" -> ColumnDataType.STRING; - case "BIGINT" -> ColumnDataType.LONG; - case "INTEGER" -> ColumnDataType.INT; - case "REAL" -> ColumnDataType.FLOAT; - case "DECIMAL" -> ColumnDataType.BIG_DECIMAL; - default -> { - try { - yield ColumnDataType.valueOf(type); - } catch (IllegalArgumentException e) { - yield null; - } - } - }; - } -} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java similarity index 87% rename from pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java index 304a16bc5742..6941f16c6f38 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/ModeAggregationFunctionRewriteOptimizerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java @@ -29,8 +29,8 @@ import static org.testng.Assert.assertEquals; -/// Covers automatic MODE type inference and unchanged numeric calls. -public class ModeAggregationFunctionRewriteOptimizerTest { +/// Covers inferred aggregate arguments and unchanged legacy calls. +public class AggregateFunctionRewriteOptimizerTest { private static final QueryOptimizer OPTIMIZER = new QueryOptimizer(); private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName("testTable") .addSingleValueDimension("stringCol", DataType.STRING) @@ -63,11 +63,16 @@ public void testModeExpressions(String original, String rewritten) { public void testLegacyCallsAndOtherRewritesAreUnchanged() { assertUnchanged("SELECT MODE(longCol, 'AVG'), MIN(stringCol), SUM(longCol), " + "MODE(stringCol, 'MAX', 'STRING'), MODE(timestampCol, 'MIN', 'TIMESTAMP') FROM testTable", SCHEMA); + TestHelper.assertEqualsQuery( + "SET autoRewriteAggregationType=true; SELECT MODE(stringCol), MIN(stringCol), MAX(longCol), SUM(longCol), " + + "COUNT(*) FROM testTable", + "SET autoRewriteAggregationType=true; SELECT MODE(stringCol, 'MIN', 'STRING'), MINSTRING(stringCol), " + + "MAXLONG(longCol), SUMLONG(longCol), COUNT(*) FROM testTable", SCHEMA); } @Test public void testUnknownTypesDoNotInitializeServerTransforms() { - assertUnchanged("SELECT MODE(unknownCol), MODE(mvStringCol), " + assertUnchanged("SELECT MODE(unknownCol), MODE(mvStringCol), MODE(NULLIF(longCol, 0)), " + "MODE(LOOKUP('baseballTeams', 'teamInteger', 'teamID', stringCol)) FROM testTable", SCHEMA); assertUnchanged("SELECT MODE(stringCol) FROM testTable", null); } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java index bb10e3a6fa7a..e59cf4ebe85f 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java @@ -3664,7 +3664,6 @@ public void testExplainPlanQueryV2() + " PinotLogicalAggregate(group=[{23}], agg#0=[COUNT()], aggType=[LEAF])\n" + " PinotLogicalTableScan(table=[[default, mytable]])\n"); assertEquals(response1Json.get("rows").get(0).get(2).asText(), "Rule Execution Times\n" - + "Rule: TypedModeRewrite -> Time:*\n" + "Rule: SortRemove -> Time:*\n" + "Rule: AggregateProjectMerge -> Time:*\n" + "Rule: AggregateProjectPullUpConstants -> Time:*\n" diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java index 61a0b33cfebb..fc52168c8d18 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java @@ -511,6 +511,7 @@ private static PinotLogicalAggregate convertAggFromIntermediateInput(Aggregate a } } } + rexList = PinotRuleUtils.inferAggregateArguments(orgAggCall, input, rexList); aggCalls.add(buildAggCall(exchange, orgAggCall, rexList, aggColumnOffset, aggType, leafReturnFinalResult)); } @@ -550,6 +551,7 @@ private static List buildAggCalls(Aggregate aggRel, List projects = new ArrayList<>(); - List names = new ArrayList<>(input.getRowType().getFieldNames()); - if (input instanceof Project) { - projects.addAll(((Project) input).getProjects()); - } else { - for (int i = 0; i < names.size(); i++) { - projects.add(RexInputRef.of(i, input.getRowType())); - } - } - RexBuilder rexBuilder = input.getCluster().getRexBuilder(); - List originalCalls = aggregate.getAggCallList(); - List rewrittenCalls = new ArrayList<>(originalCalls.size()); - boolean changed = false; - for (AggregateCall originalCall : originalCalls) { - List arguments = originalCall.getArgList(); - List rewritten = arguments; - if (originalCall.getAggregation().getKind() == SqlKind.MODE && !arguments.isEmpty() && arguments.size() < 3) { - SqlTypeName operandType = input.getRowType().getFieldList().get(arguments.get(0)).getType().getSqlTypeName(); - String type = SqlTypeName.STRING_TYPES.contains(operandType) - ? "STRING" - : operandType == SqlTypeName.TIMESTAMP ? "TIMESTAMP" : null; - if (type != null) { - rewritten = new ArrayList<>(arguments); - if (arguments.size() == 1) { - rewritten.add(addLiteral(rexBuilder, projects, names, "MIN")); - } - rewritten.add(addLiteral(rexBuilder, projects, names, type)); - changed = true; - } - } - rewrittenCalls.add(originalCall.withArgList(rewritten)); - } - if (!changed) { - return; - } - - RelNode rewrittenInput; - if (input instanceof Project) { - // Extend the existing projection: wrapping it in identity refs would hide the original reducer literals. - Project project = (Project) input; - RelDataTypeFactory.Builder rowType = input.getCluster().getTypeFactory().builder(); - for (int i = 0; i < projects.size(); i++) { - rowType.add(names.get(i), projects.get(i).getType()); - } - rewrittenInput = project.copy(project.getTraitSet(), project.getInput(), projects, rowType.build()); - } else { - rewrittenInput = LogicalProject.create(input, List.of(), projects, names); - } - call.transformTo(aggregate.copy(aggregate.getTraitSet(), rewrittenInput, aggregate.getGroupSet(), - aggregate.getGroupSets(), rewrittenCalls)); - } - - private static int addLiteral(RexBuilder rexBuilder, List projects, List names, String value) { - RexNode literal = rexBuilder.makeLiteral(value); - int index = projects.indexOf(literal); - if (index < 0) { - index = projects.size(); - projects.add(literal); - names.add("$mode$" + index); - } - return index; - } -} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java index 7e204b3fa076..38822d4dfad7 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java @@ -174,8 +174,6 @@ private PinotQueryRuleSets() { PinotAggregateFunctionRewriteRule .instanceWithDescription(PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE), - PinotModeAggregationFunctionRewriteRule - .instanceWithDescription(PlannerRuleNames.TYPED_MODE_REWRITE), // convert CASE-style filtered aggregates into true filtered aggregates // put it after AGGREGATE_REDUCE_FUNCTIONS where SUM is converted to SUM0 diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java index b7ad7a4edb2c..aba6b4dc3a06 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java @@ -27,6 +27,7 @@ import org.apache.calcite.plan.hep.HepRelVertex; 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.Exchange; import org.apache.calcite.rel.core.Filter; import org.apache.calcite.rel.core.Join; @@ -34,6 +35,7 @@ import org.apache.calcite.rel.core.RelFactories; import org.apache.calcite.rel.core.TableScan; import org.apache.calcite.rel.core.Window; +import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexCall; import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexLiteral; @@ -51,6 +53,10 @@ import org.apache.calcite.util.Util; import org.apache.pinot.calcite.rel.hint.PinotHintStrategyTable; import org.apache.pinot.common.function.sql.PinotSqlFunction; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.planner.logical.RelToPlanNodeConverter; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.spi.data.FieldSpec.DataType; public class PinotRuleUtils { @@ -97,6 +103,36 @@ public static boolean isAggregate(RelNode rel) { return unboxRel(rel) instanceof Aggregate; } + /// Appends arguments inferred by the aggregation's metadata using its original input types. Final stages must + /// resolve against the original input rather than the intermediate accumulator referenced by `arguments`. + public static List inferAggregateArguments(AggregateCall aggregateCall, RelNode originalInput, + List arguments) { + String functionName = aggregateCall.getAggregation().getName(); + if (!AggregationFunctionType.isAggregationFunction(functionName)) { + return arguments; + } + List argList = aggregateCall.getArgList(); + List inferredArguments = AggregationFunctionType.getAggregationFunctionType(functionName) + .inferArguments(argList.size(), ordinal -> { + ColumnDataType type = RelToPlanNodeConverter.convertToColumnDataType( + originalInput.getRowType().getFieldList().get(argList.get(ordinal)).getType()); + if (type.isArray() || type == ColumnDataType.OBJECT) { + return DataType.UNKNOWN; + } + return type == ColumnDataType.MAP ? DataType.MAP : type.toDataType(); + }); + if (inferredArguments.isEmpty()) { + return arguments; + } + List resolvedArguments = new ArrayList<>(arguments.size() + inferredArguments.size()); + resolvedArguments.addAll(arguments); + RexBuilder rexBuilder = originalInput.getCluster().getRexBuilder(); + for (String argument : inferredArguments) { + resolvedArguments.add(rexBuilder.makeLiteral(argument)); + } + return resolvedArguments; + } + /// utility logic to determine if a JOIN can be pushed down to the leaf-stage execution and leverage the /// segment-local info (indexing and others) to speed up the execution. /// diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java index ca30c851a013..cd0ca78fbc1e 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java @@ -236,6 +236,7 @@ private static PhysicalAggregate convertAggForGroupingSets(PhysicalAggregate phy } } } + rexList = PinotRuleUtils.inferAggregateArguments(orgAggCall, aggRel.getInput(), rexList); aggCalls.add(buildAggCall(exchange, orgAggCall, rexList, finalGroupCount, AggType.FINAL, false)); } ImmutableBitSet groupSet = ImmutableBitSet.range(finalGroupCount); @@ -322,6 +323,7 @@ private static PhysicalAggregate convertAggFromIntermediateInput(PhysicalAggrega } } } + rexList = PinotRuleUtils.inferAggregateArguments(orgAggCall, input, rexList); aggCalls.add(buildAggCall(exchange, orgAggCall, rexList, groupCount, aggType, leafReturnFinalResult)); } ImmutableBitSet.Builder groupSetBuilder = ImmutableBitSet.builder(); @@ -370,6 +372,7 @@ public static List buildAggCalls(Aggregate aggRel, List aggregates = findAggregates(plan); assertFalse(aggregates.isEmpty()); @@ -59,24 +62,56 @@ public void testDistributedModeTypes(boolean usePhysicalOptimizer) { for (AggregateNode aggregate : aggregates) { List calls = aggregate.getAggCalls(); assertEquals(calls.stream().map(RexExpression.FunctionCall::getFunctionName).toList(), - List.of("MODE", "MODE", "MODE")); + List.of("MODE", "MODE", "MODE", "MODE")); assertTypedCall(calls.get(0), "MIN", "STRING"); assertTypedCall(calls.get(1), "MAX", "TIMESTAMP"); assertEquals(calls.get(2).getFunctionOperands().size(), 2); + assertTypedCall(calls.get(3), "MAX", "STRING"); if (aggregate.getAggType().isOutputIntermediateFormat() && !aggregate.isLeafReturnFinalResult()) { sawIntermediate = true; assertEquals(aggregate.getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.OBJECT, ColumnDataType.OBJECT, ColumnDataType.OBJECT}); + new ColumnDataType[]{ColumnDataType.OBJECT, ColumnDataType.OBJECT, ColumnDataType.OBJECT, + ColumnDataType.OBJECT}); } else { sawFinal = true; assertEquals(aggregate.getDataSchema().getColumnDataTypes(), - new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP, ColumnDataType.DOUBLE}); + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.TIMESTAMP, ColumnDataType.DOUBLE, + ColumnDataType.STRING}); } } assertTrue(sawIntermediate, "Distributed MODE must exchange frequency counts"); assertTrue(sawFinal, "Distributed MODE must retain its final result types"); } + @Test(dataProvider = "physicalOptimizers") + public void testGroupedModeTypes(boolean usePhysicalOptimizer) { + List queries = List.of( + "SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */ col2, MODE(col1), MODE(ts_timestamp, 'MAX') " + + "FROM a GROUP BY col2", + "SELECT col3, GROUPING(col3), MODE(col1), MODE(ts_timestamp, 'MAX') " + + "FROM a GROUP BY GROUPING SETS ((col3), ())"); + for (int i = 0; i < queries.size(); i++) { + DispatchableSubPlan plan = _queryEnvironment.planQuery( + "SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; " + queries.get(i)); + List aggregates = findAggregates(plan); + assertFalse(aggregates.isEmpty()); + for (AggregateNode aggregate : aggregates) { + List calls = aggregate.getAggCalls(); + assertEquals(calls.size(), 2); + assertTypedCall(calls.get(0), "MIN", "STRING"); + assertTypedCall(calls.get(1), "MAX", "TIMESTAMP"); + if (aggregate.getAggType() == AggType.FINAL) { + for (int j = 0; j < calls.size(); j++) { + RexExpression.InputRef input = (RexExpression.InputRef) calls.get(j).getFunctionOperands().get(0); + assertEquals(input.getIndex(), aggregate.getGroupKeys().size() + j); + } + } + } + AggType expectedType = i == 0 ? AggType.DIRECT : AggType.FINAL; + assertTrue(aggregates.stream().anyMatch(aggregate -> aggregate.getAggType() == expectedType)); + } + } + @Test public void testInvalidTypeAnnotations() { for (String expression : List.of("MODE(col1, 'MIN', 'TIMESTAMP')", "MODE(ts_timestamp, 'MIN', 'STRING')", diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java index 8acee5ad7105..83865ef1cd27 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Locale; import java.util.Set; +import java.util.function.IntFunction; import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nullable; @@ -38,6 +39,7 @@ import org.apache.calcite.sql.type.SqlTypeName; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.Strings; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; import static com.google.common.base.Preconditions.checkArgument; @@ -48,9 +50,8 @@ /// - '$' is allowed in the name field but not in the enum name. /// /// This enum is used both in the v1 engine and multistage engine to define the allowed Pinot aggregation functions. -/// The v1 engine only relies on the 'name' field, whereas all the other fields are used in the multistage engine -/// to register the aggregation function with Calcite. This allows using a unified approach to aggregations across both -/// the v1 and multistage engines. +/// Both engines use the function name and argument inference. The SQL metadata also registers the aggregation +/// function with Calcite in the multistage engine. public enum AggregationFunctionType { // Aggregation functions for single-valued columns COUNT("count"), @@ -73,7 +74,19 @@ public enum AggregationFunctionType { OperandTypes.family(List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER), i -> i == 1), OperandTypes.family(List.of(SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER), i -> i == 1), OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER)), - ReturnTypes.explicit(SqlTypeName.OTHER), null, SqlKind.MODE), + ReturnTypes.explicit(SqlTypeName.OTHER), null, SqlKind.MODE) { + @Override + public List inferArguments(int argumentCount, IntFunction argumentTypes) { + if (argumentCount == 0 || argumentCount >= 3) { + return List.of(); + } + DataType inputType = argumentTypes.apply(0); + if (inputType != DataType.STRING && inputType != DataType.TIMESTAMP) { + return List.of(); + } + return argumentCount == 1 ? List.of("MIN", inputType.name()) : List.of(inputType.name()); + } + }, ANYVALUE("anyValue", ReturnTypes.ARG0, OperandTypes.ANY, SqlTypeName.OTHER), FIRSTWITHTIME("firstWithTime", ReturnTypes.ARG0, OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER), SqlTypeName.OTHER), @@ -375,6 +388,13 @@ public static String getNormalizedAggregationFunctionName(String functionName) { return Strings.CS.remove(StringUtils.remove(functionName, '_').toUpperCase(), "$"); } + /// Infers internal string literal arguments to append before execution. + /// Argument types are resolved lazily so functions without inferred arguments do not require expression typing. + /// Unknown or multi-value argument types are represented by [DataType#UNKNOWN]. + public List inferArguments(int argumentCount, IntFunction argumentTypes) { + return List.of(); + } + /// Returns the corresponding aggregation function type for the given function name. /// /// NOTE: Underscores in the function name are ignored. diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index bdb005ca0f3f..50cb94b4a152 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1116,7 +1116,6 @@ public static class PlannerRuleNames { public static final String AGGREGATE_UNION_TRANSPOSE = "AggregateUnionTranspose"; public static final String AGGREGATE_REDUCE_FUNCTIONS = "AggregateReduceFunctions"; public static final String AGGREGATE_FUNCTION_REWRITE = "AggregateFunctionRewrite"; - public static final String TYPED_MODE_REWRITE = "TypedModeRewrite"; public static final String AGGREGATE_CASE_TO_FILTER = "AggregateCaseToFilter"; public static final String PROJECT_FILTER_TRANSPOSE = "ProjectFilterTranspose"; public static final String PROJECT_MERGE = "ProjectMerge";