yashmayya opened a new pull request, #19416: URL: https://github.com/apache/pinot/pull/19416
## Problem `DistinctTable#mergeDistinctTable` is the cross-segment merge that `DistinctCombineOperator` runs on the main query thread. It never reported resource usage to the query accountant, and neither did its only caller chain — `DistinctResultsBlockMerger#mergeResultsBlocks` → `BaseSingleBlockCombineOperator#mergeResults`. That loop has no sampling and no termination check at all. The consequence: `QueryThreadContext.sampleUsage()` is what refreshes a thread's memory snapshot (`ResourceUsageAccountantFactory#sampleUsage`). Without it the watcher thread keeps reading a stale footprint, so it cannot pick a runaway `SELECT DISTINCT` as the kill candidate. The first sample on this path only landed in `toDataTable()` — after the table was fully built. The merge is unbounded when the query has no limit, so this is exactly where a runaway grows largest. Group-by already samples in its equivalent merge loops (`GroupByCombineOperator:143,160`, `SortedGroupByCombineOperator:158,172`), so DISTINCT was the odd one out. ## Fix Add the standard `QueryThreadContext.checkTerminationAndSampleUsagePeriodically` call to all three merge branches (unbounded, limited, limited with order-by) of the 8 concrete `DistinctTable` subclasses. `DictIdDistinctTable` and `EmptyDistinctTable` throw from `mergeDistinctTable`, so they need nothing. The counter starts at 0 and the periodic mask fires at 0, so every non-empty merge samples at least once. That is what `mergeResultsBlocks()` needed as well — a second call there would only double-sample per segment, since all of its allocation happens inside `mergeDistinctTable`. ## Behavior change With accountant-based killing enabled, a DISTINCT query that previously ran to completion can now be killed mid-merge with `SERVER_RESOURCE_LIMIT_EXCEEDED`. That is the intent. A mid-merge timeout now reports the merge scope (e.g. `Timing out on: IntDistinctTable#mergeDistinctTable`) rather than `Timed out while polling results block`. Both still surface as `EXECUTION_TIMEOUT` via `ExceptionResultsBlock`, so only the message text changes. Exceptions thrown from the merge are already handled by `BaseSingleBlockCombineOperator#mergeResultsAndAttachExecutionStats`, which prefers the terminate exception and preserves `QueryException` error codes. ## What deliberately did *not* change **The leaf scan path needs no fix.** `ProjectPlanNode` always builds a `DocIdSetOperator`, and `DocIdSetOperator#getNextBlock` already calls `QueryThreadContext.sampleUsage()` once per doc-id block. Since `DistinctOperator` and `InvertedIndexDistinctOperator#executeScanPath` both pull from that chain, `DistinctExecutor#process` is already accounted at ≤10k-doc granularity, and `BaseOperator#nextBlock` already supplies the termination check and OOM-pause hook. Adding a sample to those loops measurably raised the count for a 3-block scan from 4 to 7 and changed nothing else, so it was dropped. **`mergeDataTable` / the broker reduce path is left alone**, and is a separate follow-up. `DistinctDataTableReducer#mergeToDistinctTable` samples once per server `DataTable` before each `mergeDataTable` call and each broker-side `DistinctTable(…, DataTable)` constructor, so each unaccounted stretch is bounded by one server's payload rather than being unbounded. It is a weaker version of the same gap on a different node, and folding it in would widen this change beyond one concern. ## Testing New `DistinctTableMergeAccountingTest` — 48 cases covering the full cross-product of 8 table types × 3 merge branches, since each of those 24 call sites is hand-written and a misplaced call or a `++n`/`n++` slip would otherwise ship silently. Two assertions per case: - sampling happens on the expected cadence (exact count, not a lower bound) - a `terminate()` raised from the accountant mid-merge actually stops the merge, asserted on the exact number of values merged before the throw **All 48 fail on the unpatched code and pass with the fix.** 107 distinct/selection tests green. `spotless`, `checkstyle`, and `license` checks pass on `pinot-core`. -- 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]
