yashmayya commented on code in PR #19264: URL: https://github.com/apache/pinot/pull/19264#discussion_r4148522524
########## 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: **Blocker: the overflow buffer gets around the `numGroupsLimit` memory guard** This buffer keeps up to 10 × `numGroupsLimit` raw records (1M by default) until the merge ends. The records are not merged by key, and the size is not bounded by bytes. With sketches, one opted-in query can run a shared server out of memory: - Repro: 60 daily segments with 5k users each, `numGroupsLimit=10000`, `PERCENTILETDIGEST`, 768 MB heap. The opt-in path fails with `OutOfMemoryError` in 2 of 2 runs. Expansion passes, with a peak of about 560–590 MB. - At default settings, 1M TDigest records are about 9 GB. The expansion combine keeps at most 100k groups. It is also slow. The fold runs on one thread, and it deserializes every OBJECT value for every set before it checks that the group exists. With default settings and 20 time-partitioned segments (about 470k overflow records), `ROLLUP(3)` takes 321 ms and expansion takes 38 ms. Past the buffer, the grand total is wrong again: 30 daily segments with 40k users each give 1,100,000 instead of 1,200,000. The flags are the same as for an exact result. Fix idea: fold overflow records into bounded per-set coarse accumulators as they arrive, instead of keeping raw records. At minimum, check that the key exists before you deserialize. ########## 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: **Blocker: with ORDER BY, this capped path returns a wrong top-K** This path keeps sets by column count, then groups by hash order. It does not use the ORDER BY. Expansion keeps every group in the combine, and `finish()` keeps the top `numGroupsLimit` by the ORDER BY. Repro with default settings, one server, 12 daily segments with 10k users each: `SELECT day, userId, SUM(m) FROM t GROUP BY ROLLUP(day, userId) ORDER BY SUM(m) DESC LIMIT 50` - Expansion (same as master): the exact top 50. - Opt-in path: 34 of 50 rows differ, and 6 of the true top 50 are missing. The response is flagged, but master returns the exact result with no flags. The same order can drop a whole set. In a small repro, `GROUPING SETS ((a), (b)) ORDER BY SUM(m) DESC LIMIT 4` returns only (a) rows. Expansion returns the correct (b) rows. Also, `baseSize × numSets` is a loose bound, so this path runs on one thread even when nothing is dropped. At 64k base groups, `CUBE(3)` takes 123 ms here, 55 ms on expansion, and 54 ms on the previous head. Fix ideas: with ORDER BY, keep the top `numGroupsLimit` by the ORDER BY, like `finish()`. Estimate each set's size (for example, min(base size, product of per-column distinct counts)) and stay parallel when it fits. -- 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]
