Skip to content

fix: frame IPC messages when reading spill files so batches pin only their own bytes - #24592

Open
jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:spill-reader-compact-small-batches
Open

fix: frame IPC messages when reading spill files so batches pin only their own bytes#24592
jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:spill-reader-compact-small-batches

Conversation

@jayzhan211

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

Record batches read back from a spill file retain far more memory than was recorded for them when they were written, which is what the Record batch memory usage (...) exceeds the expected limit warning in #17340 reports. Under a memory limit this is not just a noisy log: the multi-level merge in SortExec and the spilling aggregate size their merge fan-in from max_record_batch_memory recorded at spill time, so inflated read-back batches consume memory the operator never budgeted for.

The cause is in SpillReaderStream. It reads the file in 128 KB chunks and hands them straight to arrow's StreamDecoder, which builds arrays on slices of the buffer it is given — and a slice keeps its whole backing allocation alive. That goes wrong in two ways.

Small batches pin the whole chunk. With ~5 KB batches one chunk holds ~27 messages. Each decoded batch's buffers are slices of that chunk, so each batch retains, and is accounted for, 128 KB:

chunk (128 KB allocation)
┌──────┬──────┬──────┬─────┬───────┐
│ msg1 │ msg2 │ msg3 │ ... │ msg27 │
└──────┴──────┴──────┴─────┴───────┘
   ▲
   batch1's buffers slice here, yet keep all 128 KB alive

Traced in the sort tests: a 100-row Utf8 batch read back with caps=[(404, 131072), (4316, 131072)] — 4.7 KB of data, 128 KB retained, 27× what was recorded at spill time. And it is real retention, not just accounting: while the merge holds that batch, the other 26 in the chunk stay alive too.

Straddling batches double. A message spanning two chunks cannot be sliced, so the decoder gathers it into a Vec grown by doubling and the batch keeps the spare capacity:

chunk N                      chunk N+1
┌──────┬────────────────────┬──────────────┬────────┐
│ ...  │ msgK (first part)  │ msgK (rest)  │ msgK+1 │
└──────┴────────────────────┴──────────────┴────────┘

In the spill_io bench about half of all 256 KB batches came back in a 512 KB allocation (retained=523264 data=262144). How much spare capacity a straddling batch ends up with depends on where the chunk boundary fell, which is why the reports on #17340 range from ~10% over (967744 vs 877568) up to 2×.

What changes are included in this PR?

SpillReaderStream now reassembles each IPC message into allocations sized from the message's own headers before decoding, via a small MessageFramer state machine:

  1. read the 4-byte length prefix (skipping the continuation marker) → meta_len;
  2. fill a head Vec of exactly prefix + meta_len bytes;
  3. take bodyLength from the flatbuffer metadata (arrow_ipc::root_as_message);
  4. fill a body Vec::with_capacity(body_len) from however many chunks it spans;
  5. hand [head, body] to the same StreamDecoder.

The whole body is now inside one buffer whose allocation is exactly body_len, so the decoder takes its zero-copy path and the batch pins exactly its own message:

body for msg1 (5 KB)      body for msgK (256 KB)
┌──────┐                  ┌────────────────────┐
│ msg1 │ ◀── batch1       │ msgK               │ ◀── batchK
└──────┘                  └────────────────────┘

A 5 KB batch retains 5 KB and a 256 KB batch retains 256 KB, so max_record_batch_memory recorded at write time matches what the merge actually gets back.

This costs one memcpy per message — the one the decoder already paid for straddling messages — minus the doubling reallocation, so it is not slower. spill_io bench vs main (two runs on a quiet machine): StreamReader/read_100 −7.7%, q2/lz4_frame −5.9%, all other cases within noise.

Warnings from the #17340 check (RUST_LOG=datafusion_physical_plan::spill=debug):

main this PR
memory_limit::test_stringview_external_sort (the reproducer in #17340) 26 0
whole memory_limit integration suite ~5000 0
spilling_fuzz_in_memory_constrained_env + sort_fuzz + aggregate_fuzz 4170 0

Are these changes tested?

Yes:

  • test_read_back_does_not_inflate_batch_memory: spills 50 small batches (Int32, Utf8, Utf8View, List) and asserts every read-back batch's get_record_batch_memory_size is within the margin of the written maximum. Fails on main with read-back batch retains 131072 bytes, written max was 24196.
  • test_message_framer_across_chunk_boundaries: frames and decodes an IPC stream delivered in chunks of 1, 3, 7, 64, 1000 bytes and as a whole, checking the batches are intact and each retains no more than its own message body.
  • Existing spill, sort, aggregate, repartition unit tests, the memory_limit integration tests and the extended spilling fuzz suites pass.

Are there any user-facing changes?

No API changes. Queries that spill use less memory when reading spills back, and the spurious accounting warning from #17340 no longer fires.

…their own bytes

The spill reader fed raw 128 KB read chunks to arrow's zero-copy
StreamDecoder. A batch whose message fit inside a chunk kept the whole
chunk alive, and a message spanning chunks was gathered into a Vec grown
by doubling, so read-back batches retained, and were accounted for, up to
27x the memory recorded for them at spill time.

Reassemble each IPC message into an exactly sized allocation before
decoding, using the metadata's bodyLength to size the body buffer.

Closes apache#17340
@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Aug 23, 2026
@jayzhan211
jayzhan211 requested review from 2010YOUY01 and alamb August 23, 2026 12:04
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 92.95775% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.38%. Comparing base (5134a1a) to head (25a1866).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-plan/src/spill/mod.rs 92.95% 9 Missing and 6 partials ⚠️
Additional details and impacted files
@@           Coverage Diff            @@
##             main   #24592    +/-   ##
========================================
  Coverage   81.38%   81.38%            
========================================
  Files        1116     1116            
  Lines      397960   398181   +221     
  Branches   397960   398181   +221     
========================================
+ Hits       323880   324071   +191     
- Misses      55120    55139    +19     
- Partials    18960    18971    +11     

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Memory should not blow up after Arrow IPC write-read round trip during spilling

2 participants