Skip to content

fix: Bound recursive spill-merge output construction memory - #25565

Open
Weijun-H wants to merge 12 commits into
apache:mainfrom
Weijun-H:task823-upstream-byte-target
Open

Weijun-H wants to merge 12 commits into
apache:mainfrom
Weijun-H:task823-upstream-byte-target

Conversation

@Weijun-H

@Weijun-H Weijun-H commented Sep 21, 2026 •

Copy link
Copy Markdown
Member

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?

  • Admit estimator workspace and output construction before allocation for supported List/LargeList/Struct descendants and Utf8View/BinaryView. Shared View backing data stays covered by live input grants in pure spill merges; mixed View inputs retain their original fallback.
  • Record post-GC physical IPC read bounds separately from construction statistics, including compressed/aligned read buffers. Keep progress bounds through intermediate writes, splitting and replay admission.
  • Drain pending output before replacing recursive or variable-width inputs, bounding retained source batches. Try a shallower read queue only when the same budget can directly complete a measured recursive pure-spill merge; keep the checked construction guarantee and existing fallback on failure.

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 PR c466a1780f88, and final candidate 153da46efef5. 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.

Merge/collect workload Median paired ratio Three-pair range
Struct merge 1.809x 1.460–1.888x
List merge 1.147x 1.140–1.219x
View merge 1.447x 0.973–1.473x
Sort SQL 1.167x 1.154–1.202x
array_agg SQL 1.243x 1.203–1.255x

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.

@Weijun-H
Weijun-H force-pushed the task823-upstream-byte-target branch from 807afbc to 15c67a7 Compare September 21, 2026 11:55
@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 21, 2026
@github-actions github-actions Bot added the core Core DataFusion crate label Sep 24, 2026
@Weijun-H
Weijun-H force-pushed the task823-upstream-byte-target branch from b124f7b to f9c4788 Compare September 24, 2026 09:44
@Weijun-H Weijun-H changed the title Limit sort merge output batches by byte target Limit spill merge output batches by byte target Sep 24, 2026
@Weijun-H
Weijun-H marked this pull request as ready for review September 24, 2026 10:06
@Weijun-H
Weijun-H marked this pull request as draft September 24, 2026 10:09
@codecov-commenter

codecov-commenter commented Sep 24, 2026 •

Copy link
Copy Markdown

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.

@Weijun-H
Weijun-H marked this pull request as ready for review September 24, 2026 11:20
@Weijun-H
Weijun-H requested a balanced review from Copilot September 24, 2026 11:20

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Medium severity

Open (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.

Comment on lines +509 to +511
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 andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member Author

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 andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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.
@Weijun-H

Copy link
Copy Markdown
Member Author

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.

@Weijun-H
Weijun-H requested a review from andygrove October 26, 2026 03:39
@Weijun-H Weijun-H changed the title Limit spill merge output batches by byte target fix: Bound recursive spill-merge output construction memory Oct 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants