sunchao opened a new pull request, #24740: URL: https://github.com/apache/datafusion/pull/24740
## Why are the changes needed? Closes #24739. An external sort can fail while spilling even after it has acquired enough workspace for that spill. Spill preparation returns the reserve to the parent memory pool, then cursors and encoded rows request fresh capacity. Those requests can fail if the allocation limit has decreased or another consumer has taken the released capacity. For example, a sorter holding 80 MiB of input and 16 MiB of spill workspace can encounter a new allocation ceiling of 64 MiB. Releasing the workspace leaves 80 MiB reserved, so even a 1 MiB cursor request can fail. Reusing the workspace lets the spill proceed without asking for those bytes again. This follows [#20642](https://github.com/apache/datafusion/pull/20642), which preserved the reservation for the final disk merge, and addresses the remaining release during spill preparation. ## What changes were proposed in this PR? When spilling requires a merge of separately sorted batches, the sorter keeps one parent reservation for its workspace and lets its cursor, encoded-row, and merge-buffer reservations draw from it. A private pool tracks this shared capacity, so moving bytes between these reservations does not return them to the execution pool or require them to be granted again. The workspace remains available across spills, intermediate merge passes, and retries that split an oversized spill batch. Chunked sorted output can temporarily borrow unused workspace when its accounted size exceeds the input reservation. Borrowed bytes return as batches are emitted or the stream is dropped. Any remaining growth of the sorted output still goes through the original sort consumer, preserving the parent pool's limit and fair-share checks. The sorter releases unused workspace when no further spill needs it. For a disk merge, this happens after the final pass has selected its buffer budget; live reservations remain charged. Other consumers can reclaim the idle capacity without waiting for the previous sort's output stream to be dropped. The user-facing change is that sorts affected by this reservation loss can spill successfully. SQL behavior, configuration options, and public APIs are unchanged; insufficient memory beyond the available reservations still produces an error. ## How was this PR tested? The regression `test_spill_preserves_merge_workspace_after_limit_decreases` fails with `ResourcesExhausted("allocation limit reached")` when its test-only fixture is added to upstream `ee59f628b44eb80e8a4f126288632ef51fda5dc2` without the production fix. It passes with this change. Both runs were rebuilt from the corresponding source, using the same fixture and dependency lockfile. `cargo test --locked --profile=ci -p datafusion-physical-plan -- --test-threads=4` passes all 1,839 unit tests and 11 doctests; 23 doctests are ignored. The skew regression exercises both one and two re-spills while a competing consumer holds all other pool capacity, then checks the complete output, batch-size cap, and cleanup. The extended workspace suite also passes: 10,891 tests, eight ignored, and all 505 SQL logic test files. It was run with `avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption` enabled, excluding the examples, benchmarks, and CLI packages as prescribed by the contributor checks. `cargo test --locked --profile=ci -p datafusion-cli -- --test-threads=4` passes all 107 CLI tests. `cargo fmt --all`, `git diff --check`, `cargo clippy --locked --all-targets --all-features -- -D warnings`, and the complete `./dev/rust_lint.sh` suite pass, including license, spelling, Markdown formatting, and Rust documentation checks. Additional cases compare complete sorted results and cover chunked string-view and dictionary output, overlapping output streams, intermediate disk-merge passes, parent-pool limits, and cleanup after completion, errors, or cancellation. I also ran four existing `datafusion/core/benches/sort.rs` cases with one million rows, four mixed-type sort keys, low/high cardinality, and zero/five extra payload columns. Both versions used `release-nonlto` and identical benchmark source and dependencies. An initial ten-sample comparison showed 1.6–3.0% higher mean latency with the patch. I then repeated the already-built binaries in base/patch/patch/base order with 20 samples per case; the averages of the two runs per version were: | Cardinality / extra payload columns | Upstream mean | Patched mean | Change | | --- | ---: | ---: | ---: | | Low / 0 | 117.06 ms | 120.26 ms | +2.73% | | Low / 5 | 136.85 ms | 138.13 ms | +0.94% | | High / 0 | 115.99 ms | 116.78 ms | +0.68% | | High / 5 | 134.78 ms | 137.25 ms | +1.83% | These local measurements show higher mean latency in the selected cases, with visible variation between runs. They exercise normal in-memory sorting with an unbounded pool; they do not measure spill recovery or end-to-end query performance. The bounded-pool regressions above establish the correctness benefit. AI assistance: Codex assisted with the implementation, regression tests, PR text, and the reported local checks and source reviews. -- 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]
