diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index addcf78d2df84..e3c4b2cf30f9c 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,240 @@ 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. Measured on the +/// `spill_io` benchmark, the copy is not visible end to end for batches at +/// or above the 128 KB read chunk, where skipping the doubling gather makes +/// framing the faster of the two, and costs a few percent for batches small +/// enough that several share a chunk, which are exactly the batches whose +/// accounting it fixes. +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) + } + + /// Checks that the stream ended on a message boundary. + /// + /// This mirrors [`StreamDecoder::finish`]: a stream is complete either + /// when the end-of-stream marker was read, or when it ends right after a + /// message. Framing moves the partial bytes of a truncated file out of + /// the decoder's own scratch space and into the framer, so without this + /// check a truncated spill file would silently decode as a shorter + /// stream instead of erroring. + fn finish(&self) -> Result<()> { + match &self.state { + FramerState::Finished + | FramerState::Prefix { + read: 0, + continuation: false, + .. + } => Ok(()), + _ => datafusion_common::exec_err!( + "Unexpected end of spill file: the IPC stream ends mid-message" + ), + } + } +} + +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 +342,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 +361,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 +407,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); @@ -181,6 +443,11 @@ impl Stream for SpillReaderStream { None => { this.is_done = true; + // The framer holds the bytes of an incomplete trailing + // message, so it is the one that detects truncation. + if let Err(e) = this.framer.finish() { + return Poll::Ready(Some(Err(e))); + } if let Err(e) = this.decoder.finish() { return Poll::Ready(Some(Err(e.into()))); } @@ -585,6 +852,361 @@ 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`]. + /// + /// Regression test for + /// . Without framing + /// these 50 batches read back at 131072 bytes each, against a 22992 byte + /// maximum recorded at spill time, and the whole 345 KB stream is + /// accounted for as 5.9 MB. Note that this is not a counting bug that a + /// shared-buffer aware counter could net out: holding only 5 of the 50 + /// batches still pins every read chunk, so the memory really is retained. + #[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(()) + } + + /// The workload from , + /// reduced to a single spilled run. + /// + /// `memory_limit::test_stringview_external_sort` sorts 200 batches of 1000 + /// random 100 byte strings under a 60 MB pool, producing around 10 spilled + /// runs of about 20 batches each. This spills one such run and reads it + /// back through the same [`SpillManager`] calls the external sorter and the + /// multi-level merge use: the spill records the largest batch's memory size, + /// and the merge reserves its budget from that number, so a batch that comes + /// back bigger is memory the merge is using without having reserved it. + /// + /// Each batch here is around 167 KB, larger than the 128 KB read chunk, so + /// every IPC message spans chunks. Without framing the decoder gathers such + /// a message into a `Vec` grown by doubling and the batch keeps the spare + /// capacity, so these batches read back at about 255 KB each against the + /// ~167 KB recorded at spill time - half again the budget the merge + /// reserved for them, which is the discrepancy the issue reports. + #[tokio::test] + async fn test_stringview_spill_read_back_memory_accounting() -> Result<()> { + use arrow::array::StringViewArray; + use rand::Rng; + + let schema = Arc::new(Schema::new(vec![ + Field::new("strings", DataType::Utf8View, false), + Field::new("random_numbers", DataType::Int32, false), + ])); + + // 100 random bytes per string, as in the issue's reproducer: long + // enough that the views point into a data buffer rather than inlining, + // and random so nothing dedupes. + let mut rng = rand::rng(); + let batches: Vec = (0..20) + .map(|_| { + let strings: Vec = (0..1000) + .map(|_| { + (0..100) + .map(|_| rng.random_range(0..=u8::MAX) as char) + .collect() + }) + .collect(); + let numbers: Vec = + (0..1000).map(|_| rng.random_range(0..=1000)).collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringViewArray::from(strings)) as ArrayRef, + Arc::new(Int32Array::from(numbers)) as ArrayRef, + ], + ) + .unwrap() + }) + .collect(); + + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema)); + + // The same call the external sorter spills a sorted run with: the + // returned size is what the merge budgets against. + let (spill_file, max_record_batch_memory) = spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + batches + .iter() + .map(Ok::<_, datafusion_common::DataFusionError>), + "Test", + )? + .unwrap(); + + let stream = spill_manager + .read_spill_as_stream(spill_file, Some(max_record_batch_memory))?; + 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_record_batch_memory + SPILL_BATCH_MEMORY_MARGIN, + "read-back batch retains {size} bytes, but the merge only \ + reserved for {max_record_batch_memory} bytes" + ); + } + 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); + } + } + + /// Decodes `ipc` the way the reader did before framing, returning the + /// number of batches or `Err` if the stream is not a complete one. + fn decode_unframed(ipc: &[u8]) -> std::result::Result { + let mut decoder = StreamDecoder::new(); + let mut buffer = Buffer::from(ipc); + let mut batches = 0; + while !buffer.is_empty() { + match decoder.decode(&mut buffer) { + Ok(Some(_)) => batches += 1, + Ok(None) => {} + Err(_) => return Err(()), + } + } + decoder.finish().map_err(|_| ())?; + Ok(batches) + } + + /// Same as [`decode_unframed`], but through the [`MessageFramer`]. + fn decode_framed(ipc: &[u8]) -> std::result::Result { + let mut framer = MessageFramer::new(); + let mut decoder = StreamDecoder::new(); + let mut input = Buffer::from(ipc); + let mut batches = 0; + while !input.is_empty() { + let Some(buffers) = framer.push(&mut input).map_err(|_| ())? else { + continue; + }; + for mut buffer in buffers { + while !buffer.is_empty() { + match decoder.decode(&mut buffer) { + Ok(Some(_)) => batches += 1, + Ok(None) => {} + Err(_) => return Err(()), + } + } + } + } + framer.finish().map_err(|_| ())?; + decoder.finish().map_err(|_| ())?; + Ok(batches) + } + + /// Framing must not weaken the reader's detection of a corrupt or cut + /// short spill file. The partial bytes of an incomplete trailing message + /// end up in the [`MessageFramer`] instead of the decoder's own scratch + /// space, so the framer has to report them, see [`MessageFramer::finish`]. + /// + /// Truncated at every possible offset, framed decoding must accept and + /// reject exactly what unframed decoding does. + #[test] + fn test_message_framer_reports_truncation_like_the_decoder() { + let batches: Vec = (0..3) + .map(|b| { + build_table_i32( + ("a", &(0..50).map(|i| b * 50 + i).collect::>()), + ("b", &(0..50).collect::>()), + ("c", &(0..50).collect::>()), + ) + }) + .collect(); + + let mut ipc = vec![]; + let mut writer = StreamWriter::try_new(&mut ipc, &batches[0].schema()).unwrap(); + for batch in &batches { + writer.write(batch).unwrap(); + } + writer.finish().unwrap(); + + let mut complete = 0; + for cut in 0..=ipc.len() { + let truncated = &ipc[..cut]; + let framed = decode_framed(truncated); + assert_eq!( + framed, + decode_unframed(truncated), + "framed and unframed decoding disagree on a stream cut to {cut} bytes" + ); + complete += usize::from(framed.is_ok()); + } + // Cutting anywhere but on a message boundary must be an error, so + // only the empty stream, the three batches and the end-of-stream + // marker may decode cleanly. + assert_eq!(complete, 5); + assert_eq!(decode_framed(&ipc), Ok(batches.len())); + } + + /// End-to-end: a spill file cut short mid-message must error rather than + /// silently read back as a shorter stream. + #[tokio::test] + async fn test_truncated_spill_file_errors() -> Result<()> { + let batch = build_table_i32( + ("a", &(0..100).collect::>()), + ("b", &(100..200).collect::>()), + ("c", &(200..300).collect::>()), + ); + 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(()) + } + #[tokio::test] async fn test_batch_spill_and_read_dictionary_arrays() -> Result<()> { // See https://github.com/apache/datafusion/issues/4658