adriangb opened a new pull request, #24858: URL: https://github.com/apache/datafusion/pull/24858
## Which issue does this PR close? No existing issue. This was found while investigating a production out of memory. Happy to file one if you would like it tracked for the changelog. ## Rationale for this change `GroupsAccumulatorAdapter` reports per-group memory to the `MemoryPool` from `allocation_bytes`, maintained as a delta of `AccumulatorState::size()` measured before and after accumulator work. `size()` includes the scratch `indices` vector's capacity, but that capacity is never charged: `indices` grows in the per-row push loop, which runs *before* `sizes_pre` is taken, and `indices.clear()` afterwards retains the capacity. Both measurements therefore see the same `allocated_size()` and the delta is always zero, so the capacity is charged exactly zero times, permanently. With `datafusion.execution.batch_size = 16384` a single hot group can hold a 64 KB `Vec<u32>` the pool never sees, and it gets worse as group skew increases. The same asymmetry bites at emit time: `evaluate` and `state` call `free_allocation(state.size())`, releasing capacity that was never charged, so `allocation_bytes` drifts down and saturates at zero across partial emits. In the added test the adapter reports 0 bytes while genuinely holding 224. Measured on this base with a 16384-row batch across 1000 groups, 8192 rows in group 0 and the rest spread evenly over the remaining 999, using a 16-byte accumulator: 168,096 bytes truly retained, 96,960 reported before, 71,136 bytes (42%) unaccounted. ### Production motivation This is one of a group of fixes prompted by a production process running DataFusion dying at 10.98 GB, with roughly 75% of the heap in per-group `COUNT(DISTINCT)` accumulators. This one is independent of the others: it is an accounting correction, not a memory reduction. But an aggregate whose reported size collapses toward zero across partial emits gives the pool exactly the wrong signal on the memory-pressure path. ## What changes are included in this PR? The growth is charged explicitly. A new private `indices_allocation_bytes` field records what has already been charged. Each batch totals the current capacity in the pass that already visits every group and charges only the difference, so a group whose `indices` grew once and was then cleared stays charged without being re-charged. Emitting a state drops its capacity from that total. The established invariant is `allocation_bytes == sum(state.size()) + states.allocated_size()`, which the tests assert against an oracle recomputed directly from the states. `sizes_pre` was deliberately not moved before the push loop: `groups_with_rows` is not known until after it, so a pre-measurement there would need `size()` for every group (the bottleneck the existing code comment warns about) or a per-row branch in the push loop. This change adds no `size()` call and no per-row work, only one `usize` addition per group per batch in an existing loop. ## What is the testing strategy for this PR? Three tests are added in `datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs`, all failing before this change and passing after, covering the skewed update path, the merge path, and release on partial and full emit. Each asserts the invariant above against an oracle recomputed from the states, rather than against a hardcoded byte count. `cargo check`, `cargo clippy --tests`, `cargo test` for `datafusion-functions-aggregate-common` (50 passed) and `cargo fmt --check` all pass locally. ## Are there any user-facing changes? No public API change (the new field is private) and no change to query results. Only the accounting arithmetic changed, so a memory-limited aggregate now reports its true usage to the `MemoryPool` and may spill or fail where it previously ran past its limit undetected. ## Noted while reading, not fixed here `invoke_per_accumulator` returns `?` on `take_arrays` failure with `state.indices` already populated and never cleared, so a retried call on the same adapter would double-push. An error aborts the query today so it is unreachable, but it is a latent trap if that path ever becomes recoverable. -- 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]
