adriangb opened a new pull request, #27343: URL: https://github.com/apache/datafusion/pull/27343
> [!NOTE] > Stacked on https://github.com/apache/datafusion/pull/25910. Review only the last two commits: `test: spilled aggregates over wide states in sqllogictest` and `fix: bound aggregate spill batches by the memory of the first spill`. The other commits belong to https://github.com/apache/datafusion/pull/25910. ## Which issue does this PR close? - Part of https://github.com/apache/datafusion/issues/26711. - Takes the aggregate pieces of https://github.com/apache/datafusion/pull/26987 (by Jay Zhan): the first-spill bound, the config bound, and the aggregate queries of `spill_batch_bound.slt`. ## Rationale for this change A `string_agg` over many groups fails with `ResourcesExhausted` under a 2 MB limit: ```sql SET datafusion.execution.target_partitions = 1; SET datafusion.runtime.memory_limit = '2M'; SELECT count(*), sum(length(s)) FROM ( SELECT k, string_agg(v, ',') AS s FROM (SELECT value % 40000 AS k, repeat('x', 200) AS v FROM generate_series(1, 160000)) GROUP BY k ); -- main: Resources exhausted: Additional allocation failed for SingleHashAggregateStream[0] ``` The aggregate spills runs of up to 1 MiB batches (https://github.com/apache/datafusion/pull/25853). The replay merge reserves twice the largest batch of each of at least two runs, which is more than the aggregate held when it first spilled. ## What changes are included in this PR? 1. `test:` the aggregate queries of https://github.com/apache/datafusion/pull/26987, which pass on `main` and guard against regressions (array_agg over a few large groups from https://github.com/apache/datafusion/issues/25851, string_agg at 4 MB, the legacy stream with `enable_migration_aggregate = false`). Two new 2 MB queries record the error of `main`. 2. `fix:` both aggregate streams set a bound on their `SpillManager`: `min(1 MiB, sort_spill_reservation_bytes / 4)`. The 1 MiB target of https://github.com/apache/datafusion/pull/25853 stays the default, and a small `sort_spill_reservation_bytes` lowers it, as in https://github.com/apache/datafusion/pull/26987. On the first spill, the bound drops so that the smallest merge, plus as much again for the replay, fits in half of what the aggregate held. The aggregate writes through `append_batch_with_stats`, so the merge sizes intermediate batches from the rows of each piece. The legacy stream passes the bound to its merge. The 2 MB queries return their results. ## What is the testing strategy for this PR? - `spill_batch_bound.slt` (see above). - `migrated_aggregate_spill_merge_leaves_memory_for_replay` allows 1 KiB over the pool limit. The replay merge grows a second reservation of the aggregate's consumer while the first holds a few bytes. `FairSpillPool` checks each reservation against the fair share on its own, so the total can exceed the share by those bytes. That pool behavior exists on `main`. Smaller batches let the merge reach the share exactly, which exposes it. - The `aggregate_memory_call_sequence.rs` snapshots do not change, because the merge keeps its re-spill. - The extended suite passes. Not fixed: `array_agg` over 4 groups at 6 MB still fails, on `main` and here, because a group's state is too large for the pool. ## Are there any user-facing changes? Spilling aggregates can write more, smaller spill batches, and finish under smaller memory limits. No API changes. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- 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]
