Skip to content

fix: charge accumulators' construction-time state to the memory pool in ungrouped aggregation - #24692

Open
ranflarion wants to merge 1 commit into
apache:mainfrom
ranflarion:ungrouped-agg-initial-charge
Open

fix: charge accumulators' construction-time state to the memory pool in ungrouped aggregation#24692
ranflarion wants to merge 1 commit into
apache:mainfrom
ranflarion:ungrouped-agg-initial-charge

Conversation

@ranflarion

Copy link
Copy Markdown

Which issue does this PR close?

Rationale for this change

An ungrouped aggregation can hold arbitrarily large accumulator state that the memory pool never sees. AggregateStream grows its reservation only by per-batch size() deltas (size_post.saturating_sub(size_pre) in aggregate_batch), so an accumulator that allocates its retained state in its constructor and never resizes it is charged nothing for the stream's lifetime. The grouped path already handles this: GroupsAccumulatorAdapter charges state.size() when it creates each accumulator. Embedding DataFusion 54.1.0 in a Spark accelerator we measured Spark's runtime-filter bloom_filter_agg (an 8 MiB bit array zero-filled at construction, one per concurrent task) running entirely unaccounted, so a fair-spill pool could neither fail admission nor pressure other consumers to spill.

What changes are included in this PR?

One charge in AggregateStream::new: after registering the reservation, try_grow the sum of the accumulators' initial size(). Per-batch deltas compose on top of it, so there is no double counting. test_oom's ungrouped arm is updated for the earlier failure point: with its 1-byte limit the median accumulator's initial 48 bytes now fail admission at execute_typed, so that arm asserts the construction-time ResourcesExhausted instead of collecting first; the grouped arms are unchanged.

Are these changes tested?

The updated test_oom is the regression pin: on main without the fix the ungrouped stream constructs successfully and the arm fails, with the fix it passes. The full datafusion-physical-plan lib suite passes (1784 tests).

Are there any user-facing changes?

Yes, behavioral: an ungrouped aggregation whose accumulators' construction-time state exceeds the memory limit now fails with ResourcesExhausted at execute instead of silently running past the configured limit, and memory-limited setups that relied on that overshoot may need limits raised to what the query actually uses. No API changes.

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Aug 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Ungrouped aggregation never charges accumulators' construction-time state to the memory pool

1 participant