From 25a18665662d8d92bb5fc47db47bbc3971e41915 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 23 Aug 2026 18:43:39 +0800 Subject: [PATCH] fix: frame IPC messages when reading spill files so batches pin only 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 #17340 --- datafusion/physical-plan/src/spill/mod.rs | 388 +++++++++++++++++++++- 1 file changed, 381 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index addcf78d2df84..24f37b4a5b961 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -28,6 +28,7 @@ pub use datafusion_common::utils::memory::get_record_batch_memory_size; #[doc(hidden)] pub use spill_manager::SpillManager; +use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -46,7 +47,7 @@ use arrow::ipc::{ }; use arrow::record_batch::RecordBatch; use arrow_data::ArrayDataBuilder; -use arrow_ipc::CompressionType; +use arrow_ipc::{CompressionType, root_as_message}; use datafusion_common::Result; use datafusion_common::config::SpillCompression; @@ -72,9 +73,15 @@ struct SpillReaderStream { /// see `physical_plan/sort/multi_level_merge.rs`. max_record_batch_memory: Option, - /// Holds leftover bytes from a chunk when a batch is yielded early + /// Holds leftover bytes from a chunk not yet consumed by `framer` current_buffer: Buffer, + /// Assembles the chunks of `byte_stream` into exactly sized messages + framer: MessageFramer, + + /// Framed message buffers not yet consumed by `decoder` + pending: VecDeque, + /// Keeps the file alive until the stream is dropped _spill_file: Arc, @@ -84,6 +91,213 @@ struct SpillReaderStream { // Small margin allowed to accommodate slight memory accounting variation const SPILL_BATCH_MEMORY_MARGIN: usize = 4096; +/// Reassembles an IPC stream read in arbitrary chunks into one exactly sized +/// allocation per message, so that the zero-copy [`StreamDecoder`] produces +/// batches whose buffers pin only their own message. +/// +/// # Why +/// +/// The decoder builds arrays on slices of whatever [`Buffer`] it is given, +/// and a slice keeps its whole backing allocation alive. Fed with the raw +/// chunks of the byte stream (128 KB for a file backed spill) that goes +/// wrong in two ways. +/// +/// A message that fits inside a chunk pins the entire chunk. With ~5 KB +/// batches, one chunk holds ~27 of them, and each decoded batch retains, and +/// is accounted for, 128 KB: +/// +/// ```text +/// chunk (128 KB allocation) +/// ┌──────┬──────┬──────┬─────┬───────┐ +/// │ msg1 │ msg2 │ msg3 │ ... │ msg27 │ +/// └──────┴──────┴──────┴─────┴───────┘ +/// ▲ +/// batch1's buffers slice here, yet keep all 128 KB alive +/// ``` +/// +/// A message that spans two chunks cannot be sliced, so the decoder gathers +/// it into a `Vec` grown by doubling, and the batch keeps the spare +/// capacity: a 256 KB body typically lands in a 512 KB allocation. +/// +/// ```text +/// chunk N chunk N+1 +/// ┌──────┬────────────────────┬──────────────┬────────┐ +/// │ ... │ msgK (first part) │ msgK (rest) │ msgK+1 │ +/// └──────┴────────────────────┴──────────────┴────────┘ +/// ``` +/// +/// Either way a batch uses several times the memory recorded for it at +/// spill time, breaking the `max_record_batch_memory` budgeting that the +/// multi-level merge relies on. +/// +/// # How +/// +/// Each message is copied out of the chunks into allocations sized from its +/// own headers: a head buffer (the length prefix and flatbuffer metadata, +/// whose `bodyLength` gives the body size) and, when non-empty, a body +/// buffer of exactly that size. The decoder then zero-copies from the body +/// buffer, so a batch pins exactly its own message: +/// +/// ```text +/// body for msg1 (5 KB) body for msgK (256 KB) +/// ┌──────┐ ┌────────────────────┐ +/// │ msg1 │ ◀── batch1 │ msgK │ ◀── batchK +/// └──────┘ └────────────────────┘ +/// ``` +/// +/// This costs one copy per message, which the decoder already paid for +/// spanning messages, without the doubling reallocation. +struct MessageFramer { + state: FramerState, +} + +enum FramerState { + /// Reading the 4 byte continuation marker or metadata length. + Prefix { + head: Vec, + read: usize, + continuation: bool, + }, + /// Reading the flatbuffer metadata into `head`, which already holds the + /// prefix and is allocated for `head_len` bytes. + Metadata { + head: Vec, + metadata_start: usize, + head_len: usize, + }, + /// Reading the body into `body`, which is allocated for `body_len` bytes. + Body { + head: Vec, + body: Vec, + body_len: usize, + }, + /// The end-of-stream marker was read. + Finished, +} + +impl MessageFramer { + fn new() -> Self { + Self { + state: FramerState::prefix(), + } + } + + /// Consumes bytes from `input` until a message is complete or `input` is + /// exhausted, returning the buffers of a completed message. + fn push(&mut self, input: &mut Buffer) -> Result>> { + while !input.is_empty() { + match &mut self.state { + FramerState::Prefix { + head, + read, + continuation, + } => { + let to_read = input.len().min(4 - *read); + head.extend_from_slice(&input[..to_read]); + input.advance(to_read); + *read += to_read; + if *read < 4 { + continue; + } + let word: [u8; 4] = head[head.len() - 4..].try_into().unwrap(); + if !*continuation && word == CONTINUATION_MARKER { + *continuation = true; + *read = 0; + continue; + } + let metadata_len = u32::from_le_bytes(word) as usize; + let head = std::mem::take(head); + if metadata_len == 0 { + self.state = FramerState::Finished; + return Ok(Some(vec![Buffer::from_vec(head)])); + } + let metadata_start = head.len(); + let head_len = metadata_start + metadata_len; + let mut sized = Vec::with_capacity(head_len); + sized.extend_from_slice(&head); + self.state = FramerState::Metadata { + head: sized, + metadata_start, + head_len, + }; + } + FramerState::Metadata { + head, + metadata_start, + head_len, + } => { + let to_read = input.len().min(*head_len - head.len()); + head.extend_from_slice(&input[..to_read]); + input.advance(to_read); + if head.len() < *head_len { + continue; + } + let message = + root_as_message(&head[*metadata_start..]).map_err(|e| { + datafusion_common::exec_datafusion_err!( + "Invalid IPC message in spill file: {e}" + ) + })?; + let body_len = + usize::try_from(message.bodyLength()).map_err(|_| { + datafusion_common::exec_datafusion_err!( + "Invalid IPC message body length in spill file: {}", + message.bodyLength() + ) + })?; + let head = std::mem::take(head); + if body_len == 0 { + self.state = FramerState::prefix(); + return Ok(Some(vec![Buffer::from_vec(head)])); + } + self.state = FramerState::Body { + head, + body: Vec::with_capacity(body_len), + body_len, + }; + } + FramerState::Body { + head, + body, + body_len, + } => { + let to_read = input.len().min(*body_len - body.len()); + body.extend_from_slice(&input[..to_read]); + input.advance(to_read); + if body.len() < *body_len { + continue; + } + let (head, body) = (std::mem::take(head), std::mem::take(body)); + self.state = FramerState::prefix(); + return Ok(Some(vec![ + Buffer::from_vec(head), + Buffer::from_vec(body), + ])); + } + FramerState::Finished => { + return datafusion_common::exec_err!( + "Unexpected bytes after the end of the IPC stream in spill file" + ); + } + } + } + Ok(None) + } +} + +impl FramerState { + fn prefix() -> Self { + Self::Prefix { + head: Vec::with_capacity(8), + read: 0, + continuation: false, + } + } +} + +/// Marks a length prefix in the IPC stream format, see `arrow_ipc`. +const CONTINUATION_MARKER: [u8; 4] = [0xff; 4]; + impl SpillReaderStream { fn new( schema: SchemaRef, @@ -101,6 +315,8 @@ impl SpillReaderStream { max_record_batch_memory, is_done: false, current_buffer: Buffer::from(&[]), + framer: MessageFramer::new(), + pending: VecDeque::new(), _spill_file: spill_file, schema_validated: false, }) @@ -118,8 +334,13 @@ impl Stream for SpillReaderStream { } loop { - if !this.current_buffer.is_empty() { - match this.decoder.decode(&mut this.current_buffer) { + // Decode the framed messages first + if let Some(buffer) = this.pending.front_mut() { + if buffer.is_empty() { + this.pending.pop_front(); + continue; + } + match this.decoder.decode(buffer) { Ok(Some(batch)) => { // One-time schema validation on the first decoded batch. // The IPC stream embeds the writer's schema in its header; @@ -159,17 +380,31 @@ impl Stream for SpillReaderStream { return Poll::Ready(Some(Ok(batch))); } Ok(None) => { - // The chunk didn't form a complete message. Arrow consumed the partial bytes - // into its internal scratch pad, leaving our current_buffer completely empty. - // We do nothing and fall through to fetch more data. + // A schema or dictionary message, or the buffer was + // only part of a message. Carry on with the next one. } Err(e) => { this.is_done = true; return Poll::Ready(Some(Err(e.into()))); } } + continue; } + // Then frame the next message out of the current chunk + if !this.current_buffer.is_empty() { + match this.framer.push(&mut this.current_buffer) { + Ok(Some(buffers)) => this.pending.extend(buffers), + Ok(None) => {} + Err(e) => { + this.is_done = true; + return Poll::Ready(Some(Err(e))); + } + } + continue; + } + + // Finally fetch another chunk match futures::ready!(this.byte_stream.as_mut().poll_next(cx)) { Some(Ok(chunk)) => { this.current_buffer = Buffer::from(chunk); @@ -585,6 +820,145 @@ mod tests { Ok(()) } + /// Reading a spill file back must not inflate the batches' memory + /// footprint: the decoder is zero-copy, so without framing every small + /// batch would keep a whole read chunk alive, see [`MessageFramer`]. + #[tokio::test] + async fn test_read_back_does_not_inflate_batch_memory() -> Result<()> { + use arrow::array::{ListArray, StringViewArray}; + use arrow::buffer::OffsetBuffer; + + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + Field::new("v", DataType::Utf8View, true), + Field::new( + "l", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + false, + ), + ])); + + // Small batches: many fit in one 128 KB read chunk. + let batches: Vec = (0..50) + .map(|b| { + let n = 100; + let ints = Int32Array::from_iter_values((0..n).map(|i| b * n + i)); + let strs = StringArray::from_iter_values( + (0..n).map(|i| format!("string value number {i}")), + ); + let views = StringViewArray::from_iter((0..n).map(|i| { + (i % 3 != 0).then(|| format!("a longer view value {b}/{i}")) + })); + let values = Int32Array::from_iter_values(0..n * 2); + let offsets = + OffsetBuffer::from_lengths(std::iter::repeat_n(2, n as usize)); + let list = ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + offsets, + Arc::new(values), + None, + ); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(ints), + Arc::new(strs), + Arc::new(views), + Arc::new(list), + ], + ) + .unwrap() + }) + .collect(); + + let max_written = batches + .iter() + .map(get_record_batch_memory_size) + .max() + .unwrap(); + + 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(); + + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + let read_back = collect(stream).await?; + assert_eq!(read_back.len(), batches.len()); + + for (written, read) in batches.iter().zip(&read_back) { + assert_eq!(written, read); + let size = get_record_batch_memory_size(read); + assert!( + size <= max_written + SPILL_BATCH_MEMORY_MARGIN, + "read-back batch retains {size} bytes, written max was {max_written}" + ); + } + Ok(()) + } + + /// Frames an IPC stream delivered in chunks of `chunk_size` bytes and + /// decodes it, checking that every batch is intact and that its buffers + /// are backed by an allocation no larger than its own message body. + fn frame_and_decode(ipc: &[u8], chunk_size: usize, expected: &[RecordBatch]) { + let mut framer = MessageFramer::new(); + let mut decoder = StreamDecoder::new(); + let mut decoded = vec![]; + for chunk in ipc.chunks(chunk_size) { + let mut input = Buffer::from(chunk); + while !input.is_empty() { + let Some(buffers) = framer.push(&mut input).unwrap() else { + continue; + }; + let body_size = buffers.last().unwrap().len(); + for mut buffer in buffers { + while !buffer.is_empty() { + if let Some(batch) = decoder.decode(&mut buffer).unwrap() { + let retained = get_record_batch_memory_size(&batch); + assert!( + retained <= body_size, + "batch retains {retained} bytes for a {body_size} byte body" + ); + decoded.push(batch); + } + } + } + } + } + decoder.finish().unwrap(); + assert!(matches!(framer.state, FramerState::Finished)); + assert_eq!(decoded, expected); + } + + #[test] + fn test_message_framer_across_chunk_boundaries() { + let batches: Vec = (0..5) + .map(|b| { + let n = 10 * (b + 1); + build_table_i32( + ("a", &(0..n).collect::>()), + ("b", &(n..2 * n).collect::>()), + ("c", &(2 * n..3 * n).collect::>()), + ) + }) + .collect(); + let schema = batches[0].schema(); + + let mut ipc = vec![]; + let mut writer = StreamWriter::try_new(&mut ipc, &schema).unwrap(); + for batch in &batches { + writer.write(batch).unwrap(); + } + writer.finish().unwrap(); + + for chunk_size in [1, 3, 7, 64, 1000, ipc.len()] { + frame_and_decode(&ipc, chunk_size, &batches); + } + } + #[tokio::test] async fn test_batch_spill_and_read_dictionary_arrays() -> Result<()> { // See https://github.com/apache/datafusion/issues/4658