Fix GROUP BY DISTINCTCOUNT ClassCastException in mergeDataTablesOnly (#18842)

mergeDataTablesOnly reused getIndexedTable, which unconditionally called
indexedTable.finish(true, true). For a GROUP BY query with an OBJECT-typed
intermediate aggregate (DISTINCTCOUNT, DISTINCTCOUNTHLL, etc.), that
storeFinalResult=true call:
- mutated DataSchema._columnDataTypes from OBJECT to the final-result type
  in place (IndexedTable.java:170-174), and
- replaced each row's value with extractFinalResult(value) (Set → Integer
  for DISTINCTCOUNT) at IndexedTable.java:195.

The subsequent buildIntermediateDataTable read DataSchema's cached
_storedColumnDataTypes — populated earlier from the pre-finalize OBJECT
schema and never invalidated by finish — so the OBJECT branch fired and
handed the now-Integer value into
BaseDistinctAggregateAggregationFunction.serializeIntermediateResult(Set),
which threw ClassCastException.

Fix: move finish() out of getIndexedTable and let each caller pick the
right mode. reduceResult keeps finish(true, true) (downstream consumers
expect final scalars). mergeDataTablesOnly calls finish(true, false) so
aggregate values stay as intermediates and round-trip through
buildIntermediateDataTable correctly.

Adds a testGroupByDistinctCountObjectRoundTrip regression: server emits
intermediate OBJECT-encoded Sets, merge produces a re-injectable
intermediate DataTable, and reducing that intermediate matches a direct
reduce of the same servers.

Author: navina

pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java

@@ -133,11 +133,11 @@ public DataTable mergeDataTablesOnly(String tableName, DataSchema dataSchema,
     // When servers are configured to return final aggregate state, the input DataTables hold final
     // (not intermediate) values, so the merge-only contract — "produce an intermediate DataTable that
     // can be re-merged via the normal reduce path" — cannot be honored.
-    if (_queryContext.isServerReturnFinalResult()) {
+    if (_queryContext.isServerReturnFinalResult() || _queryContext.isServerReturnFinalResultKeyUnpartitioned()) {
       throw new UnsupportedOperationException(
           "Merge-only reduction is not supported when servers return final aggregate results "
-              + "(server.returnFinalResult / isServerReturnFinalResult); input would be final-typed, "
-              + "not intermediate.");
+              + "(serverReturnFinalResult / serverReturnFinalResultKeyUnpartitioned); input would be "
+              + "final-typed, not intermediate.");
     }
     dataSchema = ReducerDataSchemaUtils.canonicalizeDataSchemaForGroupBy(_queryContext, dataSchema);
     try {
@@ -147,6 +147,8 @@ public DataTable mergeDataTablesOnly(String tableName, DataSchema dataSchema,
       Collection<DataTable> dataTables = dataTableMap.values();
       // Reuse the regular reduce's merge: builds the IndexedTable of group keys + intermediate agg state.
       IndexedTable indexedTable = getIndexedTable(dataSchema, dataTables, reducerContext);
+      // Keep aggregate values as intermediates so the output can be re-merged through the regular reduce path.
+      indexedTable.finish(false, false);
       DataTable mergedDataTable = buildIntermediateDataTable(dataSchema, indexedTable);
       if (indexedTable.isTrimmed() && _queryContext.isUnsafeTrim()) {
         mergedDataTable.getMetadata().put(MetadataKey.GROUPS_TRIMMED.getName(), "true");
@@ -175,6 +177,10 @@ private void reduceResult(BrokerResponseNative brokerResponseNative, DataSchema
       BrokerMetrics brokerMetrics) {
     // NOTE: This step will modify the data schema and also return final aggregate results.
     IndexedTable indexedTable = getIndexedTable(dataSchema, dataTables, reducerContext);
+    // Sort + finalize: mutate _dataSchema column types to final-result types and replace each row's value
+    // with extractFinalResult(...). Required by the regular reduce path: the downstream consumers
+    // (PostAggregationHandler, HavingFilterHandler, ResultTable serialization) expect final scalars.
+    indexedTable.finish(true, true);
     if (indexedTable.isTrimmed() && _queryContext.isUnsafeTrim()) {
       brokerResponseNative.setGroupsTrimmed(true);
     }
@@ -408,7 +414,12 @@ public void runJob() {
       }
     }
 
-    indexedTable.finish(true, true);
+    // NOTE: finish(...) is invoked by the caller, not here. The two callers want different semantics:
+    //   - reduceResult (regular reduce → ResultTable) needs finish(true, true) to extract final results.
+    //   - mergeDataTablesOnly (merge-only intermediate output) needs finish(true, false) so the aggregate
+    //     values stay as intermediates (e.g. Set for DISTINCTCOUNT, not Integer); otherwise the next
+    //     buildIntermediateDataTable's serializeIntermediateResult(...) call casts the final scalar back to
+    //     the intermediate type and throws ClassCastException.
     return indexedTable;
   }
 

pinot-core/src/test/java/org/apache/pinot/core/query/reduce/MergeDataTablesOnlyTest.java

@@ -118,6 +118,34 @@ public void testGroupByRoundTrip() {
     assertRoundTrip(query, serverTables);
   }
 
+  /// Verifies that a GROUP BY query selecting DISTINCTCOUNT keeps OBJECT aggregate state re-mergeable.
+  @Test
+  public void testGroupByDistinctCountObjectRoundTrip()
+      throws IOException {
+    String query = "SELECT col1, DISTINCTCOUNT(col2) FROM testTable GROUP BY col1";
+    BrokerRequest brokerRequest = CalciteSqlCompiler.compileToBrokerRequest(query);
+    BrokerResponseNative baseline = reduce(brokerRequest, toMap(buildDistinctCountServerTables(query)));
+    DataTable merged = merge(brokerRequest, toMap(buildDistinctCountServerTables(query)));
+    assertNotNull(merged, "merge produced null");
+    BrokerResponseNative viaMerge = reduce(brokerRequest, singletonMap(merged));
+    assertResultTablesEquivalent(baseline.getResultTable(), viaMerge.getResultTable());
+  }
+
+  /// Two per-server intermediate DataTables for selecting `DISTINCTCOUNT(col2)` grouped by `col1`.
+  private static List<DataTable> buildDistinctCountServerTables(String query)
+      throws IOException {
+    AggregationFunction aggFunction = aggFunctions(query)[0];
+    DataSchema schema = new DataSchema(new String[]{"col1", "distinctcount(col2)"},
+        new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.OBJECT});
+    return List.of(
+        buildGroupByWithObject(schema, aggFunction,
+            new Object[]{1, new IntOpenHashSet(new int[]{1, 2, 3})},
+            new Object[]{2, new IntOpenHashSet(new int[]{4, 5})}),
+        buildGroupByWithObject(schema, aggFunction,
+            new Object[]{1, new IntOpenHashSet(new int[]{3, 4, 5})},
+            new Object[]{3, new IntOpenHashSet(new int[]{6, 7})}));
+  }
+
   @Test
   public void testDistinctRoundTrip() {
     String query = "SELECT DISTINCT col1 FROM testTable";
@@ -151,6 +179,18 @@ public void testGroupByServerReturnFinalResultRejected() {
     assertThrows(UnsupportedOperationException.class, () -> merge(brokerRequest, map));
   }
 
+  @Test
+  public void testGroupByServerReturnFinalResultKeyUnpartitionedRejected() {
+    BrokerRequest brokerRequest = CalciteSqlCompiler.compileToBrokerRequest(
+        "SELECT col1, DISTINCTCOUNT(col2) FROM testTable GROUP BY col1");
+    brokerRequest.getPinotQuery()
+        .putToQueryOptions(Broker.Request.QueryOptionKey.SERVER_RETURN_FINAL_RESULT_KEY_UNPARTITIONED, "true");
+    DataSchema schema = new DataSchema(new String[]{"col1", "distinctcount(col2)"},
+        new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT});
+    Map<ServerRoutingInstance, DataTable> map = singletonMap(buildGroupBy(schema, new Object[][]{{1, 2}}));
+    assertThrows(UnsupportedOperationException.class, () -> merge(brokerRequest, map));
+  }
+
   @Test
   public void testGroupByLimitZeroRoundTrip() {
     // LIMIT 0 on group-by: mergeDataTablesOnly does NOT apply LIMIT (intermediate is unlimited),
@@ -489,6 +529,27 @@ private static DataTable buildGroupBy(DataSchema schema, Object[][] rows) {
     }
   }
 
+  /**
+   * Builds a GROUP BY DataTable with one INT key column followed by one OBJECT aggregate column whose
+   * value is the intermediate {@code Set} (or other aggregate state) the server would have emitted.
+   * Mirrors what {@code GroupByResultsBlock.getDataTable()} produces on the server side for the
+   * non-null-handling path: scalar key written directly, OBJECT column written via
+   * {@code serializeIntermediateResult}.
+   */
+  private static DataTable buildGroupByWithObject(DataSchema schema, AggregationFunction aggFunction,
+      Object[]... rows)
+      throws IOException {
+    DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(schema);
+    for (Object[] row : rows) {
+      builder.startRow();
+      // Single INT key column at index 0; OBJECT aggregate intermediate at index 1.
+      builder.setColumn(0, (int) row[0]);
+      builder.setColumn(1, aggFunction.serializeIntermediateResult(row[1]));
+      builder.finishRow();
+    }
+    return builder.build();
+  }
+
   private static DataTable buildObjectRow(DataSchema schema, AggregationFunction aggFunction, Object intermediate)
       throws IOException {
     DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(schema);