yashmayya opened a new pull request, #19419: URL: https://github.com/apache/pinot/pull/19419
Bounds the memory a `DISTINCT` query uses on an MSE leaf stage, by flushing the accumulated distinct values downstream once they reach a threshold and letting the partitioned intermediate stage finish the de-duplication. This is the DISTINCT counterpart to `StreamingGroupByCombineOperator` (#18035, #18510). ### Why group-by was covered but DISTINCT was not MSE plans `SELECT DISTINCT a, b` as an `Aggregate` with zero aggregate calls. `ServerPlanRequestVisitor#visitAggregate` emits a leaf query of `SELECT a, b FROM t GROUP BY a, b`, and `NonAggregationGroupByToDistinctQueryRewriter` (registered in `ServerPlanRequestUtils`) then rewrites that back into `SELECT DISTINCT a, b`. So the leaf runs `DistinctOperator` / `DistinctCombineOperator`, not the group-by machinery. `CombinePlanNode` only routed selection-only and group-by aggregation queries to a streaming combine operator, and `isDistinctQuery` is a third, disjoint query class — so DISTINCT could never reach one. ### How large the leaf table actually gets The planner pushes `LIMIT` (plus `OFFSET`) into the LEAF aggregate for any distinct query that has one — `PinotAggregateExchangeNodeInsertRule#isGroupTrimmingEnabled` is unconditionally true for aggregates with no aggregate calls. The checked-in goldens show it: `SELECT DISTINCT col1, col2 FROM a LIMIT 10` produces `aggType=[LEAF], limit=[10]`. So the leaf table is normally bounded at `LIMIT`, not unbounded — it is only unbounded under `is_enable_group_trim='false'`. The problem is that `LIMIT` is routinely far larger than the memory worth spending on one leaf (`SELECT DISTINCT user_id ... LIMIT 10000000` still OOMs a server), and `numGroupsLimit` does not apply to the distinct path at all. The only existing lever is the `is_skip_leaf_stage_group_by` hint, which is all-or-nothing. ### Change New `StreamingDistinctCombineOperator`. It merges per-segment blocks into an accumulated `DistinctTable` and flushes it once it reaches the threshold, then starts a fresh one. - Query option: `streamingDistinctFlushThreshold` - Broker config: `pinot.broker.mse.streaming.distinct.flush.threshold`, default `-1` (off) The FINAL stage de-duplicates the partial flushes because the hash exchange is keyed on the distinct columns and set union is associative, commutative and idempotent. `CombinePlanNode` selects it only when all of these hold: | Guard | Why | |---|---| | threshold > 0 and `isDistinctQuery` | | | no ORDER BY | an ordered distinct keeps a bounded top-`LIMIT` heap, so it is already memory-bounded | | leaf `LIMIT` > threshold | at or below it there is no memory to save, and the early-exit short-circuit is worth more | | not `serverReturnFinalResult` | otherwise no aggregate sits above the leaf to de-duplicate across flush windows | ### Trade-off, stated explicitly When it engages, the leaf gives up the cross-segment `isQuerySatisfied()` early exit: flushing empties the accumulator long before it can reach `LIMIT`, so every segment gets scanned. That is deliberate — a lower memory ceiling for more scan work — and it is called out in the code rather than left for someone to discover. ### Notes for reviewers Unlike the group-by operator, no `detachFromWorkerThreadState` hook is needed: every per-segment distinct operator (`DistinctOperator`, `DictionaryBasedDistinctOperator`, `InvertedIndexDistinctOperator`, `JsonIndexDistinctOperator`) returns a fully materialized table — the dictionary-based executors resolve dict IDs to real values in `getResult()` — so a block handed to the consumer is self-contained. The operator does override `isQuerySatisfied` to `false`, matching the group-by operator. `BaseStreamingCombineOperator#processSegments` calls it on the producing thread with a block it has *already* published, and this operator adopts that block as its accumulator and mutates it. Letting the worker read `isSatisfied()` off it afterwards would be a data race on a plain `HashSet` mid-`add`/`rehash`. Nothing is lost — the consumer still evaluates satisfaction itself, and the cross-segment early exit is given up by construction anyway (above). Two limits are documented on the class rather than papered over: 1. It bounds only the server-level table. Per-segment `DistinctTable`s are still built in full (group-by caps its per-segment maps with `numGroupsLimit`; distinct has no equivalent), so peak leaf heap is roughly `flushThreshold` plus one table per in-flight worker thread. When a single segment already holds more than `flushThreshold` values, its block is adopted and flushed with nothing merged in, so the leaf does no cross-segment de-duplication at all. Requiring a merge before each flush would let the accumulator grow to two segments' worth and raise the very ceiling this exists to lower — bounding the per-segment tables is the real fix, and is left as follow-up. 2. The `maxRows*InDistinct` guardrails are evaluated inside `DistinctResultsBlockMerger#mergeResultsBlocks`, which never runs for the adopted first block of a window — so in the no-merge regime above none of them fire. Where merging does happen, `maxRowsInDistinct` and `maxExecutionTimeMsInDistinct` stay per-query bounds because the scanned-doc count is deliberately carried across flush windows; `maxRowsWithoutChangeInDistinct` still degrades. Feature is off by default, so behaviour is unchanged unless the option or broker config is set. ### Testing - `StreamingDistinctCombineOperatorTest` — 16 tests over real segments: multi-flush correctness, values repeated across flush windows (the idempotence argument), an accumulator that merges across segments before flushing, all four per-segment operators and all three `getResult()` materialization paths, null handling across a flush boundary, empty results, the `CombinePlanNode` guard matrix including the non-distinct negative case, the query-option wiring through the real plan entry point, and `maxRowsInDistinct` remaining a per-query bound with its early-termination reason surviving to the terminal metadata block. - `MultiStageBrokerRequestHandlerTest` — broker default injection, `putIfAbsent` precedence, and no injection when unset. - `QueryOptionsUtilsTest` — option parsing bounds (also backfills the group-by option, which had none). - `MultiStageEngineIntegrationTest` — end-to-end vs a baseline, comparing row multisets, with an `explainAskingServers` assertion that the streaming operator was actually selected, plus negative cases for each guard. Note none of these queries may carry `ORDER BY`: the planner pushes collations into the leaf aggregate for zero-agg-call aggregates, which disqualifies the operator and would make the comparison vacuous. -- 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]
