Fix GROUP BY DISTINCTCOUNT ClassCastException in mergeDataTablesOnly

[navina]

Background

BrokerReduceService exposes two reduce paths:

• reduceResult(...) — the normal path: merge per-server results and finalize.
• mergeDataTablesOnly(...) — merge-only: produce an intermediate DataTable that can be re-merged later via the normal reduce path. Used by callers that need to combine partial scatters without finalizing.

The merge-only path's contract is that aggregate values stay as intermediates end-to-end.

What's wrong

OSS's single-stage broker request handler auto-sets SERVER_RETURN_FINAL_RESULT=true when the routing table targets exactly one server (BaseSingleStageBrokerRequestHandler.java:844). That optimization causes servers to return finalized scalars (e.g. Integer for DISTINCTCOUNT, Long for DISTINCTCOUNTHLL) instead of intermediate sketches — which directly contradicts the merge-only contract: the input is already final-typed, not intermediate.

Two ways this manifested in GroupByDataTableReducer.mergeDataTablesOnly:

1. No guard for the single-server-final-result mode. When _queryContext.isServerReturnFinalResult() is true, the input DataTables hold final values, so merge-only cannot honor its contract. The previous code silently proceeded and produced a malformed intermediate.

2. getIndexedTable over-finalized internally. Even when servers correctly returned intermediates, the shared helper unconditionally called indexedTable.finish(true, true). For a GROUP BY query with an OBJECT-typed intermediate aggregate (DISTINCTCOUNT, DISTINCTCOUNTHLL, etc.), storeFinalResult=true:

   • 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 lazily-populated _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

• Reject the single-server final-result mode explicitly in mergeDataTablesOnly: throw UnsupportedOperationException with a clear message naming SERVER_RETURN_FINAL_RESULT. Callers expecting an intermediate must either disable this option or use reduceResult instead.
• 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.

Tests

Adds testGroupByDistinctCountObjectRoundTrip in MergeDataTablesOnlyTest: servers emit intermediate OBJECT-encoded Sets, merge produces a re-injectable intermediate DataTable, and reducing that intermediate matches a direct reduce of the same per-server inputs.

Release Notes

Fixes a ClassCastException in mergeDataTablesOnly when merging GROUP BY results that include OBJECT-typed intermediate aggregates (e.g. DISTINCTCOUNT, DISTINCTCOUNTHLL). mergeDataTablesOnly now also throws UnsupportedOperationException up front when called against a query whose servers will return final results (e.g. the single-server SERVER_RETURN_FINAL_RESULT auto-set), instead of silently producing a malformed intermediate.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 83.33333% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 64.82%. Comparing base (302783f) to head (cdfb22b).
⚠️ Report is 60 commits behind head on master.

Files with missing lines	Patch %	Lines
...not/core/query/reduce/GroupByDataTableReducer.java	83.33%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #18842      +/-   ##
============================================
+ Coverage     64.77%   64.82%   +0.04%     
- Complexity     1322     1347      +25     
============================================
  Files          3393     3396       +3     
  Lines        211022   212504    +1482     
  Branches      33135    33484     +349     
============================================
+ Hits         136687   137746    +1059     
- Misses        63322    63597     +275     
- Partials      11013    11161     +148     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-21	64.82% <83.33%> (+0.04%)	⬆️
temurin	64.82% <83.33%> (+0.04%)	⬆️
unittests	64.81% <83.33%> (+0.04%)	⬆️
unittests1	56.78% <83.33%> (-0.21%)	⬇️
unittests2	37.23% <0.00%> (+0.06%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Copilot]

Pull request overview

Fixes a ClassCastException in the broker-side merge-only reduce path for GROUP BY queries containing OBJECT-typed intermediate aggregates (e.g. DISTINCTCOUNT), ensuring merge-only output remains re-mergeable as an intermediate DataTable.

Changes:

• Adjust GroupByDataTableReducer so IndexedTable.finish(...) is invoked by callers with the correct semantics: merge-only keeps intermediate aggregate state, while normal reduce finalizes scalars.
• Add a regression test covering GROUP BY + DISTINCTCOUNT OBJECT-intermediate round-tripping through mergeDataTablesOnly and then the normal reduce path.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java	Moves finish(...) responsibility to callers and uses finish(true, false) for merge-only to preserve intermediate aggregate values.
pinot-core/src/test/java/org/apache/pinot/core/query/reduce/MergeDataTablesOnlyTest.java	Adds a regression test and helper to build OBJECT-aggregate group-by DataTables that emulate server intermediate encoding.