xiangfu0 commented on code in PR #19264:
URL: https://github.com/apache/pinot/pull/19264#discussion_r4149064456
##########
pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java:
##########
@@ -87,6 +121,530 @@ 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.
+ ///
+ /// When the worst-case derived output fits `numGroupsLimit`, the grouping
sets are partitioned round-robin
+ /// across up to `numTasks` workers, each with task-local maps that are
unioned without contention. Otherwise,
+ /// the derive processes sets coarsest-first into one capped map, so it
never materializes base groups times
+ /// the number of sets. 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);
+ // A base group can contribute to every set. Bound the derived groups
before materializing them, while
+ // reserving at least one slot per set so a coarse set such as the grand
total cannot be starved.
+ int derivedUpperBound = (int) Math.min((long) baseTable.size() * numSets,
Integer.MAX_VALUE);
+ int derivedCap = Math.max(queryContext.getNumGroupsLimit(), numSets);
+ int capacityBound = Math.min(derivedUpperBound, derivedCap);
+ int initialCapacity = getIndexedTableInitialCapacity(capacityBound,
capacityBound,
+ queryContext.getMinInitialIndexedTableCapacity());
+
+ List<Map.Entry<Key, Record>> baseEntries = new
ArrayList<>(baseTable.getRecordEntries());
+ if (derivedUpperBound > derivedCap &&
useDeterministicIndexedTable(queryContext)) {
+ baseEntries.sort(Map.Entry.comparingByKey());
+ }
+ // A full-union set (one that contains every union column, i.e. the
identity grouping) maps base groups to
+ // derived groups 1:1: keys are unique so no merge ever runs on its
records, and its stored intermediates can
+ // safely be the base objects themselves (after the pre-serialize pass
below, no other task touches them).
+ boolean[] isFullUnionSet = new boolean[numSets];
+ for (int s = 0; s < numSets; s++) {
+ boolean fullUnion = true;
+ for (int col = 0; col < numUnionColumns; col++) {
+ fullUnion &= setContains[s][col];
+ }
+ isFullUnionSet[s] = fullUnion;
+ }
+
+ /// Serialize each OBJECT base intermediate exactly ONCE, on exactly one
thread (parallel by base-entry
+ /// range, so no object is shared between pre-pass tasks). This is a
correctness requirement, not just a
+ /// perf win: serialization can MUTATE the accumulator (TDigest#compress
rewrites centroid arrays; theta /
+ /// CPC / tuple sketch accumulators flush their pending lists in
getResult), so concurrent serialization of
+ /// the same base object by multiple derive tasks corrupts it. After this
pass, derive tasks read only the
+ /// immutable bytes and deserialize their own private copies.
+ SerializedIntermediateResult[][] serializedIntermediates =
+ preSerializeObjectIntermediates(baseEntries, aggregationFunctions,
numUnionColumns, numTasks,
+ queryContext, executorService);
+
+ // Keep the contention-free parallel path when even the worst-case fan-out
fits the group limit. Otherwise,
+ // derive coarsest-first into one capped map; building every per-set map
before trimming can use
+ // baseTable.size() * numSets records of heap even though only derivedCap
groups can be returned.
+ if (derivedUpperBound > derivedCap) {
Review Comment:
Fixed in dca0a82. The instance planner now sums conservative dictionary
bounds for every grouping set across all local segments, including the
grand-total set. BASE derivation is selected only when the complete output fits
numGroupsLimit. Ordered queries above that bound use the expansion path and its
normal top-K finish; safe ordered queries can still use BASE. The proven
per-set bound also allows the parallel derive when baseSize × numSets is loose.
The cap-boundary planner test and both sides of ordered integration parity
passed locally (11 focused core tests and 3 selected integration tests
overall). I have not rerun the 64k-group benchmark on this head, so the PR
description no longer claims a current-head speedup.
##########
pinot-core/src/main/java/org/apache/pinot/core/data/table/GroupingSetsBaseIndexedTable.java:
##########
@@ -0,0 +1,93 @@
+/**
+ * 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.data.table;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.query.request.context.QueryContext;
+
+
+/// Combine table for the BASE groups of a grouping-sets base-aggregation
query. Like
+/// [UnboundedConcurrentIndexedTable] with a result-size cap, but instead of
silently dropping records whose
+/// (new) base key arrives after the cap is hit, it RETAINS them in an
overflow buffer. The combine later folds
+/// the overflow records into the already-derived grouping-set groups
(existing groups only), mirroring the
+/// expansion path's behavior under the group limit: the grand total and
coarse subtotals stay exact, and only
+/// the overflowing fine-set groups are lost.
+///
+/// The overflow buffer is bounded (a multiple of the result size); if even
the buffer overflows, further
+/// unknown-key records are dropped and [#isOverflowTruncated()] reports it so
the response can be flagged as
+/// trimmed.
+public class GroupingSetsBaseIndexedTable extends
UnboundedConcurrentIndexedTable {
+ /// Bounds the overflow buffer relative to the result size, so a
pathological key space cannot exhaust heap.
+ private static final int MAX_OVERFLOW_FACTOR = 10;
+
+ private final int _maxOverflowRecords;
+ private final ConcurrentLinkedQueue<Record> _overflowRecords = new
ConcurrentLinkedQueue<>();
+ private final AtomicInteger _numOverflowRecords = new AtomicInteger();
+ private volatile boolean _full;
+ private volatile boolean _overflowTruncated;
+
+ public GroupingSetsBaseIndexedTable(DataSchema dataSchema, QueryContext
queryContext, int resultSize,
+ int initialCapacity, ExecutorService executorService) {
+ super(dataSchema, false, queryContext, resultSize, initialCapacity,
executorService);
+ _maxOverflowRecords = (int) Math.min((long) resultSize *
MAX_OVERFLOW_FACTOR, Integer.MAX_VALUE);
+ }
+
+ @Override
+ protected void upsertWithoutOrderBy(Key key, Record record) {
+ if (_full) {
+ if (!updateExistingRecordIfPresent(key, record)) {
+ // Unknown base key after the cap: retain the record so its
contribution to the coarse grouping sets
+ // (grand total, subtotals) is not lost. A record raced in by another
thread right after the miss is
+ // still correct: the later overflow fold merges into existing groups.
+ if (_numOverflowRecords.incrementAndGet() <= _maxOverflowRecords) {
+ _overflowRecords.add(record);
Review Comment:
Fixed in dca0a82. The 10× raw overflow buffer and its single-threaded fold
are removed. Before building segment operators, the planner conservatively
bounds all local BASE keys and the total derived keys under numGroupsLimit; a
raised groupingSetsBaseAggregationMaxGroups cannot bypass this hard limit.
Unknown BASE keys past the cap fail the query, and a post-worker size check
catches concurrent inserts at the final slot. The unexpected-overflow unit test
and forced-secondary-gate grand-total integration test passed locally. I have
not rerun the 768 MB TDigest reproduction; current-head CI is pending.
--
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]
