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

[jayzhan211]

Which issue does this PR close?

• Closes Memory should not blow up after Arrow IPC write-read round trip during spilling #17340.

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.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.30769% with 27 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.45%. Comparing base (1064661) to head (d677b5f).

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/spill/mod.rs	92.30%	13 Missing and 14 partials ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #24592    +/-   ##
========================================
  Coverage   81.44%   81.45%            
========================================
  Files        1118     1118            
  Lines      399602   399948   +346     
  Branches   399602   399948   +346     
========================================
+ Hits       325458   325773   +315     
- Misses      55146    55161    +15     
- Partials    18998    19014    +16     

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

[alamb]

Thank you @jayzhan211 -- the description on this PR is beautiful

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

I think another potentially user visible tradeoff of this PR is that it copies data into the right sized allocations, which does improve memory accounting, but at the cost of an extra memcpy for each byte read.

Can we add the test from #17340 that shows this improving memory accounting? I worry that without such a test we could easily lose this feature during a refactoring.

I also remember another recent PR about memory accounting trying to take into account shared buffers

• #24319

It seems like this function might already do what we want
https://docs.rs/datafusion/latest/datafusion/common/utils/memory/fn.get_record_batch_memory_size.html

Another thing claude code flagged is that this code doesn't seem to error on truncated input files (where on main it does)

On main, a spill file cut off mid-message leaves the partial bytes inside StreamDecoder's internal scratch, so the existing decoder.finish() call at EOF returns an error. On this PR, those partial bytes live in the MessageFramer instead
— and the framer's state is never checked at EOF.

Here is the reproducer

  #[tokio::test]
  async fn test_truncated_spill_file_errors() -> Result<()> {
      let batch = build_table_i32(
          ("a", &(0..100).collect::<Vec<_>>()),
          ("b", &(100..200).collect::<Vec<_>>()),
          ("c", &(200..300).collect::<Vec<_>>()),
      );
      let schema = batch.schema();
      let batches = vec![batch.clone(), batch];

      let env = Arc::new(RuntimeEnv::default());
      let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
      let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema));
      let spill_file = spill_manager
          .spill_record_batch_and_finish(&batches, "Test")?
          .unwrap();

      // Truncate the file mid-message: drop the 8-byte end-of-stream marker
      // plus part of the last batch's body.
      let path = spill_file.path().unwrap().to_path_buf();
      let len = std::fs::metadata(&path)?.len();
      let f = std::fs::OpenOptions::new().write(true).open(&path)?;
      f.set_len(len - 8 - 13)?;
      drop(f);

      let stream = spill_manager.read_spill_as_stream(spill_file, None)?;
      let result = collect(stream).await;
      assert!(
          result.is_err(),
          "truncated spill file should error, got Ok with {} batches",
          result.as_ref().map(|b| b.len()).unwrap_or(0)
      );
      Ok(())
  }

[alamb]

I think this line effectively copies each input byte twice (once into input and then once into body)

[jayzhan211]

Tests added, not yet find better way to eliminate the memcpy, and the cost seems low 🤔

[alamb]

Thank you @jayzhan211

This code seems good to me and well reasoned. However, since I don't use the spill writer all that much I don't really know how much, if at all, the extra copy will help/hurt of if it is a good tradeoff for better memory usage

Maybe @rluvaton or @2010YOUY01 who I know have spent more time on this might be able to weigh in as well

[jayzhan211]

Thanks for the review! I’ll wait another day to see if there’s any more input.