fix: preserve external sort workspace across spilling

[sunchao]

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, 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. MergeMemoryPool, provided by datafusion_execution::memory_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 multi-run spills, intermediate merge passes, and retries that split an oversized spill batch. Single-batch spills also retain this workspace, including a spill triggered by the next input batch and the final leftover batch after an earlier spill.

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 small multi-batch concatenation path still releases idle workspace before resizing the ordinary sort reservation, because that reservation cannot borrow merge workspace. This path can still require a fresh parent allocation.

The user-facing change is that sorts affected by the covered reservation losses can spill successfully. SQL behavior and configuration options are unchanged; insufficient memory beyond the available reservations still produces an error. MergeMemoryPool and WorkspaceLoan are additive public APIs in datafusion_execution::memory_pool, documented with a runnable example and the single-parent-consumer accounting policy.

How was this PR tested?

The single-batch regressions were first run with the original spill branch logic: the reduced-limit case failed with ResourcesExhausted("allocation limit reached"), and the live-loan case detected that the retained reservation had been released. Both pass after removing the empty/single-batch releases. The reduced-limit test covers insertion-triggered spilling and finalization, complete sorted output, the batch-size cap, and cleanup; the cancellation test exercises the actual single-input dispatch and checks the live workspace loan.

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.

The updated source was validated with Rust 1.97.0:

• Extended workspace coverage: 10,918 Rust tests passed and eight were ignored, including 1,913 physical-plan tests and 113 execution tests. Features were avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption, with examples, benchmarks, and CLI excluded as prescribed by the contributor checks.
• All 505 SQL logic test files passed with four test threads.
• All 107 CLI tests passed.
• The execution and physical-plan documentation suites passed 29 doctests, with 23 ignored, including the new public-pool example.
• cargo fmt --all, git diff --check, cargo clippy --locked --all-targets --all-features -- -D warnings, and the complete ./dev/rust_lint.sh suite passed, including license, spelling, Markdown, and Rust documentation checks.

The first extended run passed 117 fuzz cases but hit the process's 1,024-open-file limit in test_sort_10k_mem. The same test executable passed that case with the limit raised to 65,536. The remaining workspace suites were then completed with the higher limit while skipping the already-exercised fuzz module; the five SQL test-driver unit tests and SQL logic tests were run separately. No source change was needed for this environment limit.

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.

On the earlier 94c1c16 revision, 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.

The updated source was also compared with upstream main at the PR base (d1fc4b334) using the existing low-cardinality Utf8View sort cases. Both revisions used Rust 1.97.0, release-nonlto, identical benchmark source and dependencies, independent target directories, and verified source/executable hashes. Two rounds ran in main/patch/patch/main order with four Tokio workers, 30 samples per case, one second of warmup, and five seconds of measurement:

Rows	Main mean	Updated mean	Change
100K	5.549 ms	5.425 ms	-2.24%
1M	61.098 ms	59.625 ms	-2.41%

The means average the two rounds. No regression was detected in these two cases; the 100K differences were not statistically significant, and the 1M reverse-order comparison fell within Criterion's noise threshold. These normal, unbounded-memory sorts do not establish a general speedup, measure constrained spilling, or isolate the review changes from the rest of the PR.

AI assistance: Codex assisted with the implementation, regression tests, PR text, and the reported local checks and source reviews.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 91.12554% with 41 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.63%. Comparing base (d1fc4b3) to head (804fbca).
⚠️ Report is 69 commits behind head on main.

Files with missing lines	Patch %	Lines
...ion/execution/src/memory_pool/merge_memory_pool.rs	90.00%	11 Missing and 21 partials ⚠️
...usion/physical-plan/src/sorts/multi_level_merge.rs	91.17%	0 Missing and 9 partials ⚠️

Additional details and impacted files

@@ Coverage Diff @@## main #24740 +/- ##
==========================================
+ Coverage 81.47% 81.63% +0.16% 
==========================================
Files 1122 1124 +2 Lines 403582 410317 +6735 Branches 403582 410317 +6735 ==========================================
+ Hits 328821 334974 +6153 - Misses 55512 55636 +124 - Partials 19249 19707 +458 

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[jayzhan211]

Thanks @sunchao , I have a suggestion:

The fix doesn't apply when spilling with a single in-memory batch

in_mem_sort_stream is called with is_output_stream = false only from
sort_and_spill_in_mem_batches, but the three release_unused() calls at
sort.rs:608, sort.rs:626 and sort.rs:635 are unconditional. On the output
path they're already redundant — sort() released at sort.rs:389 — so they
take effect only on the spill path, handing the workspace back to the
execution pool right before it's needed.

Probing the borrow() at sort.rs:775:

spill with 2 in-mem batches: want=950272 got=950272 # works
spill with 1 in-mem batch: want=950272 got=0 # workspace already released

Two consequences on that path: the chunked-output try_resize falls back to the
parent pool, and the reserve_memory_for_merge() at the end of the spill has to
re-acquire the full sort_spill_reservation_bytes from the parent — so it can
still fail with ResourcesExhausted under contention, which is what #24739 is
about. in_mem_batches.len() == 1 at spill time is common: a spill triggered by
the second batch, or sort() flushing a single leftover batch after an earlier
spill.

Suggested fix — drop the two calls that are dead-or-harmful:

 if self.in_mem_batches.is_empty() {
- self.merge_pool.release_unused();
let empty_stream =

 if self.in_mem_batches.len() == 1 {
- self.merge_pool.release_unused();
let batch = self.in_mem_batches.swap_remove(0);

The third one (the sort_in_place_threshold_bytes concat branch) is genuinely
different and should stay: concat_batches grows self.reservation, an
execution-pool consumer that can't borrow merge workspace, so the idle floor has
to be returned for that growth to have anywhere to come from. Worth a comment,
since it now reads as inconsistent with the two above:

 if self.reservation.size() < self.sort_in_place_threshold_bytes {
+ // Unlike the paths above, `concat_batches` grows `self.reservation`,+ // an execution-pool consumer that cannot borrow merge workspace.+ // Return the idle floor so that growth has somewhere to come from.
self.merge_pool.release_unused();

With those changes the probe reads got=950272 on both spill passes and all 108
sorts:: tests stay green.

Note that none of the three calls is currently covered — removing all three
leaves the suite at 109/109 green — so this needs a regression test. Roughly
check_chunked_string_view_workspace with the pool sized for one batch instead
of two, so the spill happens while exactly one batch is buffered:

// ... same Utf8View batches / ordering as check_chunked_string_view_workspace ...let input_bytes = get_reserved_bytes_for_record_batch(&batches[0])?;// Room for exactly one input batch plus the merge workspace.let capacity = options.sort_spill_reservation_bytes + input_bytes;// ... build sorter over `pool` ...
sorter.insert_batch(batches[0].clone()).await?;assert_eq!(pool.reserved(), capacity);
sorter.insert_batch(batches[1].clone()).await?;// spills with one buffered batchassert!(sorter.spilled_before());assert_eq!(sorter.in_mem_batches.len(),1);let stream = sorter.sort().await?;drop(sorter);let output:Vec<RecordBatch> = stream.try_collect().await?;assert_eq!(concat_batches(&schema,&output)?.num_rows(),2* rows);assert_released(&pool,&runtime).await;

To make it a true regression test it needs an assertion that the borrow actually
happened rather than just that the sort succeeded — e.g. that
pool.state.lock().unwrap().denied doesn't increase across the spill, or
exposing the loan size the way test_chunked_sort_returns_live_workspace_loan_on_drop
does.

[jayzhan211]

Should we also move MergeMemoryPool to datafusion_execution crate

[sunchao]

Moved in 804fbca9c. MergeMemoryPool and WorkspaceLoan now live in datafusion_execution::memory_pool, with public API documentation and a tested example. Updated the sort and merge callers to use that shared implementation.

[sunchao]

@jayzhan211 Thanks for the detailed repro. Fixed your single-batch spill finding in 804fbca9c.

Removed release_unused() from the empty and single-batch branches. Kept it in the concatenation branch and added the explanation you suggested: concatenation can grow the ordinary sort reservation, which cannot borrow merge workspace. That path can still require a fresh parent grant, as documented in the PR description.

Added test_single_batch_spill_preserves_workspace_after_limit_decreases, covering both a spill triggered by inserting the second batch and the final leftover single-batch spill. The reduced limit makes the first spill fail if its workspace is released and must be reacquired.

Also updated test_single_batch_spill_returns_live_workspace_loan_on_drop to exercise in_mem_sort_stream(false, false), assert the remaining loan through the available workspace credit, and verify cleanup when the stream is dropped early. Both regressions fail with the old branch logic and pass with the fix.

Local workspace and SQL logic tests, formatting, Clippy, and the full lint script passed; validation details are in the PR description.

[jayzhan211]

Thanks @sunchao , LGTM!