[Weijun-H]

Which issue does this PR close?

• Related to Clarify memory-limit guarantees for external-sort spilling #25183.

Rationale for this change

Queries that sort wide values under a memory limit can exceed their budget during spill merge or return ResourcesExhausted even when the spill inputs individually fit. Output construction must be admitted before Arrow allocates the output buffers.

What changes are included in this PR?

• Reserve spill-merge output-construction workspace before interleave. If the preferred fixed-width output does not fit, re-spill using the minimum spill-derived target; a single row that still cannot reserve its estimated output returns a memory error.
• Preserve the configured row batch size for fixed-width output when memory permits. Variable-width output uses a byte target and finds its largest supported prefix with a single forward scan.
• Use the same Arrow allocation model for spill admission and output estimates, including aligned bitmap and fixed-size-binary buffers. Target-enabled estimation failures now return before allocation instead of allocating first and growing the reservation afterward.
• Retain admitted workspace across ordinary sort output batches. Aggregate replay releases it after each batch and reuses the larger of replay and construction headroom. Unsupported schemas keep the existing unaccounted path.

What is the testing strategy for this PR?

Coverage includes fixed-width preferred/minimum targets, aligned validity/Boolean/fixed-size-binary buffers, byte-array and null outputs, exact prefix boundaries, the 7-byte failure / 8-byte success boundary for a non-null Int64 row, estimator failure classification and reservation release, unsupported-schema fallback, output-headroom re-spill, retained workspace across retries, multi-pass spill merge, exact output batch metrics, and constrained-memory sort fuzzing.

cargo fmt --all, datafusion-physical-plan all-target/all-feature clippy, all 2,331 datafusion-physical-plan library tests, release-mode builder tests, and targeted spill-merge tests pass locally.

Are there any user-facing changes?

Supported flat spill-merge schemas may emit more, smaller output batches when the preferred target cannot be reserved. Strict pre-allocation accounting may cause additional re-spill at tight memory limits, or fail a query that previously succeeded only because output allocation was not fully accounted. There are no public API or configuration changes.

This change does not guarantee reservation coverage while an asynchronous consumer retains an already-yielded output batch.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 89.90043% with 142 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.53%. Comparing base (6574a8c) to head (850a3b2).
⚠️ Report is 514 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/sorts/builder.rs	88.95%	71 Missing and 29 partials ⚠️
...usion/physical-plan/src/sorts/multi_level_merge.rs	91.64%	19 Missing and 16 partials ⚠️
...atafusion/physical-plan/src/spill/spill_manager.rs	88.88%	2 Missing and 1 partial ⚠️
datafusion/physical-plan/src/sorts/sort.rs	66.66%	0 Missing and 2 partials ⚠️
...afusion/physical-plan/src/sorts/streaming_merge.rs	93.93%	2 Missing ⚠️

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.
📢 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.

[Copilot]

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

Open (1)

• Round byte-array values-buffer estimate before allocation admission · New

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.

[andygrove]

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:

1. The byte target caps output rows at the largest spilled batch even for fixed-width schemas. test_sort_spill now emits 7 batches of at most 1500 rows instead of 2, and test_sort_fetch_memory_calculation emits 4 × 100 instead of 1 × 400. For fixed-width columns the output size of batch_size rows is known from the schema. Could the target be max(largest spilled batch, fixed-width bytes for batch_size rows), or apply only when variable-width columns are present?
2. The target comes from get_sliced_size, which counts validity as ceil(n / 8), but the estimate rounds validity up to 64 bytes. So an output batch identical to the largest input batch gets split: a nullable Int32 batch 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.
3. 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_spill pass. Could we keep them exact, and assert the new exact counts in the two tests where the output really changes?
4. largest_prefix_under_target rescans 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_size 1000) 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.
5. The try_grow after building in finish_record_batch can 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 a debug_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.

[Weijun-H]

Thanks for the detailed review! @andygrove Addressed all five points:

1. Restored fixed-width batch sizes, with matching reservations and a smaller target when memory is tight.
2. Matched spill targets to Arrow interleave allocation estimates.
3. Restored exact batch and metric assertions.
4. Replaced repeated variable-width prefix scans with incremental accumulation, preserving the fixed-width fast path.
5. Removed post-allocation reservation growth; estimation failures now fail before allocation.

Also addressed the nits and documented the stricter memory-accounting behavior. Ready for another look.

[andygrove]

Thanks @Weijun-H, I re-checked this at 850a3b2 and all five points from my earlier review are resolved.

1. Fixed-width output keeps the configured row count. Initial runs record batch_size as their row limit, so the target is max(largest spilled batch, fixed-width bytes for batch_size rows). In a CLI A/B against main (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).
2. 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).
3. The exact batch and metric assertions are back. The 32 sort and spill fuzz tests pass with --features extended_tests.
4. 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.
5. 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 in builder.rs, three hunks in multi_level_merge.rs, one in spill_manager.rs. main has 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.rs still has two looser assertions (first_batch.num_rows() <= rows and spill_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.

[Weijun-H]

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.