Repository navigation
Conversation
807afbc to
15c67a7
Compare
b124f7b to
f9c4788
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25565 +/- ##
==========================================
+ Coverage 82.42% 82.53% +0.11%
==========================================
Files 1138 1141 +3
Lines 435501 441203 +5702
Branches 435501 441203 +5702
==========================================
+ Hits 358955 364167 +5212
- Misses 54841 55041 +200
- Partials 21705 21995 +290 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Byte-array output memory can be underestimated before allocation, defeating the new admission guarantee.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 1
What changed in this PR
Adds byte-targeted spill-merge output batching and reserves construction workspace to improve sorting under constrained memory.
Changes:
- Adds byte-aware output batching and reservation retention.
- Re-spills inputs when output headroom cannot fit.
- Updates focused spill, memory, and metrics tests.
| File | Description |
|---|---|
datafusion/physical-plan/src/sorts/builder.rs |
Implements byte-targeted batch construction. |
datafusion/physical-plan/src/sorts/merge.rs |
Passes output workspace into batch construction. |
datafusion/physical-plan/src/sorts/multi_level_merge.rs |
Adds output-headroom admission and re-spilling. |
datafusion/physical-plan/src/sorts/streaming_merge.rs |
Extends merge builder configuration. |
datafusion/physical-plan/src/sorts/sort.rs |
Updates sort batching tests. |
datafusion/physical-plan/src/sorts/sort/spill_tests.rs |
Tests retained workspace behavior. |
datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs |
Adjusts constrained-memory assertions. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| self.validity_buffer_bytes(column_idx, rows_to_emit)? | ||
| .checked_add((rows_to_emit + 1).checked_mul(size_of::<T::Offset>())?)? | ||
| .checked_add(values_len) |
andygrove
left a comment
There was a problem hiding this comment.
Thanks @Weijun-H for taking on spill-merge output accounting. Reserving the output-construction headroom at admission is the right place to charge it. I checked the new tests against the change: with byte targeting disabled, 14 of the new or changed builder, merge and spill_tests tests fail, so they really exercise it. I didn't find any underflow or re-spill loop.
A few things before I'd approve:
- The byte target caps output rows at the largest spilled batch even for fixed-width schemas.
test_sort_spillnow emits 7 batches of at most 1500 rows instead of 2, andtest_sort_fetch_memory_calculationemits 4 × 100 instead of 1 × 400. For fixed-width columns the output size ofbatch_sizerows is known from the schema. Could the target bemax(largest spilled batch, fixed-width bytes for batch_size rows), or apply only when variable-width columns are present? - The target comes from
get_sliced_size, which counts validity asceil(n / 8), but the estimate rounds validity up to 64 bytes. So an output batch identical to the largest input batch gets split: a nullableInt32batch of 1000 rows has a target of 4125 bytes and an estimate of 4128, and is emitted as 999 + 1 rows. Using the same estimator for both would avoid this. - The loosened assertions don't seem necessary. With the original exact checks restored on this branch, the five sort fuzz tests and
should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_to_spillpass. Could we keep them exact, and assert the new exact counts in the two tests where the output really changes? largest_prefix_under_targetrescans the prefix on each probe. In a local release-mode microbenchmark (2M rows of(Int64, Utf8), target of one 1000-row batch) the builder took about 135 ms versus about 15 ms without a target. Emitting the same 2000 batches without a target (batch_size1000) takes about 16 ms, so the cost is in the byte estimation, not the extra batches. Byte prefix sums computed once per build should remove most of it.- The
try_growafter building infinish_record_batchcan only fail once the batch is already allocated, which is the pattern #25183 warns about. In my testing on arrow 60 the estimate was always an upper bound, so adebug_assert!seems enough.
Nits, fine as follow-ups: debug! instead of warn! in the reservation retry loop, since it fires on every output batch under memory pressure; the pool_extra parameter of spill_merge_output_construction_uses_the_real_pool is only called with one value; and a note in the description that sorts close to the memory limit may now re-spill more, or fail where main succeeded, because this memory was previously used without being reserved.
Happy to re-review after 1–3.
|
Thanks for the detailed review! @andygrove Addressed all five points:
Also addressed the nits and documented the stricter memory-accounting behavior. Ready for another look. |
andygrove
left a comment
There was a problem hiding this comment.
Thanks @Weijun-H, I re-checked this at 850a3b2 and all five points from my earlier review are resolved.
- Fixed-width output keeps the configured row count. Initial runs record
batch_sizeas their row limit, so the target ismax(largest spilled batch, fixed-width bytes for batch_size rows). In a CLI A/B againstmain(one partition, 3M-row string sort and 6M-row i64 sort, 40 MB down to 1.5 MB), output batch counts are identical wherever both succeed (367 and 733). - Spill sizes now include the bitmap and fixed-size-binary alignment padding, so the target matches the estimate (nullable
Int32, 1000 rows: 4125 to 4128 bytes). - The exact batch and metric assertions are back. The 32 sort and spill fuzz tests pass with
--features extended_tests. - Prefix estimation is a single forward scan. In a release-mode probe (2M rows of
(Int64, Utf8), target of one 1000-row batch) building takes about 25 ms against about 11 ms without a target, down from 135 ms against 15 ms. Fixed-width schemas show no difference. - The reservation is taken before interleave, and the check afterwards is a
debug_assert!. The 475 sort, spill and aggregate lib tests pass in debug, so the assertion ran.
I also swept the memory limit against main. I found no limit where main succeeds and this PR fails. The string sort at 2.5 MB emits 367 batches against 733 on main. The i64 sort at 2 MB with a 1 MB spill reservation spills 274 times against 244, which is the stricter accounting the description already mentions.
Remaining before I approve:
- The PR now conflicts with
main: imports inbuilder.rs, three hunks inmulti_level_merge.rs, one inspill_manager.rs.mainhas since merged #25820, #25372, #26366 and #23565, which touch the same area, so please rebase and re-run the sort, spill and fuzz tests. - Nit:
spill_tests.rsstill has two looser assertions (first_batch.num_rows() <= rowsandspill_file_count > spills_before). They look justified because the output changes there, but exact values would be better if they are easy to compute.
The Rust (fork) failure is the datafusion-cli S3 tests and looks unrelated.
Preserve max-row spill statistics and intermediate sizing from main. Split output workspace before retaining the input grant, and retain construction workspace through intermediate spills while final aggregate replay releases it. Tighten the two spill assertions and align new aggregate memory-call snapshots with strict output construction accounting without changing pressure limits.
|
Thanks @andygrove ! Merged main and resolved the conflicts in c466a17, tightened both assertions, and updated the accounting tests and snapshots for the separate output reservation. Memory limits, results, and rewrite-count assertions are unchanged. |

Which issue does this PR close?
Rationale for this change
Spill merges can retain exhausted input batches while reading replacements and allocate output buffers or nested-array scratch before reserving their memory. Wide List/Struct/View values can exceed the intended merge budget or fail despite individually admissible inputs.
What changes are included in this PR?
What is the testing strategy for this PR?
The complete affected physical-plan suite passed 2722 tests, including admission/retry/release, readback statistics, short-string Lists, null/empty children, View buffers, checked arithmetic, grant-preserving final-merge selection and multi-thread readers. Existing multi-batch backpressure tests check queue capacities one and two. The test-only signature correction reran all nine admission cases. A final header-only follow-up adopts current main's ASF SPDX format in three new files; CI-version license checking passes on those files with the current configuration. The full repository lint suite passed, including standard/all-feature Clippy, documentation, packaging and HTML checks.
The earlier workspace extended run passed 13,097 Rust tests and 542 SQL files. Those results are retained as historical evidence, distinct from the fresh affected-suite checks.
Three rotating local optimized non-LTO rounds use identical input, schema, output and budgets against pinned main and the original PR. Measurements precede the test-only parameter-packaging correction and a license-header-only follow-up; executable production code is unchanged in the final candidate. The fixed revisions are main
350ca76505c9, original PRc466a1780f88, and final candidate153da46efef5. The versioned synthetic eight-run fixture and View spill regressions provide deterministic inputs and correctness checks. The three-way harness and raw observations are not included in this PR; they can be shared on request. Initial-write+merge median paired ratios (candidate/original) are Struct 1.207x, List 1.135x, View 1.412x; SQL collect ratios are sort 1.167x, array_agg 1.243x. SQL collect includes initial writing/statistics and merge/replay. Sort and array_agg spill bytes match the original, and the eight-run List no longer needs an intermediate rewrite.These costs are material. The early-drain and byte-target rules produce 1024/1792/102 output batches for Struct/List/View, compared with 64 originally. View's short-run variation crosses parity, so this is not a stable production-throughput estimate. The rounds share one power state internally; measurements from earlier battery-powered candidates are not used to claim a direct improvement over this AC-powered round.
The eight-run fixture can be run with
cargo test --profile release-nonlto -p datafusion-physical-plan --lib progress_tests::measured_eight_run_merge -- --nocapture. Its built-in merge clock is distinct from the separate harness's initial-write+merge clock; the separate three-way harness uses the measured intervals reported above.Are there any user-facing changes?
Supported recursive spill merges reserve their construction space before allocating; insufficient single-row space returns a resource error. The guarantee stops at output construction; it does not cover the entire lifetime of outputs retained by consumers after delivery, whole-query RSS or all sort/aggregate memory. Write-side splitting (#25910) and aggregate spill batching (#25853) are separate changes; this PR does not establish their combined guarantees.