ranflarion opened a new issue, #24691: URL: https://github.com/apache/datafusion/issues/24691
### Describe the bug The ungrouped aggregation path never charges accumulator state that is allocated at construction to the memory pool. `AggregateStream::new` registers its `MemoryReservation` after `create_accumulators` with nothing charged, and afterwards the reservation only grows by per-batch deltas: `aggregate_batch` measures `accum.size()` before and after each `update_batch`/`merge_batch` and adds `size_post.saturating_sub(size_pre)` (`datafusion/physical-plan/src/aggregates/aggregate_stream.rs:300-303` and `467-473` on current `main`). An accumulator that allocates its retained state in its constructor and never resizes it is therefore invisible: both measurements include the same allocation, the delta stays zero for the stream's lifetime, and the pool believes the operator holds nothing. The grouped path does not have this hole: `GroupsAccumulatorAdapter` charges `state.size()` when it creates each accumulator (`add_allocation(state.size())` in `functions-aggregate-common`'s `groups_accumulator.rs`), so the same accumulator behind a `GROUP BY` is accounted and the ungrouped path is the inconsistent one. Built-in accumulators mostly start near-empty and grow with input, so the deltas cover them; the full cost lands on UDAF accumulators with fixed-size state. We hit it embedding DataFusion 54.1.0 in a Spark accelerator (apache/datafusion-comet lineage): Spark's runtime-filter `bloom_filter_agg` zero-fills its whole bit array at construction, 8 MiB at Spark's default `maxNumBits` of 67108864, and the ungrouped aggregation Spark injects for runtime filters ran one such filter per concurrent task, all invisible to the fair-spill pool. The pool could neither fail admission nor pressure other consumers to spill while the process held cores x 8 MiB it did not know about. ### To Reproduce Unit-shape repro against `main`: an ungrouped `AggregateExec` over a UDAF whose accumulator allocates a 1 MiB `Vec` in its constructor and reports it honestly in `Accumulator::size()`, executed with a memory limit well under 1 MiB. `execute` succeeds and the stream runs to completion with the pool's reserved bytes never exceeding a few hundred. The existing `test_oom` in `aggregates/mod.rs` demonstrates the same blindness from the other side: its ungrouped arm runs a median aggregate against a 1-byte pool, and the stream constructs successfully despite the accumulator's initial `size()` of 48 bytes already exceeding the limit; the test only errors later because median's state happens to grow per batch. ### Expected behavior `AggregateStream::new` should charge the sum of the accumulators' initial `size()` to its reservation, mirroring what `GroupsAccumulatorAdapter` does at accumulator creation. Per-batch deltas then compose on top of the initial charge with no double counting. This is a behavior change for memory-limited setups: an ungrouped aggregation whose construction-time state exceeds the limit now fails admission with `ResourcesExhausted` at `execute` instead of silently running past the configured limit. ### Additional context We run this fix in production against 54.1.0 and can offer it as a PR (one `try_grow` in `AggregateStream::new` plus updating `test_oom`'s ungrouped arm, which becomes the regression pin: it fails on unpatched `main` and passes with the charge). -- 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]
