gortiz commented on code in PR #19264:
URL: https://github.com/apache/pinot/pull/19264#discussion_r4106188236


##########
pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java:
##########
@@ -172,6 +184,24 @@ protected void processSegments() {
     }
   }
 
+  /// Resolves the merge-time record layout from the first result block's 
schema. A base-aggregation grouping-set
+  /// segment emits base groups whose schema has no synthetic $groupingId 
column, so it has exactly
+  /// `numUnionColumns + numAggregationFunctions` columns; the expansion path 
(and non-grouping-set queries)
+  /// carries the full `$groupingId`-including layout. Deriving the mode from 
the schema (rather than the query
+  /// context alone) also covers the MV-column carve-out, where a grouping-set 
query still uses expansion.
+  /// Must be called under the same synchronization that guards the one-time 
`_indexedTable` creation.
+  private void resolveMergeLayout(GroupByResultsBlock resultsBlock) {

Review Comment:
   **MAJOR: layout is decided from the FIRST block only, and never checked 
again**
   
   The base-vs-expansion choice is made per segment (DefaultGroupByExecutor: 
hasMVGroupByExpression comes from that segment's column metadata), but the 
combine takes the layout from whichever block arrives first. If segments 
disagree (e.g. a column that is SV in old segments and MV in newer ones, or a 
custom GroupByExecutor that returns the default false), blocks with the other 
width are upserted with the wrong _numKeyColumns. That is silent corruption or 
an AIOOBE. Suggest: make the decision query-level (from the table schema), or 
at least check each block's schema width in processSegments and fail fast.



##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java:
##########
@@ -89,20 +93,29 @@ public DefaultGroupByExecutor(QueryContext queryContext, 
AggregationFunction[] a
       // isDictionaryEncoded() flag rather than gating on dictionary nullness 
alone.
       hasNoDictionaryGroupByExpression |= !columnContext.isDictionaryEncoded();
     }
-    // Grouping-set queries expand each row into one group per grouping set, 
so they always use the
-    // multi-value (int[][]) executor path even though the union group-by 
columns are single-valued.
-    boolean groupingSets = queryContext.isGroupingSets();
-    _hasMVGroupByExpression = hasMVGroupByExpression || groupingSets;
+    /// A grouping-set query using the legacy per-row expansion path expands 
each row into one group per
+    /// grouping set, so it always uses the multi-value (int[][]) executor 
path even though the union group-by
+    /// columns are single-valued. With base aggregation (the default), the 
segment instead aggregates the base
+    /// grouping (union columns) exactly like a plain GROUP BY -- no forced MV 
path, no grouping-set key
+    /// generator -- and emits the base groups; the combine phase merges them 
and derives the individual grouping
+    /// sets in parallel. This keeps total work minimal (each row grouped 
once, not once per set) regardless of
+    /// cardinality. Base aggregation is disabled only when a group-by column 
is multi-valued: an MV column fans
+    /// a row across its values in the base grouping, so rolling that column 
up would over-count the row (see
+    /// [QueryContext#isGroupingSetsBaseAggregation]).
+    boolean useBaseAggregation = queryContext.isGroupingSets() && 
queryContext.isGroupingSetsBaseAggregation()

Review Comment:
   **MAJOR correctness: numGroupsLimit now corrupts the ROLLUP totals**
   
   GroupByUtils.createIndexedTableForCombineOperator caps the grouping-set 
combine table at numGroupsLimit. For grouping sets the combine table is capped 
at numGroupsLimit. On the base path it now holds BASE groups, and the segment 
key generator is capped the same way. When the limit is hit, a row whose base 
key is new gets dropped from EVERY set, including the grand total and the 
coarse subtotals. On the expansion path the grand-total/coarse groups are 
created early and keep aggregating, so the totals stayed exact and only the 
fine groups were lost. Now the rows most users look at (totals) become wrong 
under the limit. Both paths set numGroupsLimitReached, but this is a quality 
regression on by default. No test covers limit-hit + base path; 
DeriveGroupingSetsTest only exercises the derive, after the base cap.



##########
pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java:
##########
@@ -87,6 +100,230 @@ public static GroupByResultsBlock 
buildGroupingSetsResultsBlock(QueryContext que
     return resultsBlock;
   }
 
+  /// Derives the individual grouping sets from a merged BASE-grouping 
[IndexedTable] (union columns aggregated
+  /// once, like a plain GROUP BY), for a GROUP BY GROUPING SETS / ROLLUP / 
CUBE query using base aggregation.
+  /// Each base group is projected into every grouping set: its rolled-up 
(non-participating) columns are set to
+  /// `null`, the `$groupingId` discriminator is inserted after the union 
columns, and the base group's
+  /// aggregation intermediates are merged into the derived group. This moves 
the per-set fan-out from O(rows) to
+  /// O(base groups) and runs it here -- after the row-collapsing base merge 
-- across the combine's threads,
+  /// rather than expanding every scanned row.
+  ///
+  /// The grouping sets are partitioned round-robin across up to `numTasks` 
worker tasks; each task derives its
+  /// owned sets over all base entries into a task-local plain HashMap. Task 
key spaces are disjoint (every
+  /// derived key carries its set ordinal), so the heavy merge work runs 
contention-free -- no concurrent table,
+  /// no bin locks, no shared size counters -- and the task maps are then 
unioned single-threaded into the result
+  /// table with plain stores (never merges). Profiling showed the previous 
shared-concurrent-table derive spent
+  /// most of its time in ConcurrentHashMap machinery rather than in the 
projection itself.
+  ///
+  /// Clone discipline: a base group's intermediate flows into every grouping 
set (and is read concurrently by
+  /// other tasks), while [AggregationFunction#merge] mutates/returns its 
arguments. Each derived group's stored
+  /// accumulator is therefore a clone (see [#cloneIntermediate]), and OBJECT 
intermediates are re-cloned per
+  /// merge so a merge can never mutate the shared base accumulator another 
task is still reading. Scalar
+  /// intermediates are immutable and clone-free.
+  public static IndexedTable 
deriveGroupingSetsFromMergedBaseTable(IndexedTable baseTable, QueryContext 
queryContext,
+      int numTasks, ExecutorService executorService) {
+    AggregationFunction[] aggregationFunctions = 
queryContext.getAggregationFunctions();
+    assert aggregationFunctions != null;
+    int numAggregationFunctions = aggregationFunctions.length;
+    List<int[]> groupingSets = queryContext.getGroupingSets();
+    int numSets = groupingSets.size();
+    int numUnionColumns = queryContext.getGroupByExpressions().size();
+    // Per grouping set: membership mask over the union columns (true = 
participates, false = rolled up to NULL).
+    boolean[][] setContains = new boolean[numSets][numUnionColumns];
+    for (int s = 0; s < numSets; s++) {
+      for (int columnIndex : groupingSets.get(s)) {
+        setContains[s][columnIndex] = true;
+      }
+    }
+
+    // Grouping-set output schema: the base schema with the synthetic 
$groupingId INT column inserted right after
+    // the union group-by columns (mirroring GroupByOperator's grouping-set 
schema layout).
+    DataSchema groupingSetsSchema = 
insertGroupingIdColumn(baseTable.getDataSchema(), numUnionColumns);
+    // The derive must not drop groups: it is a bounded transformation of the 
already-bounded base groups (whose
+    // count is capped at numGroupsLimit per segment), so the derived count is 
at most numGroupsLimit * numSets --
+    // a finite amount. Capping the derived table at numGroupsLimit would drop 
derived groups NON-DETERMINISTICALLY
+    // under the parallel upsert (whichever threads fill the quota first win), 
which can starve an entire
+    // low-magnitude grouping set such as the grand total. Use an unbounded 
result size here and defer the real
+    // ORDER BY + LIMIT to the broker (per-set trim below still bounds it when 
explicitly configured).
+    int derivedUpperBound = (int) Math.min((long) baseTable.size() * numSets, 
Integer.MAX_VALUE);
+    int initialCapacity = getIndexedTableInitialCapacity(derivedUpperBound, 
derivedUpperBound,
+        queryContext.getMinInitialIndexedTableCapacity());
+
+    List<Map.Entry<Key, Record>> baseEntries = new 
ArrayList<>(baseTable.getRecordEntries());
+    int numTaskSlots = Math.max(1, Math.min(numTasks, numSets));
+    List<Map<Key, Record>> taskMaps;
+    if (numTaskSlots == 1) {
+      taskMaps = List.of(deriveSets(baseEntries, setContains, 0, 1, 
numUnionColumns, numAggregationFunctions,
+          aggregationFunctions));
+    } else {
+      List<Future<Map<Key, Record>>> futures = new ArrayList<>(numTaskSlots);
+      for (int t = 0; t < numTaskSlots; t++) {
+        int taskIndex = t;
+        futures.add(executorService.submit(() -> deriveSets(baseEntries, 
setContains, taskIndex, numTaskSlots,
+            numUnionColumns, numAggregationFunctions, aggregationFunctions)));
+      }
+      taskMaps = new ArrayList<>(numTaskSlots);
+      try {
+        for (Future<Map<Key, Record>> future : futures) {
+          taskMaps.add(future.get());
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException("Interrupted while deriving grouping sets", 
e);
+      } catch (ExecutionException e) {
+        for (Future<Map<Key, Record>> future : futures) {
+          future.cancel(true);
+        }
+        throw new RuntimeException("Caught exception while deriving grouping 
sets", e.getCause());
+      }
+    }
+
+    // Union the disjoint task maps into the result table. Since task key 
spaces are disjoint, the union never
+    // merges records. In the common case, adopt the largest task map as the 
table's backing map and putAll the
+    // rest into it, so the largest map's keys are never re-hashed. The 
deterministic accurate-group-by mode
+    // needs a sorted (skip-list) backing map for deterministic iteration 
order, so it keeps the upsert-based
+    // union into the table produced by getTrimDisabledIndexedTable 
(single-threaded, numThreads = 1).
+    IndexedTable derivedTable;
+    if (useDeterministicIndexedTable(queryContext)) {
+      derivedTable = getTrimDisabledIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          initialCapacity, 1, executorService);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        for (Map.Entry<Key, Record> entry : taskMap.entrySet()) {
+          derivedTable.upsert(entry.getKey(), entry.getValue());
+        }
+      }
+    } else {
+      Map<Key, Record> mergedMap = taskMaps.get(0);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap.size() > mergedMap.size()) {
+          mergedMap = taskMap;
+        }
+      }
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap != mergedMap) {
+          mergedMap.putAll(taskMap);
+        }
+      }
+      derivedTable = new SimpleIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          Integer.MAX_VALUE, Integer.MAX_VALUE, mergedMap, executorService);
+    }
+
+    /// Optional server-side per-set trim: when configured (and there is an 
ORDER BY), keep at most K groups
+    /// within each grouping set (bucketed by $groupingId), bounding this 
server's derived output for
+    /// high-cardinality unions. Bucketing per set means a global top-K cannot 
starve a low-magnitude set. This
+    /// is an approximate top-K (deferred exact ORDER BY + LIMIT still runs at 
the broker); unset -> keep all.
+    int serverTrimSize = queryContext.getGroupingSetServerTrimSize();
+    if (serverTrimSize > 0 && derivedTable.size() > (long) serverTrimSize * 
numSets) {
+      derivedTable.finish(false);
+      TableResizer tableResizer = new TableResizer(groupingSetsSchema, 
queryContext);
+      List<IntermediateRecord> kept = 
tableResizer.trimTableByGroupingSet(derivedTable, serverTrimSize,
+          numUnionColumns);
+      
ServerMetrics.get().addMeteredGlobalValue(ServerMeter.AGGREGATE_TIMES_GROUPS_TRIMMED,
 1);
+      return buildIndexedTableFromRecords(groupingSetsSchema, queryContext, 
kept, executorService);
+    }
+    return derivedTable;
+  }
+
+  /// Builds a trim-disabled grouping-set [IndexedTable] pre-populated with 
the given (already unique) records.
+  /// Used to materialize the per-set-trim survivors back into the table the 
combine returns.
+  private static IndexedTable buildIndexedTableFromRecords(DataSchema 
dataSchema, QueryContext queryContext,
+      List<IntermediateRecord> records, ExecutorService executorService) {
+    int numRecords = records.size();
+    int initialCapacity =
+        getIndexedTableInitialCapacity(numRecords, numRecords, 
queryContext.getMinInitialIndexedTableCapacity());
+    // Populated single-threaded, so use a plain-HashMap-backed table 
(numThreads = 1).
+    IndexedTable table = getTrimDisabledIndexedTable(dataSchema, false, 
queryContext, Integer.MAX_VALUE,
+        initialCapacity, 1, executorService);
+    for (IntermediateRecord record : records) {
+      table.upsert(record._key, record._record);
+    }
+    return table;
+  }
+
+  /// One derive task: projects every base group into the grouping sets owned 
by `taskIndex` (set `s` is owned
+  /// when `s % numTaskSlots == taskIndex`) and aggregates the derived groups 
into a task-local map. Each task
+  /// runs single-threaded over its own map, so a plain HashMap suffices; task 
key spaces are disjoint because
+  /// every derived key ends with its set ordinal.
+  ///
+  /// The first record stored for a derived key holds cloned intermediates (it 
becomes the group's accumulator,
+  /// mutated by later merges); on merge, OBJECT intermediates from the shared 
base group are cloned again so the
+  /// merge can never mutate a base accumulator that other tasks are still 
reading. Scalar intermediates are
+  /// immutable and pass through uncloned.
+  private static Map<Key, Record> deriveSets(List<Map.Entry<Key, Record>> 
baseEntries, boolean[][] setContains,
+      int taskIndex, int numTaskSlots, int numUnionColumns, int 
numAggregationFunctions,
+      AggregationFunction[] aggregationFunctions) {
+    int numSets = setContains.length;
+    // A single set can produce at most one derived group per base entry, so 
this covers the largest owned set
+    // (typically the near-identity one) without rehashing; coarser owned sets 
add far fewer distinct groups.
+    Map<Key, Record> taskMap = new 
HashMap<>(HashUtil.getHashMapCapacity(baseEntries.size()));
+    // Flyweight probe: reuse one key buffer for lookups (HashMap.get does not 
retain its argument) and copy it
+    // into a fresh array only when inserting a new derived group. Most 
projections into coarse sets hit an
+    // existing group, so this avoids a key-array + Key allocation per 
projection on the merge path.
+    Object[] probeValues = new Object[numUnionColumns + 1];
+    Key probeKey = new Key(probeValues);
+    for (Map.Entry<Key, Record> baseEntry : baseEntries) {

Review Comment:
   **MAJOR: the derive loop has no termination/accounting checks**
   
   deriveSets does baseGroups × numSets/numTasks work (plus a serde clone for 
OBJECT aggregations) and never calls 
QueryThreadContext.checkTerminationAndSampleUsagePeriodically, unlike the 
segment-merge loop in GroupByCombineOperator. So the query timeout, the OOM 
killer and cancellation cannot stop it or see its allocations. Also, 
future.get() at line 168 has no timeout. On InterruptedException the sibling 
futures are not cancelled; only the ExecutionException branch cancels them.



##########
pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java:
##########
@@ -87,6 +100,230 @@ public static GroupByResultsBlock 
buildGroupingSetsResultsBlock(QueryContext que
     return resultsBlock;
   }
 
+  /// Derives the individual grouping sets from a merged BASE-grouping 
[IndexedTable] (union columns aggregated
+  /// once, like a plain GROUP BY), for a GROUP BY GROUPING SETS / ROLLUP / 
CUBE query using base aggregation.
+  /// Each base group is projected into every grouping set: its rolled-up 
(non-participating) columns are set to
+  /// `null`, the `$groupingId` discriminator is inserted after the union 
columns, and the base group's
+  /// aggregation intermediates are merged into the derived group. This moves 
the per-set fan-out from O(rows) to
+  /// O(base groups) and runs it here -- after the row-collapsing base merge 
-- across the combine's threads,
+  /// rather than expanding every scanned row.
+  ///
+  /// The grouping sets are partitioned round-robin across up to `numTasks` 
worker tasks; each task derives its
+  /// owned sets over all base entries into a task-local plain HashMap. Task 
key spaces are disjoint (every
+  /// derived key carries its set ordinal), so the heavy merge work runs 
contention-free -- no concurrent table,
+  /// no bin locks, no shared size counters -- and the task maps are then 
unioned single-threaded into the result
+  /// table with plain stores (never merges). Profiling showed the previous 
shared-concurrent-table derive spent
+  /// most of its time in ConcurrentHashMap machinery rather than in the 
projection itself.
+  ///
+  /// Clone discipline: a base group's intermediate flows into every grouping 
set (and is read concurrently by
+  /// other tasks), while [AggregationFunction#merge] mutates/returns its 
arguments. Each derived group's stored
+  /// accumulator is therefore a clone (see [#cloneIntermediate]), and OBJECT 
intermediates are re-cloned per
+  /// merge so a merge can never mutate the shared base accumulator another 
task is still reading. Scalar
+  /// intermediates are immutable and clone-free.
+  public static IndexedTable 
deriveGroupingSetsFromMergedBaseTable(IndexedTable baseTable, QueryContext 
queryContext,
+      int numTasks, ExecutorService executorService) {
+    AggregationFunction[] aggregationFunctions = 
queryContext.getAggregationFunctions();
+    assert aggregationFunctions != null;
+    int numAggregationFunctions = aggregationFunctions.length;
+    List<int[]> groupingSets = queryContext.getGroupingSets();
+    int numSets = groupingSets.size();
+    int numUnionColumns = queryContext.getGroupByExpressions().size();
+    // Per grouping set: membership mask over the union columns (true = 
participates, false = rolled up to NULL).
+    boolean[][] setContains = new boolean[numSets][numUnionColumns];
+    for (int s = 0; s < numSets; s++) {
+      for (int columnIndex : groupingSets.get(s)) {
+        setContains[s][columnIndex] = true;
+      }
+    }
+
+    // Grouping-set output schema: the base schema with the synthetic 
$groupingId INT column inserted right after
+    // the union group-by columns (mirroring GroupByOperator's grouping-set 
schema layout).
+    DataSchema groupingSetsSchema = 
insertGroupingIdColumn(baseTable.getDataSchema(), numUnionColumns);
+    // The derive must not drop groups: it is a bounded transformation of the 
already-bounded base groups (whose
+    // count is capped at numGroupsLimit per segment), so the derived count is 
at most numGroupsLimit * numSets --
+    // a finite amount. Capping the derived table at numGroupsLimit would drop 
derived groups NON-DETERMINISTICALLY
+    // under the parallel upsert (whichever threads fill the quota first win), 
which can starve an entire
+    // low-magnitude grouping set such as the grand total. Use an unbounded 
result size here and defer the real
+    // ORDER BY + LIMIT to the broker (per-set trim below still bounds it when 
explicitly configured).
+    int derivedUpperBound = (int) Math.min((long) baseTable.size() * numSets, 
Integer.MAX_VALUE);

Review Comment:
   **MAJOR risk: server output grows to base × numSets with no cap**
   
   The derived table is unbounded (resultSize MAX_VALUE). Before, the combine 
table for grouping sets was capped at numGroupsLimit TOTAL groups. Now each 
server can send up to numGroupsLimit × numSets rows (CUBE over 5 columns = 32 
sets → 3.2M rows at the default 100k limit). The base table, the baseEntries 
copy and all task maps are alive together, so peak heap is about (2 + numSets) 
× base. Broker reduce and DataTable size grow in the same way. Because this is 
on by default, I would want either a cap that is safe per set, or the feature 
off by default for the first release.



##########
pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java:
##########
@@ -87,6 +100,230 @@ public static GroupByResultsBlock 
buildGroupingSetsResultsBlock(QueryContext que
     return resultsBlock;
   }
 
+  /// Derives the individual grouping sets from a merged BASE-grouping 
[IndexedTable] (union columns aggregated
+  /// once, like a plain GROUP BY), for a GROUP BY GROUPING SETS / ROLLUP / 
CUBE query using base aggregation.
+  /// Each base group is projected into every grouping set: its rolled-up 
(non-participating) columns are set to
+  /// `null`, the `$groupingId` discriminator is inserted after the union 
columns, and the base group's
+  /// aggregation intermediates are merged into the derived group. This moves 
the per-set fan-out from O(rows) to
+  /// O(base groups) and runs it here -- after the row-collapsing base merge 
-- across the combine's threads,
+  /// rather than expanding every scanned row.
+  ///
+  /// The grouping sets are partitioned round-robin across up to `numTasks` 
worker tasks; each task derives its
+  /// owned sets over all base entries into a task-local plain HashMap. Task 
key spaces are disjoint (every
+  /// derived key carries its set ordinal), so the heavy merge work runs 
contention-free -- no concurrent table,
+  /// no bin locks, no shared size counters -- and the task maps are then 
unioned single-threaded into the result
+  /// table with plain stores (never merges). Profiling showed the previous 
shared-concurrent-table derive spent
+  /// most of its time in ConcurrentHashMap machinery rather than in the 
projection itself.
+  ///
+  /// Clone discipline: a base group's intermediate flows into every grouping 
set (and is read concurrently by
+  /// other tasks), while [AggregationFunction#merge] mutates/returns its 
arguments. Each derived group's stored
+  /// accumulator is therefore a clone (see [#cloneIntermediate]), and OBJECT 
intermediates are re-cloned per
+  /// merge so a merge can never mutate the shared base accumulator another 
task is still reading. Scalar
+  /// intermediates are immutable and clone-free.
+  public static IndexedTable 
deriveGroupingSetsFromMergedBaseTable(IndexedTable baseTable, QueryContext 
queryContext,
+      int numTasks, ExecutorService executorService) {
+    AggregationFunction[] aggregationFunctions = 
queryContext.getAggregationFunctions();
+    assert aggregationFunctions != null;
+    int numAggregationFunctions = aggregationFunctions.length;
+    List<int[]> groupingSets = queryContext.getGroupingSets();
+    int numSets = groupingSets.size();
+    int numUnionColumns = queryContext.getGroupByExpressions().size();
+    // Per grouping set: membership mask over the union columns (true = 
participates, false = rolled up to NULL).
+    boolean[][] setContains = new boolean[numSets][numUnionColumns];
+    for (int s = 0; s < numSets; s++) {
+      for (int columnIndex : groupingSets.get(s)) {
+        setContains[s][columnIndex] = true;
+      }
+    }
+
+    // Grouping-set output schema: the base schema with the synthetic 
$groupingId INT column inserted right after
+    // the union group-by columns (mirroring GroupByOperator's grouping-set 
schema layout).
+    DataSchema groupingSetsSchema = 
insertGroupingIdColumn(baseTable.getDataSchema(), numUnionColumns);
+    // The derive must not drop groups: it is a bounded transformation of the 
already-bounded base groups (whose
+    // count is capped at numGroupsLimit per segment), so the derived count is 
at most numGroupsLimit * numSets --
+    // a finite amount. Capping the derived table at numGroupsLimit would drop 
derived groups NON-DETERMINISTICALLY
+    // under the parallel upsert (whichever threads fill the quota first win), 
which can starve an entire
+    // low-magnitude grouping set such as the grand total. Use an unbounded 
result size here and defer the real
+    // ORDER BY + LIMIT to the broker (per-set trim below still bounds it when 
explicitly configured).
+    int derivedUpperBound = (int) Math.min((long) baseTable.size() * numSets, 
Integer.MAX_VALUE);
+    int initialCapacity = getIndexedTableInitialCapacity(derivedUpperBound, 
derivedUpperBound,
+        queryContext.getMinInitialIndexedTableCapacity());
+
+    List<Map.Entry<Key, Record>> baseEntries = new 
ArrayList<>(baseTable.getRecordEntries());
+    int numTaskSlots = Math.max(1, Math.min(numTasks, numSets));
+    List<Map<Key, Record>> taskMaps;
+    if (numTaskSlots == 1) {
+      taskMaps = List.of(deriveSets(baseEntries, setContains, 0, 1, 
numUnionColumns, numAggregationFunctions,
+          aggregationFunctions));
+    } else {
+      List<Future<Map<Key, Record>>> futures = new ArrayList<>(numTaskSlots);
+      for (int t = 0; t < numTaskSlots; t++) {
+        int taskIndex = t;
+        futures.add(executorService.submit(() -> deriveSets(baseEntries, 
setContains, taskIndex, numTaskSlots,
+            numUnionColumns, numAggregationFunctions, aggregationFunctions)));
+      }
+      taskMaps = new ArrayList<>(numTaskSlots);
+      try {
+        for (Future<Map<Key, Record>> future : futures) {
+          taskMaps.add(future.get());
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException("Interrupted while deriving grouping sets", 
e);
+      } catch (ExecutionException e) {
+        for (Future<Map<Key, Record>> future : futures) {
+          future.cancel(true);
+        }
+        throw new RuntimeException("Caught exception while deriving grouping 
sets", e.getCause());
+      }
+    }
+
+    // Union the disjoint task maps into the result table. Since task key 
spaces are disjoint, the union never
+    // merges records. In the common case, adopt the largest task map as the 
table's backing map and putAll the
+    // rest into it, so the largest map's keys are never re-hashed. The 
deterministic accurate-group-by mode
+    // needs a sorted (skip-list) backing map for deterministic iteration 
order, so it keeps the upsert-based
+    // union into the table produced by getTrimDisabledIndexedTable 
(single-threaded, numThreads = 1).
+    IndexedTable derivedTable;
+    if (useDeterministicIndexedTable(queryContext)) {
+      derivedTable = getTrimDisabledIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          initialCapacity, 1, executorService);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        for (Map.Entry<Key, Record> entry : taskMap.entrySet()) {
+          derivedTable.upsert(entry.getKey(), entry.getValue());
+        }
+      }
+    } else {
+      Map<Key, Record> mergedMap = taskMaps.get(0);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap.size() > mergedMap.size()) {
+          mergedMap = taskMap;
+        }
+      }
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap != mergedMap) {
+          mergedMap.putAll(taskMap);
+        }
+      }
+      derivedTable = new SimpleIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          Integer.MAX_VALUE, Integer.MAX_VALUE, mergedMap, executorService);
+    }
+
+    /// Optional server-side per-set trim: when configured (and there is an 
ORDER BY), keep at most K groups
+    /// within each grouping set (bucketed by $groupingId), bounding this 
server's derived output for
+    /// high-cardinality unions. Bucketing per set means a global top-K cannot 
starve a low-magnitude set. This
+    /// is an approximate top-K (deferred exact ORDER BY + LIMIT still runs at 
the broker); unset -> keep all.
+    int serverTrimSize = queryContext.getGroupingSetServerTrimSize();
+    if (serverTrimSize > 0 && derivedTable.size() > (long) serverTrimSize * 
numSets) {
+      derivedTable.finish(false);
+      TableResizer tableResizer = new TableResizer(groupingSetsSchema, 
queryContext);
+      List<IntermediateRecord> kept = 
tableResizer.trimTableByGroupingSet(derivedTable, serverTrimSize,
+          numUnionColumns);
+      
ServerMetrics.get().addMeteredGlobalValue(ServerMeter.AGGREGATE_TIMES_GROUPS_TRIMMED,
 1);
+      return buildIndexedTableFromRecords(groupingSetsSchema, queryContext, 
kept, executorService);
+    }
+    return derivedTable;
+  }
+
+  /// Builds a trim-disabled grouping-set [IndexedTable] pre-populated with 
the given (already unique) records.
+  /// Used to materialize the per-set-trim survivors back into the table the 
combine returns.
+  private static IndexedTable buildIndexedTableFromRecords(DataSchema 
dataSchema, QueryContext queryContext,
+      List<IntermediateRecord> records, ExecutorService executorService) {
+    int numRecords = records.size();
+    int initialCapacity =
+        getIndexedTableInitialCapacity(numRecords, numRecords, 
queryContext.getMinInitialIndexedTableCapacity());
+    // Populated single-threaded, so use a plain-HashMap-backed table 
(numThreads = 1).
+    IndexedTable table = getTrimDisabledIndexedTable(dataSchema, false, 
queryContext, Integer.MAX_VALUE,
+        initialCapacity, 1, executorService);
+    for (IntermediateRecord record : records) {
+      table.upsert(record._key, record._record);
+    }
+    return table;
+  }
+
+  /// One derive task: projects every base group into the grouping sets owned 
by `taskIndex` (set `s` is owned
+  /// when `s % numTaskSlots == taskIndex`) and aggregates the derived groups 
into a task-local map. Each task
+  /// runs single-threaded over its own map, so a plain HashMap suffices; task 
key spaces are disjoint because
+  /// every derived key ends with its set ordinal.
+  ///
+  /// The first record stored for a derived key holds cloned intermediates (it 
becomes the group's accumulator,
+  /// mutated by later merges); on merge, OBJECT intermediates from the shared 
base group are cloned again so the
+  /// merge can never mutate a base accumulator that other tasks are still 
reading. Scalar intermediates are
+  /// immutable and pass through uncloned.
+  private static Map<Key, Record> deriveSets(List<Map.Entry<Key, Record>> 
baseEntries, boolean[][] setContains,
+      int taskIndex, int numTaskSlots, int numUnionColumns, int 
numAggregationFunctions,
+      AggregationFunction[] aggregationFunctions) {
+    int numSets = setContains.length;
+    // A single set can produce at most one derived group per base entry, so 
this covers the largest owned set
+    // (typically the near-identity one) without rehashing; coarser owned sets 
add far fewer distinct groups.
+    Map<Key, Record> taskMap = new 
HashMap<>(HashUtil.getHashMapCapacity(baseEntries.size()));

Review Comment:
   **perf: every task map is sized for all base entries**
   
   Each of the numTaskSlots tasks pre-sizes its HashMap to baseEntries.size(), 
even when it owns only coarse sets (grand total = 1 group). That allocates 
numTaskSlots × base-size bucket arrays for nothing. Size it from the largest 
owned set instead (e.g. the base size only if the task owns a set that contains 
every union column, otherwise a smaller estimate). Also, partitioning by SET 
limits parallelism to numSets and is very unbalanced: for ROLLUP(a,b) one task 
does the near-identity set and the others finish at once.



##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByExecutor.java:
##########
@@ -51,4 +51,13 @@ public interface GroupByExecutor {
   GroupKeyGenerator getGroupKeyGenerator();
 
   GroupByResultHolder[] getGroupByResultHolders();
+
+  /// For a grouping-set query, whether this executor aggregated only the BASE 
(union) grouping so the caller
+  /// must derive the individual grouping-set records from the base groups (see
+  /// 
[org.apache.pinot.core.util.GroupByUtils#buildGroupingSetsResultsBlockFromBaseGroups]),
 rather than the

Review Comment:
   **Stale javadoc link**
   
   GroupByUtils#buildGroupingSetsResultsBlockFromBaseGroups does not exist. The 
derive is GroupByUtils#deriveGroupingSetsFromMergedBaseTable, and it runs in 
the combine, not in the executor's caller.



##########
pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java:
##########
@@ -84,10 +84,14 @@ protected IndexedTable(DataSchema dataSchema, boolean 
hasFinalInput, QueryContex
 
     List<ExpressionContext> groupByExpressions = 
queryContext.getGroupByExpressions();
     assert groupByExpressions != null;
-    /// Includes the synthetic $groupingId key column for GROUP BY GROUPING 
SETS / ROLLUP / CUBE queries, so
-    /// that rows from different grouping sets are keyed (and therefore 
merged) independently.
-    _numKeyColumns = queryContext.getNumGroupByKeyColumns();
     _aggregationFunctions = queryContext.getAggregationFunctions();
+    /// Derive the key-column count from the schema (total columns minus 
aggregation columns) rather than the
+    /// query context, so it matches the actual record layout. This normally 
equals
+    /// `queryContext.getNumGroupByKeyColumns()` -- which includes the 
synthetic $groupingId column for GROUP BY
+    /// GROUPING SETS / ROLLUP / CUBE so rows from different grouping sets are 
keyed (and merged) independently --
+    /// but for a base-aggregation grouping-set combine table the records are 
BASE groups WITHOUT the $groupingId
+    /// column, so the key count is the union-column count. Deriving from the 
schema keeps both layouts correct.
+    _numKeyColumns = dataSchema.size() - _aggregationFunctions.length;

Review Comment:
   **Two sources of truth for the key-column count**
   
   IndexedTable now gets _numKeyColumns from the schema, but the TableResizer 
it builds 3 lines below still uses queryContext.getNumGroupByKeyColumns() 
(union+1). The base combine table with an ORDER BY therefore has a resizer that 
reads aggregations one column too far to the right. It does no harm today only 
because trim is disabled for grouping sets and the base table is never 
finish(sort)ed. This is a trap for a later change. Pass the key count into 
TableResizer, or do not build a resizer for the base table.



##########
pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java:
##########
@@ -422,6 +422,67 @@ public List<IntermediateRecord> 
trimInSegmentResultsByGroupingSet(GroupKeyGenera
     return result;
   }
 
+  /// Per-set bucketed trim over already-built [IntermediateRecord]s (the 
base-aggregation derive path produces
+  /// these directly rather than via a [GroupKeyGenerator]). Keeps the top 
`perSetSize` records WITHIN each
+  /// grouping set, bucketed by the `discriminatorColumnIndex` ($groupingId) 
value, with the same across-set
+  /// anti-starvation guarantee and approximation caveats as 
[#trimInSegmentResultsByGroupingSet].
+  public List<IntermediateRecord> 
trimInSegmentRecordsByGroupingSet(Collection<IntermediateRecord> records,

Review Comment:
   **Dead code**
   
   trimInSegmentRecordsByGroupingSet has no callers in the PR. It is an older 
iteration of the server trim. Remove it (and its nearly identical twin 
trimTableByGroupingSet could share the heap logic).



##########
pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java:
##########
@@ -87,6 +100,230 @@ public static GroupByResultsBlock 
buildGroupingSetsResultsBlock(QueryContext que
     return resultsBlock;
   }
 
+  /// Derives the individual grouping sets from a merged BASE-grouping 
[IndexedTable] (union columns aggregated
+  /// once, like a plain GROUP BY), for a GROUP BY GROUPING SETS / ROLLUP / 
CUBE query using base aggregation.
+  /// Each base group is projected into every grouping set: its rolled-up 
(non-participating) columns are set to
+  /// `null`, the `$groupingId` discriminator is inserted after the union 
columns, and the base group's
+  /// aggregation intermediates are merged into the derived group. This moves 
the per-set fan-out from O(rows) to
+  /// O(base groups) and runs it here -- after the row-collapsing base merge 
-- across the combine's threads,
+  /// rather than expanding every scanned row.
+  ///
+  /// The grouping sets are partitioned round-robin across up to `numTasks` 
worker tasks; each task derives its
+  /// owned sets over all base entries into a task-local plain HashMap. Task 
key spaces are disjoint (every
+  /// derived key carries its set ordinal), so the heavy merge work runs 
contention-free -- no concurrent table,
+  /// no bin locks, no shared size counters -- and the task maps are then 
unioned single-threaded into the result
+  /// table with plain stores (never merges). Profiling showed the previous 
shared-concurrent-table derive spent
+  /// most of its time in ConcurrentHashMap machinery rather than in the 
projection itself.
+  ///
+  /// Clone discipline: a base group's intermediate flows into every grouping 
set (and is read concurrently by
+  /// other tasks), while [AggregationFunction#merge] mutates/returns its 
arguments. Each derived group's stored
+  /// accumulator is therefore a clone (see [#cloneIntermediate]), and OBJECT 
intermediates are re-cloned per
+  /// merge so a merge can never mutate the shared base accumulator another 
task is still reading. Scalar
+  /// intermediates are immutable and clone-free.
+  public static IndexedTable 
deriveGroupingSetsFromMergedBaseTable(IndexedTable baseTable, QueryContext 
queryContext,
+      int numTasks, ExecutorService executorService) {
+    AggregationFunction[] aggregationFunctions = 
queryContext.getAggregationFunctions();
+    assert aggregationFunctions != null;
+    int numAggregationFunctions = aggregationFunctions.length;
+    List<int[]> groupingSets = queryContext.getGroupingSets();
+    int numSets = groupingSets.size();
+    int numUnionColumns = queryContext.getGroupByExpressions().size();
+    // Per grouping set: membership mask over the union columns (true = 
participates, false = rolled up to NULL).
+    boolean[][] setContains = new boolean[numSets][numUnionColumns];
+    for (int s = 0; s < numSets; s++) {
+      for (int columnIndex : groupingSets.get(s)) {
+        setContains[s][columnIndex] = true;
+      }
+    }
+
+    // Grouping-set output schema: the base schema with the synthetic 
$groupingId INT column inserted right after
+    // the union group-by columns (mirroring GroupByOperator's grouping-set 
schema layout).
+    DataSchema groupingSetsSchema = 
insertGroupingIdColumn(baseTable.getDataSchema(), numUnionColumns);
+    // The derive must not drop groups: it is a bounded transformation of the 
already-bounded base groups (whose
+    // count is capped at numGroupsLimit per segment), so the derived count is 
at most numGroupsLimit * numSets --
+    // a finite amount. Capping the derived table at numGroupsLimit would drop 
derived groups NON-DETERMINISTICALLY
+    // under the parallel upsert (whichever threads fill the quota first win), 
which can starve an entire
+    // low-magnitude grouping set such as the grand total. Use an unbounded 
result size here and defer the real
+    // ORDER BY + LIMIT to the broker (per-set trim below still bounds it when 
explicitly configured).
+    int derivedUpperBound = (int) Math.min((long) baseTable.size() * numSets, 
Integer.MAX_VALUE);
+    int initialCapacity = getIndexedTableInitialCapacity(derivedUpperBound, 
derivedUpperBound,
+        queryContext.getMinInitialIndexedTableCapacity());
+
+    List<Map.Entry<Key, Record>> baseEntries = new 
ArrayList<>(baseTable.getRecordEntries());
+    int numTaskSlots = Math.max(1, Math.min(numTasks, numSets));
+    List<Map<Key, Record>> taskMaps;
+    if (numTaskSlots == 1) {
+      taskMaps = List.of(deriveSets(baseEntries, setContains, 0, 1, 
numUnionColumns, numAggregationFunctions,
+          aggregationFunctions));
+    } else {
+      List<Future<Map<Key, Record>>> futures = new ArrayList<>(numTaskSlots);
+      for (int t = 0; t < numTaskSlots; t++) {
+        int taskIndex = t;
+        futures.add(executorService.submit(() -> deriveSets(baseEntries, 
setContains, taskIndex, numTaskSlots,
+            numUnionColumns, numAggregationFunctions, aggregationFunctions)));
+      }
+      taskMaps = new ArrayList<>(numTaskSlots);
+      try {
+        for (Future<Map<Key, Record>> future : futures) {
+          taskMaps.add(future.get());
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException("Interrupted while deriving grouping sets", 
e);
+      } catch (ExecutionException e) {
+        for (Future<Map<Key, Record>> future : futures) {
+          future.cancel(true);
+        }
+        throw new RuntimeException("Caught exception while deriving grouping 
sets", e.getCause());
+      }
+    }
+
+    // Union the disjoint task maps into the result table. Since task key 
spaces are disjoint, the union never
+    // merges records. In the common case, adopt the largest task map as the 
table's backing map and putAll the
+    // rest into it, so the largest map's keys are never re-hashed. The 
deterministic accurate-group-by mode
+    // needs a sorted (skip-list) backing map for deterministic iteration 
order, so it keeps the upsert-based
+    // union into the table produced by getTrimDisabledIndexedTable 
(single-threaded, numThreads = 1).
+    IndexedTable derivedTable;
+    if (useDeterministicIndexedTable(queryContext)) {
+      derivedTable = getTrimDisabledIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          initialCapacity, 1, executorService);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        for (Map.Entry<Key, Record> entry : taskMap.entrySet()) {
+          derivedTable.upsert(entry.getKey(), entry.getValue());
+        }
+      }
+    } else {
+      Map<Key, Record> mergedMap = taskMaps.get(0);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap.size() > mergedMap.size()) {
+          mergedMap = taskMap;
+        }
+      }
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap != mergedMap) {
+          mergedMap.putAll(taskMap);
+        }
+      }
+      derivedTable = new SimpleIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          Integer.MAX_VALUE, Integer.MAX_VALUE, mergedMap, executorService);
+    }
+
+    /// Optional server-side per-set trim: when configured (and there is an 
ORDER BY), keep at most K groups
+    /// within each grouping set (bucketed by $groupingId), bounding this 
server's derived output for
+    /// high-cardinality unions. Bucketing per set means a global top-K cannot 
starve a low-magnitude set. This
+    /// is an approximate top-K (deferred exact ORDER BY + LIMIT still runs at 
the broker); unset -> keep all.
+    int serverTrimSize = queryContext.getGroupingSetServerTrimSize();
+    if (serverTrimSize > 0 && derivedTable.size() > (long) serverTrimSize * 
numSets) {

Review Comment:
   **groupingSetsServerTrimSize: trim does not set groupsTrimmed, and the name 
does not match the behaviour**
   
   (1) When this trim drops groups, the result is approximate, but mergeResults 
never calls setGroupsTrimmed(true), so the broker and the client do not know. 
(2) QueryContext converts the option with getTableCapacity(limit, value) = 
max(5*limit, value). The value is a minimum, not the 'keep at most this many' 
limit that the CommonConstants javadoc describes. Name it 
groupingSetsMinServerTrimSize, or change the doc. (3) Integer.parseInt at 
QueryContext:622 is not validated, so a bad value fails late inside combine 
with a generic error.



##########
pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java:
##########
@@ -167,6 +193,18 @@ protected GroupByResultsBlock getNextBlock() {
     /// the final ORDER BY + LIMIT across all sets.
     if (_queryContext.isGroupingSets()) {
       /// The $groupingId discriminator is the key column immediately after 
the union group-by columns.
+      if (groupByExecutor.isGroupingSetsBaseAggregation()) {

Review Comment:
   **Base path silently bypasses the per-set segment trim**
   
   When the base path is used, minSegmentGroupTrimSize / 
getGroupingSetSegmentTrimSize no longer apply, because base groups cannot be 
trimmed safely. This is correct, but it is a behaviour change for clusters that 
set a segment trim. Say so in the PR description and add a test with ORDER BY + 
minSegmentGroupTrimSize.



##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:
##########
@@ -252,6 +253,26 @@ public int getNumExtraGroupByKeyColumns() {
     return isGroupingSets() ? 1 : 0;
   }
 
+  /// Returns whether grouping-set queries should aggregate the base grouping 
(union columns) once per segment
+  /// and derive the individual grouping-set records from those base groups, 
instead of expanding every input
+  /// row into one group per grouping set. Enabled by default; the 
`groupingSetsBaseAggregation=false` query
+  /// option forces the legacy per-row expansion path. See 
[CommonConstants.Broker.Request.QueryOptionKey].
+  ///
+  /// Only the standard 
[org.apache.pinot.core.operator.combine.GroupByCombineOperator] performs the 
merge-time
+  /// derive that turns the emitted base groups back into grouping-set 
records. Base aggregation is therefore
+  /// disabled whenever a different combine path would run: filtered 
aggregations (distinct shared group-key
+  /// generator), the streaming combine used on MSE leaf stages 
(`streamingGroupByFlushThreshold > 0`, which
+  /// flushes incrementally and cannot derive over the fully-merged base 
table), and the sorted-aggregate combine
+  /// ([#shouldSortAggregateUnderSafeTrim()]). In all these cases the segment 
expands per row so the emitted rows
+  /// already carry the $groupingId discriminator.
+  public boolean isGroupingSetsBaseAggregation() {

Review Comment:
   **Default-on + the list of exclusions lives in 2-3 places**
   
   The decision is spread over QueryContext (filtered aggs, streaming, 
sort-aggregate), DefaultGroupByExecutor (MV) and the combine (schema width). 
The exclusions are needed because the other combine operators cannot derive. If 
someone adds a new combine path later, it will silently receive base blocks. 
Consider one query-level resolver, plus a check in the other combine operators 
that rejects blocks without $groupingId. The good news: the server output 
layout is the same on both paths, so mixed-version brokers are not affected.



##########
pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java:
##########
@@ -87,6 +100,230 @@ public static GroupByResultsBlock 
buildGroupingSetsResultsBlock(QueryContext que
     return resultsBlock;
   }
 
+  /// Derives the individual grouping sets from a merged BASE-grouping 
[IndexedTable] (union columns aggregated
+  /// once, like a plain GROUP BY), for a GROUP BY GROUPING SETS / ROLLUP / 
CUBE query using base aggregation.
+  /// Each base group is projected into every grouping set: its rolled-up 
(non-participating) columns are set to
+  /// `null`, the `$groupingId` discriminator is inserted after the union 
columns, and the base group's
+  /// aggregation intermediates are merged into the derived group. This moves 
the per-set fan-out from O(rows) to
+  /// O(base groups) and runs it here -- after the row-collapsing base merge 
-- across the combine's threads,
+  /// rather than expanding every scanned row.
+  ///
+  /// The grouping sets are partitioned round-robin across up to `numTasks` 
worker tasks; each task derives its
+  /// owned sets over all base entries into a task-local plain HashMap. Task 
key spaces are disjoint (every
+  /// derived key carries its set ordinal), so the heavy merge work runs 
contention-free -- no concurrent table,
+  /// no bin locks, no shared size counters -- and the task maps are then 
unioned single-threaded into the result
+  /// table with plain stores (never merges). Profiling showed the previous 
shared-concurrent-table derive spent
+  /// most of its time in ConcurrentHashMap machinery rather than in the 
projection itself.
+  ///
+  /// Clone discipline: a base group's intermediate flows into every grouping 
set (and is read concurrently by
+  /// other tasks), while [AggregationFunction#merge] mutates/returns its 
arguments. Each derived group's stored
+  /// accumulator is therefore a clone (see [#cloneIntermediate]), and OBJECT 
intermediates are re-cloned per
+  /// merge so a merge can never mutate the shared base accumulator another 
task is still reading. Scalar
+  /// intermediates are immutable and clone-free.
+  public static IndexedTable 
deriveGroupingSetsFromMergedBaseTable(IndexedTable baseTable, QueryContext 
queryContext,
+      int numTasks, ExecutorService executorService) {
+    AggregationFunction[] aggregationFunctions = 
queryContext.getAggregationFunctions();
+    assert aggregationFunctions != null;
+    int numAggregationFunctions = aggregationFunctions.length;
+    List<int[]> groupingSets = queryContext.getGroupingSets();
+    int numSets = groupingSets.size();
+    int numUnionColumns = queryContext.getGroupByExpressions().size();
+    // Per grouping set: membership mask over the union columns (true = 
participates, false = rolled up to NULL).
+    boolean[][] setContains = new boolean[numSets][numUnionColumns];
+    for (int s = 0; s < numSets; s++) {
+      for (int columnIndex : groupingSets.get(s)) {
+        setContains[s][columnIndex] = true;
+      }
+    }
+
+    // Grouping-set output schema: the base schema with the synthetic 
$groupingId INT column inserted right after
+    // the union group-by columns (mirroring GroupByOperator's grouping-set 
schema layout).
+    DataSchema groupingSetsSchema = 
insertGroupingIdColumn(baseTable.getDataSchema(), numUnionColumns);
+    // The derive must not drop groups: it is a bounded transformation of the 
already-bounded base groups (whose
+    // count is capped at numGroupsLimit per segment), so the derived count is 
at most numGroupsLimit * numSets --
+    // a finite amount. Capping the derived table at numGroupsLimit would drop 
derived groups NON-DETERMINISTICALLY
+    // under the parallel upsert (whichever threads fill the quota first win), 
which can starve an entire
+    // low-magnitude grouping set such as the grand total. Use an unbounded 
result size here and defer the real
+    // ORDER BY + LIMIT to the broker (per-set trim below still bounds it when 
explicitly configured).
+    int derivedUpperBound = (int) Math.min((long) baseTable.size() * numSets, 
Integer.MAX_VALUE);
+    int initialCapacity = getIndexedTableInitialCapacity(derivedUpperBound, 
derivedUpperBound,
+        queryContext.getMinInitialIndexedTableCapacity());
+
+    List<Map.Entry<Key, Record>> baseEntries = new 
ArrayList<>(baseTable.getRecordEntries());
+    int numTaskSlots = Math.max(1, Math.min(numTasks, numSets));
+    List<Map<Key, Record>> taskMaps;
+    if (numTaskSlots == 1) {
+      taskMaps = List.of(deriveSets(baseEntries, setContains, 0, 1, 
numUnionColumns, numAggregationFunctions,
+          aggregationFunctions));
+    } else {
+      List<Future<Map<Key, Record>>> futures = new ArrayList<>(numTaskSlots);
+      for (int t = 0; t < numTaskSlots; t++) {
+        int taskIndex = t;
+        futures.add(executorService.submit(() -> deriveSets(baseEntries, 
setContains, taskIndex, numTaskSlots,
+            numUnionColumns, numAggregationFunctions, aggregationFunctions)));
+      }
+      taskMaps = new ArrayList<>(numTaskSlots);
+      try {
+        for (Future<Map<Key, Record>> future : futures) {
+          taskMaps.add(future.get());
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException("Interrupted while deriving grouping sets", 
e);
+      } catch (ExecutionException e) {
+        for (Future<Map<Key, Record>> future : futures) {
+          future.cancel(true);
+        }
+        throw new RuntimeException("Caught exception while deriving grouping 
sets", e.getCause());
+      }
+    }
+
+    // Union the disjoint task maps into the result table. Since task key 
spaces are disjoint, the union never
+    // merges records. In the common case, adopt the largest task map as the 
table's backing map and putAll the
+    // rest into it, so the largest map's keys are never re-hashed. The 
deterministic accurate-group-by mode
+    // needs a sorted (skip-list) backing map for deterministic iteration 
order, so it keeps the upsert-based
+    // union into the table produced by getTrimDisabledIndexedTable 
(single-threaded, numThreads = 1).
+    IndexedTable derivedTable;
+    if (useDeterministicIndexedTable(queryContext)) {
+      derivedTable = getTrimDisabledIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          initialCapacity, 1, executorService);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        for (Map.Entry<Key, Record> entry : taskMap.entrySet()) {
+          derivedTable.upsert(entry.getKey(), entry.getValue());
+        }
+      }
+    } else {
+      Map<Key, Record> mergedMap = taskMaps.get(0);
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap.size() > mergedMap.size()) {
+          mergedMap = taskMap;
+        }
+      }
+      for (Map<Key, Record> taskMap : taskMaps) {
+        if (taskMap != mergedMap) {
+          mergedMap.putAll(taskMap);
+        }
+      }
+      derivedTable = new SimpleIndexedTable(groupingSetsSchema, false, 
queryContext, Integer.MAX_VALUE,
+          Integer.MAX_VALUE, Integer.MAX_VALUE, mergedMap, executorService);
+    }
+
+    /// Optional server-side per-set trim: when configured (and there is an 
ORDER BY), keep at most K groups
+    /// within each grouping set (bucketed by $groupingId), bounding this 
server's derived output for
+    /// high-cardinality unions. Bucketing per set means a global top-K cannot 
starve a low-magnitude set. This
+    /// is an approximate top-K (deferred exact ORDER BY + LIMIT still runs at 
the broker); unset -> keep all.
+    int serverTrimSize = queryContext.getGroupingSetServerTrimSize();
+    if (serverTrimSize > 0 && derivedTable.size() > (long) serverTrimSize * 
numSets) {
+      derivedTable.finish(false);
+      TableResizer tableResizer = new TableResizer(groupingSetsSchema, 
queryContext);
+      List<IntermediateRecord> kept = 
tableResizer.trimTableByGroupingSet(derivedTable, serverTrimSize,
+          numUnionColumns);
+      
ServerMetrics.get().addMeteredGlobalValue(ServerMeter.AGGREGATE_TIMES_GROUPS_TRIMMED,
 1);
+      return buildIndexedTableFromRecords(groupingSetsSchema, queryContext, 
kept, executorService);
+    }
+    return derivedTable;
+  }
+
+  /// Builds a trim-disabled grouping-set [IndexedTable] pre-populated with 
the given (already unique) records.
+  /// Used to materialize the per-set-trim survivors back into the table the 
combine returns.
+  private static IndexedTable buildIndexedTableFromRecords(DataSchema 
dataSchema, QueryContext queryContext,
+      List<IntermediateRecord> records, ExecutorService executorService) {
+    int numRecords = records.size();
+    int initialCapacity =
+        getIndexedTableInitialCapacity(numRecords, numRecords, 
queryContext.getMinInitialIndexedTableCapacity());
+    // Populated single-threaded, so use a plain-HashMap-backed table 
(numThreads = 1).
+    IndexedTable table = getTrimDisabledIndexedTable(dataSchema, false, 
queryContext, Integer.MAX_VALUE,
+        initialCapacity, 1, executorService);
+    for (IntermediateRecord record : records) {
+      table.upsert(record._key, record._record);
+    }
+    return table;
+  }
+
+  /// One derive task: projects every base group into the grouping sets owned 
by `taskIndex` (set `s` is owned
+  /// when `s % numTaskSlots == taskIndex`) and aggregates the derived groups 
into a task-local map. Each task
+  /// runs single-threaded over its own map, so a plain HashMap suffices; task 
key spaces are disjoint because
+  /// every derived key ends with its set ordinal.
+  ///
+  /// The first record stored for a derived key holds cloned intermediates (it 
becomes the group's accumulator,
+  /// mutated by later merges); on merge, OBJECT intermediates from the shared 
base group are cloned again so the
+  /// merge can never mutate a base accumulator that other tasks are still 
reading. Scalar intermediates are
+  /// immutable and pass through uncloned.
+  private static Map<Key, Record> deriveSets(List<Map.Entry<Key, Record>> 
baseEntries, boolean[][] setContains,
+      int taskIndex, int numTaskSlots, int numUnionColumns, int 
numAggregationFunctions,
+      AggregationFunction[] aggregationFunctions) {
+    int numSets = setContains.length;
+    // A single set can produce at most one derived group per base entry, so 
this covers the largest owned set
+    // (typically the near-identity one) without rehashing; coarser owned sets 
add far fewer distinct groups.
+    Map<Key, Record> taskMap = new 
HashMap<>(HashUtil.getHashMapCapacity(baseEntries.size()));
+    // Flyweight probe: reuse one key buffer for lookups (HashMap.get does not 
retain its argument) and copy it
+    // into a fresh array only when inserting a new derived group. Most 
projections into coarse sets hit an
+    // existing group, so this avoids a key-array + Key allocation per 
projection on the merge path.
+    Object[] probeValues = new Object[numUnionColumns + 1];
+    Key probeKey = new Key(probeValues);
+    for (Map.Entry<Key, Record> baseEntry : baseEntries) {
+      Object[] baseKeys = baseEntry.getKey().getValues();
+      Object[] baseValues = baseEntry.getValue().getValues();
+      for (int s = taskIndex; s < numSets; s += numTaskSlots) {
+        boolean[] contains = setContains[s];
+        for (int col = 0; col < numUnionColumns; col++) {
+          probeValues[col] = contains[col] ? baseKeys[col] : null;
+        }
+        probeValues[numUnionColumns] = s;
+        Record existing = taskMap.get(probeKey);
+        if (existing == null) {
+          Object[] keyValues = probeValues.clone();
+          Object[] values = new Object[numUnionColumns + 1 + 
numAggregationFunctions];
+          System.arraycopy(keyValues, 0, values, 0, numUnionColumns + 1);
+          for (int i = 0; i < numAggregationFunctions; i++) {
+            values[numUnionColumns + 1 + i] =
+                cloneIntermediate(aggregationFunctions[i], 
baseValues[numUnionColumns + i]);
+          }
+          taskMap.put(new Key(keyValues), new Record(values));
+        } else {
+          Object[] values = existing.getValues();
+          for (int i = 0; i < numAggregationFunctions; i++) {
+            int valueIndex = numUnionColumns + 1 + i;
+            values[valueIndex] = 
AggregationFunctionUtils.merge(aggregationFunctions[i], values[valueIndex],
+                cloneIntermediate(aggregationFunctions[i], 
baseValues[numUnionColumns + i]));

Review Comment:
   **perf: a serde clone for every OBJECT merge**
   
   Each projection of an OBJECT intermediate (DISTINCTCOUNT, percentiles, AVG 
pair…) makes a serialize + deserialize round-trip, even for the first insert. 
This probably explains the parity or regression at high cardinality. One cheap 
improvement: a set that contains ALL union columns maps base → derived 1:1, so 
it never merges. It can use the base object directly with no clone (other tasks 
only read it). That removes the most expensive set's clones. Also, is the clone 
of the second merge argument really needed? The merge contract normally mutates 
only the first argument. If some functions return or mutate the second, write 
that down and fix those functions.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to