fix: preserve external sort workspace across spilling - #24740
Conversation
Keep acquired spill workspace available to the sorter's cursor, row, and batch reservations instead of returning it to the execution pool before the spill merge. Reuse idle workspace for chunked sorted output while keeping additional output growth under the original sort consumer. Release the idle reserve after the final output path has selected its budget, preserving ownership and accounting for live reservations. Add regressions for reduced memory availability, intermediate and skewed merge retries, parent pool limits, and cleanup on completion or cancellation. Closes apache#24739.
7592082 to
94c1c16
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24740 +/- ##
==========================================
+ Coverage 81.47% 81.48% +0.01%
==========================================
Files 1122 1123 +1
Lines 403582 404083 +501
Branches 403582 404083 +501
==========================================
+ Hits 328821 329275 +454
- Misses 55512 55521 +9
- Partials 19249 19287 +38 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Thanks @sunchao , I have a suggestion: The fix doesn't apply when spilling with a single in-memory batch
Probing the Two consequences on that path: the chunked-output 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 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 Note that none of the three calls is currently covered — removing all three // ... 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 batch
assert!(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 |
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. 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_decreasesfails withResourcesExhausted("allocation limit reached")when its test-only fixture is added to upstreamee59f628b44eb80e8a4f126288632ef51fda5dc2without 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=4passes 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_encryptionenabled, excluding the examples, benchmarks, and CLI packages as prescribed by the contributor checks.cargo test --locked --profile=ci -p datafusion-cli -- --test-threads=4passes 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.shsuite 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.rscases with one million rows, four mixed-type sort keys, low/high cardinality, and zero/five extra payload columns. Both versions usedrelease-nonltoand 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: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.