feat: memory accounting for WindowAggExec#23207
Draft
Dandandan wants to merge 1 commit into
Draft
Conversation
`WindowAggExec` buffers the entire input into `WindowAggStream.batches` and then `concat_batches` it into one contiguous copy (a transient ~2x peak) before computing, with no consultation of the task memory pool. On a large unbounded-window input this can OOM-kill the process instead of failing gracefully, and its usage is invisible to pool accounting. This registers a `MemoryConsumer` for the stream and reserves: - each buffered input batch as it arrives, using a `RecordBatchMemoryCounter` so buffers shared across batches (e.g. zero-copy slices from an upstream Sort/Repartition) are counted once rather than once per batch; and - the freshly concatenated batch in `compute_aggregates`, reflecting the transient doubling during the final compute. When the pool limit is exceeded the operator now returns `ResourcesExhausted` attributed to its consumer, matching `ExternalSorter`/`AggregateExec`. The default pool is unbounded, so default behavior is unchanged. This is accounting/backpressure only — it does not by itself reduce peak memory or add spilling. A unit test asserts the operator errors with `ResourcesExhausted` (naming its consumer) under a tiny memory pool; the `window` sqllogictest suite passes under the default pool. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
WindowAggExecbuffers the entire input intoWindowAggStream.batchesand thenconcat_batchesit into one contiguous copy (a transient ~2x peak) before computing — with no consultation of the task memory pool (there is currently noMemoryReservation/MemoryConsumeranywhere in the window operators). On a large unbounded-window input this can be OOM-killed by the OS rather than failing gracefully, and the usage is invisible to pool accounting /EXPLAIN ANALYZE.What changes are included in this PR?
Registers a
MemoryConsumerfor the stream (mirroringExternalSorter) and reserves:RecordBatchMemoryCounterso buffers shared across batches (e.g. zero-copy slices from an upstream Sort/Repartition) are counted once, not once per batch; andcompute_aggregates, reflecting the transient doubling during the final compute.When the pool limit is exceeded, the operator returns
ResourcesExhaustedattributed to its consumer, matchingExternalSorter/AggregateExec.Scope note: this is accounting + backpressure only. It does not by itself reduce peak memory or add spilling (that would be follow-up work); it makes the existing usage visible and fail gracefully.
Are these changes tested?
Yes:
ResourcesExhausted(naming itsWindowAggStreamconsumer) under a tiny (100-byte) memory pool.windowsqllogictest suite passes under the default pool.Are there any user-facing changes?
The default memory pool is unbounded, so default behavior is unchanged. Queries running against a bounded pool that previously over-committed may now fail deterministically with
ResourcesExhaustedinstead of being OOM-killed — the intended behavior.