andygrove opened a new pull request, #25814: URL: https://github.com/apache/datafusion/pull/25814
## Which issue does this PR close? - Part of #25758 - Related to #25423 and #25804 ## Rationale for this change On `branch-55`, a final aggregate that has spilled can fail with `ResourcesExhausted` while replaying its spill files. The spill merge reserves as much memory as the pool will give it, and the `OrderedFinalAggregateStream` that consumes the merged rows cannot spill, so when the merge takes the whole allowance the replay's first resize fails. #25423 is the same failure in `aggregate_memory_spill.slt`. DataFusion Comet, which uses 55.1.0, hits this under memory pressure (apache/datafusion-comet#6254). A `GROUP BY` over 2M distinct 128-character strings fails every time with a 24 MiB per-task budget. ## What changes are included in this PR? This PR backports #25383 from @sunchao to the `branch-55` line. When the merge feeds an aggregate replay, it only takes spill files that would still fit if it reserved twice their buffers, and releases the extra before the replay starts. This PR is stacked on #25805, the backport of #24740, which #25383 builds on. The first commit is the one from #25805, and only the second is new. I will rebase once #25805 merges. On top of #25805, `multi_level_merge.rs` and `streaming_merge.rs` match `main` at #25383 exactly. Two changes were needed for `branch-55`: - `ordered_single_stream.rs` (#24259) does not exist on `branch-55`, so its one-line change is dropped. There, an ordered `Single` aggregate uses `GroupedHashAggregateStream`, which this PR covers, so the `ordered_single` case of `migrated_aggregate_spill_merge_leaves_memory_for_replay` expects `StreamType::GroupedHash`. - `test_order_is_retained_when_spilling` keeps `branch-55`'s `new_spill_ctx` and drops the accumulator phase time assertion, which needs #24423. It uses the same 1024-byte limit as #25383. ## Are these changes tested? Yes. The backport includes the tests from #25383. On this branch I ran: - `cargo test -p datafusion-physical-plan --lib` - the `aggregate_memory_spill`, `ordered_aggregate_spill`, `nested_loop_join_spill` and `sort_merge_join_spill` sqllogictests - `cargo fmt --all -- --check` and `cargo clippy -p datafusion-physical-plan -p datafusion-execution --all-targets --all-features -- -D warnings` I also built Comet against this branch. Its reproduction of apache/datafusion-comet#6254 fails on 55.1.0 and passes here, with the spill merge fan-in left unlimited. ## Are there any user-facing changes? Aggregations that spill leave memory for processing the merged rows, so they fail less often under memory pressure. This can mean smaller merges and more merge passes. No public API changes. -- 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]
