Support string and timestamp MODE aggregations

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<Expression> selectList = pinotQuery.getSelectList();
     if (selectList != null) {
       for (Expression expression : selectList) {
-        maybeRewriteAggregateFunction(expression, schema);
+        maybeRewriteAggregateFunction(expression, schema, autoRewrite);
       }
     }
 
     List<Expression> groupByList = pinotQuery.getGroupByList();
     if (groupByList != null) {
       for (Expression expression : groupByList) {
-        maybeRewriteAggregateFunction(expression, schema);
+        maybeRewriteAggregateFunction(expression, schema, autoRewrite);
       }
     }
 
     List<Expression> orderByList = pinotQuery.getOrderByList();
     if (orderByList != null) {
       for (Expression expression : orderByList) {
-        maybeRewriteAggregateFunction(expression, schema);
+        maybeRewriteAggregateFunction(expression, schema, autoRewrite);
       }
     }
 
-    maybeRewriteAggregateFunction(pinotQuery.getFilterExpression(), schema);
-    maybeRewriteAggregateFunction(pinotQuery.getHavingExpression(), schema);
+    maybeRewriteAggregateFunction(pinotQuery.getFilterExpression(), schema, autoRewrite);
+    maybeRewriteAggregateFunction(pinotQuery.getHavingExpression(), schema, autoRewrite);
   }
 
-  private void maybeRewriteAggregateFunction(@Nullable Expression expression, Schema schema) {
+  private void maybeRewriteAggregateFunction(@Nullable Expression expression, Schema schema, boolean autoRewrite) {
     if (expression == null || !expression.isSetFunctionCall()) {
       return;
     }
 
     Function function = expression.getFunctionCall();
+    List<Expression> operands = function.getOperands();
+    for (Expression operand : operands) {
+      // Infer arguments within scalar expressions while preserving the existing top-level variant rewrites.
+      maybeRewriteAggregateFunction(operand, schema, false);
+    }
     String functionName = function.getOperator();
     if (!AggregationFunctionType.isAggregationFunction(functionName)) {
       return;
     }
 
+    List<String> inferredArguments = AggregationFunctionType.getAggregationFunctionType(functionName)
+        .inferArguments(operands.size(), i -> {
+          ColumnDataType type = getOperandType(operands.get(i), schema);
+          if (type == null || type.isArray() || type == ColumnDataType.OBJECT) {
+            return DataType.UNKNOWN;
+          }
+          return type == ColumnDataType.MAP ? DataType.MAP : type.toDataType();
+        });
+    if (!inferredArguments.isEmpty()) {
+      List<Expression> 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<Expression> arguments = function.getOperands();
+    String name = FunctionRegistry.canonicalize(function.getOperator());
+    switch (name) {
+      case "cast":
+        return literalType(arguments, 1);
+      case "jsonextractscalar":
+      case "jsonextractscalarfast":
+      case "jsonextractscalarfirstmatch":
+      case "jsonextractscalarfory":
+        return literalType(arguments, 2);
+      case "case":
+        // CASE stores alternating condition/result pairs followed by an optional ELSE result.
+        ColumnDataType resultType = ColumnDataType.UNKNOWN;
+        for (int i = 1; i < arguments.size(); i += 2) {
+          resultType = commonType(resultType, getOperandType(arguments.get(i), schema));
+        }
+        if (arguments.size() % 2 == 1) {
+          resultType = commonType(resultType, getOperandType(arguments.get(arguments.size() - 1), schema));
+        }
+        return resultType;
+      case "datetimeconvert":
+        if (arguments.size() < 3 || !arguments.get(2).isSetLiteral()
+            || !arguments.get(2).getLiteral().isSetStringValue()) {
+          return null;
+        }
+        DateTimeFieldSpec.TimeFormat format =
+            new DateTimeFormatSpec(arguments.get(2).getLiteral().getStringValue()).getTimeFormat();
+        return format == DateTimeFieldSpec.TimeFormat.EPOCH || format == DateTimeFieldSpec.TimeFormat.TIMESTAMP
+            ? ColumnDataType.LONG
+            : ColumnDataType.STRING;
+      default:
+        ColumnDataType[] argumentTypes = new ColumnDataType[arguments.size()];
+        for (int i = 0; i < arguments.size(); i++) {
+          argumentTypes[i] = getOperandType(arguments.get(i), schema);
+          if (argumentTypes[i] == null) {
+            return null;
+          }
+        }
+        FunctionInfo functionInfo = FunctionRegistry.lookupFunctionInfo(name, argumentTypes);
+        return functionInfo != null ? FunctionUtils.getColumnDataType(functionInfo.getMethod().getReturnType()) : null;
+    }
+  }
+
+  @Nullable
+  private static ColumnDataType commonType(@Nullable ColumnDataType left, @Nullable ColumnDataType right) {
+    if (left == null || right == null) {
+      return null;
+    }
+    if (left == ColumnDataType.UNKNOWN) {
+      return right;
+    }
+    return right == ColumnDataType.UNKNOWN || left == right ? left : null;
+  }
+
+  @Nullable
+  private static ColumnDataType literalType(List<Expression> 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;
+        }
+      }
+    };
+  }
 }

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]);
+        }
       }
     }
 

pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/ModeAggregationFunctionTest.java

@@ -19,13 +19,24 @@
 
 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.queries.FluentQueryTest;
 import org.apache.pinot.spi.config.table.FieldConfig;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.apache.pinot.spi.utils.PinotDataType;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.expectThrows;
+
 
 public class ModeAggregationFunctionTest extends AbstractAggregationFunctionTest {
 
@@ -40,6 +51,58 @@ Object[] scenarios() {
     };
   }
 
+  @DataProvider
+  Object[] typedScenarios() {
+    return new Object[]{new Scenario(DataType.STRING, true), new Scenario(DataType.STRING, false),
+        new Scenario(DataType.TIMESTAMP, true), new Scenario(DataType.TIMESTAMP, false)};
+  }
+
+  @Test(dataProvider = "typedScenarios")
+  void typedModeMergesCountsAndHandlesNulls(Scenario scenario) {
+    String type = scenario._dataType.name();
+    String first = scenario._dataType == DataType.STRING ? "apple" : "2026-09-03 10:11:12.123";
+    String shared = scenario._dataType == DataType.STRING ? "banana" : "2026-09-04 10:11:12.456";
+    String last = scenario._dataType == DataType.STRING ? "cherry" : "2026-09-05 10:11:12.789";
+    // Each instance has a different local mode; the shared value wins only after merging full counts.
+    scenario.getDeclaringTable(true)
+        .onFirstInstance("myField", first, first, first, shared, shared, "null")
+        .andOnSecondInstance("myField", last, last, last, shared, shared, "null")
+        .whenQuery("select mode(myField, 'MIN', '" + type + "') as mode from testTable")
+        .thenResultTextIs("mode[" + type + "]\n" + shared)
+        .whenQuery("select mode(myField, 'MIN', '" + type + "') from testTable where myField is null")
+        .thenResultIs(new Object[]{null})
+        .whenQuery("select myField, mode(myField, 'MIN', '" + type + "') from testTable "
+            + "group by myField order by myField")
+        .thenResultIs(new Object[]{first, first}, new Object[]{shared, shared}, new Object[]{last, last},
+            new Object[]{null, null});
+  }
+
+  @DataProvider
+  Object[][] typedStates() {
+    return new Object[][]{{"STRING", "", "zebra"}, {"TIMESTAMP", 9_007_199_254_740_992L, 9_007_199_254_740_993L}};
+  }
+
+  @Test(dataProvider = "typedStates")
+  void typedModePreservesSerializedValuesAndTieReducers(String type, Object smaller, Object larger) {
+    ModeAggregationFunction min = typedMode(type, "MIN");
+    AggregationFunction.SerializedIntermediateResult serialized =
+        min.serializeIntermediateResult(Map.of(smaller, 2L, larger, 1L));
+    Map<?, Long> 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"));
+  }
+
+  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 {
     private final DataType _dataType;
     private final boolean _dictionary;

pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java

@@ -0,0 +1,86 @@
+/**
+ * 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;
+
+
+/// 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)
+      .addSingleValueDimension("longCol", DataType.LONG)
+      .addMultiValueDimension("mvStringCol", DataType.STRING)
+      .addDateTime("timestampCol", DataType.TIMESTAMP, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS")
+      .build();
+
+  @DataProvider
+  public Object[][] modeExpressions() {
+    return new Object[][]{
+        {"MODE(stringCol)", "MODE(stringCol, 'MIN', 'STRING')"},
+        {"MODE(timestampCol, 'MAX')", "MODE(timestampCol, 'MAX', 'TIMESTAMP')"},
+        {"MODE(CONCAT(stringCol, 'suffix'))", "MODE(CONCAT(stringCol, 'suffix'), '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')"},
+        {"fromTimestamp(MODE(timestampCol))", "fromTimestamp(MODE(timestampCol, 'MIN', 'TIMESTAMP'))"}
+    };
+  }
+
+  @Test(dataProvider = "modeExpressions")
+  public void testModeExpressions(String original, String rewritten) {
+    String prefix = "SELECT ";
+    TestHelper.assertEqualsQuery(prefix + original + " FROM testTable", prefix + rewritten + " FROM testTable", SCHEMA);
+  }
+
+  @Test
+  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), MODE(NULLIF(longCol, 0)), "
+        + "MODE(LOOKUP('baseballTeams', 'teamInteger', 'teamID', stringCol)) FROM testTable", SCHEMA);
+    assertUnchanged("SELECT MODE(stringCol) FROM testTable", null);
+  }
+
+  private static void assertUnchanged(String sql, Schema schema) {
+    PinotQuery original = CalciteSqlParser.compileToPinotQuery(sql);
+    PinotQuery optimized = original.deepCopy();
+    OPTIMIZER.optimize(optimized, schema);
+    assertEquals(optimized, original);
+  }
+}