From 3b4425cb4b0c7e0688248f6c2571c08307c3fdb1 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Mon, 31 Aug 2026 16:09:02 +0700 Subject: [PATCH 1/3] perf: reuse per-partition scratch in the shuffle write path The multi-partition write path pays two allocation costs per partition per event: BufBatchWriter starts from an empty byte buffer and regrows it toward the 1MB write-buffer size, and the interleave iterator materializes the whole partition index list as a fresh usize-pair vector (16 bytes per row) plus a batch-ref vector over all buffered batches. The task-level writer now owns one byte buffer and recycles it through the sequential partition loops, drained between partitions with its capacity kept, so an event allocates the buffer once instead of once per partition and holds exactly one regardless of partition count. The interleave path builds the batch-ref list once per write cycle and converts indices per output chunk into a small reusable scratch instead of widening the full list up front. Wire format and produced batches are unchanged, pinned by byte-level tests. Write time drops ~27% on a 2000-partition zstd shuffle bench; the larger effect is allocator pressure under concurrent tasks, which a single-task bench understates. Part of #5002. --- .../src/partitioners/multi_partition.rs | 17 ++- .../partitioned_batch_iterator.rs | 118 +++++++++++++++--- native/shuffle/src/shuffle_writer.rs | 3 + .../shuffle/src/writers/buf_batch_writer.rs | 74 ++++++++++- .../writers/local/local_partition_writer.rs | 16 ++- native/shuffle/src/writers/local/spill.rs | 6 + 6 files changed, 211 insertions(+), 23 deletions(-) diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 37ce85d1cd2..c9c84dcd1c1 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -567,12 +567,17 @@ impl MultiPartitionShuffleRepartitioner { with_trace("shuffle_spill", self.tracing_enabled, || { let num_output_partitions = self.partition_indices.len(); let write_result = { - let mut partitioned_batches = self.partitioned_batches(); + let partitioned_batches = self.partitioned_batches(); + // Build the batch-ref slice once and share it across all partitions. + let batch_refs = partitioned_batches.batch_refs(); (0..num_output_partitions).try_for_each(|partition_id| { self.partition_writer.write( partition_id, - &mut partitioned_batches - .produce(partition_id, &self.metrics.interleave_time), + &mut partitioned_batches.produce( + &batch_refs, + partition_id, + &self.metrics.interleave_time, + ), &self.metrics, ) }) @@ -630,15 +635,17 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartition with_trace("shuffle_write", self.tracing_enabled, || { let start_time = Instant::now(); - let mut partitioned_batches = self.partitioned_batches(); + let partitioned_batches = self.partitioned_batches(); self.pinned_buffers.clear(); let num_output_partitions = self.partition_indices.len(); + // Build the batch-ref slice once and share it across all partitions. + let batch_refs = partitioned_batches.batch_refs(); #[allow(clippy::needless_range_loop)] for i in 0..num_output_partitions { self.partition_writer.finish_partition( i, - &mut partitioned_batches.produce(i, &self.metrics.interleave_time), + &mut partitioned_batches.produce(&batch_refs, i, &self.metrics.interleave_time), &self.metrics, )?; } diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs index f124d98ff24..1ac066b8d6b 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -42,14 +42,22 @@ impl PartitionedBatchesProducer { } } + /// References to all buffered batches. Build this once per write cycle and share it + /// across every partition's [`Self::produce`] call instead of rebuilding a fresh + /// `Vec<&RecordBatch>` over all buffered batches for each partition. + pub(super) fn batch_refs(&self) -> Vec<&RecordBatch> { + self.buffered_batches.iter().collect() + } + pub(super) fn produce<'a>( - &'a mut self, + &'a self, + refs: &'a [&'a RecordBatch], partition_id: usize, interleave_time: &'a Time, ) -> PartitionedBatchIterator<'a> { PartitionedBatchIterator::new( &self.partition_indices[partition_id], - &self.buffered_batches, + refs, self.batch_size, interleave_time, ) @@ -58,9 +66,13 @@ impl PartitionedBatchesProducer { /// Iterates over the shuffled record batches belonging to a single output partition. pub(crate) struct PartitionedBatchIterator<'a> { - record_batches: Vec<&'a RecordBatch>, + record_batches: &'a [&'a RecordBatch], batch_size: usize, - indices: Vec<(usize, usize)>, + indices: &'a [(u32, u32)], + /// Scratch for the current chunk's indices widened to what `interleave_record_batch` + /// expects. Reused across chunks so each partition costs one small allocation + /// (capacity at most `batch_size`) rather than re-materializing its whole index list. + chunk_scratch: Vec<(usize, usize)>, pos: usize, interleave_time: &'a Time, } @@ -68,29 +80,26 @@ pub(crate) struct PartitionedBatchIterator<'a> { impl<'a> PartitionedBatchIterator<'a> { fn new( indices: &'a [(u32, u32)], - buffered_batches: &'a [RecordBatch], + record_batches: &'a [&'a RecordBatch], batch_size: usize, interleave_time: &'a Time, ) -> Self { if indices.is_empty() { // Avoid unnecessary allocations when the partition is empty return Self { - record_batches: vec![], + record_batches: &[], batch_size, - indices: vec![], + indices: &[], + chunk_scratch: vec![], pos: 0, interleave_time, }; } - let record_batches = buffered_batches.iter().collect::>(); - let current_indices = indices - .iter() - .map(|(i_batch, i_row)| (*i_batch as usize, *i_row as usize)) - .collect::>(); Self { record_batches, batch_size, - indices: current_indices, + indices, + chunk_scratch: Vec::with_capacity(batch_size.min(indices.len())), pos: 0, interleave_time, } @@ -106,9 +115,14 @@ impl Iterator for PartitionedBatchIterator<'_> { } let indices_end = std::cmp::min(self.pos + self.batch_size, self.indices.len()); - let indices = &self.indices[self.pos..indices_end]; + self.chunk_scratch.clear(); + self.chunk_scratch.extend( + self.indices[self.pos..indices_end] + .iter() + .map(|(i_batch, i_row)| (*i_batch as usize, *i_row as usize)), + ); let mut timer = self.interleave_time.timer(); - let result = interleave_record_batch(&self.record_batches, indices); + let result = interleave_record_batch(self.record_batches, &self.chunk_scratch); timer.stop(); match result { Ok(batch) => { @@ -122,3 +136,77 @@ impl Iterator for PartitionedBatchIterator<'_> { } } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use std::sync::Arc; + + fn batches() -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + (0..3) + .map(|b| { + let values: Vec = (0..5).map(|r| b * 100 + r).collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap() + }) + .collect() + } + + /// Chunked index conversion must interleave exactly like converting the whole partition's + /// index list up front, including the short tail chunk, and share one batch-ref slice + /// across partitions. + #[test] + fn chunked_interleave_matches_full_conversion() { + let buffered = batches(); + let indices: Vec<(u32, u32)> = vec![ + (0, 0), + (2, 4), + (1, 1), + (0, 3), + (2, 0), + (1, 4), + (0, 1), + (2, 2), + (1, 0), + (0, 4), + ]; + let batch_size = 4; // chunks of 4, 4, and a tail of 2 + let producer = PartitionedBatchesProducer::new( + buffered.clone(), + vec![indices.clone(), Vec::new()], + batch_size, + ); + let refs = producer.batch_refs(); + let time = Time::default(); + + let produced: Vec = producer + .produce(&refs, 0, &time) + .collect::>() + .unwrap(); + + let expected_refs: Vec<&RecordBatch> = buffered.iter().collect(); + let full: Vec<(usize, usize)> = indices + .iter() + .map(|(b, r)| (*b as usize, *r as usize)) + .collect(); + let expected: Vec = full + .chunks(batch_size) + .map(|chunk| interleave_record_batch(&expected_refs, chunk).unwrap()) + .collect(); + + assert_eq!(produced, expected); + assert_eq!(produced.last().unwrap().num_rows(), 2, "tail chunk"); + + let empty: Vec = producer + .produce(&refs, 1, &time) + .collect::>() + .unwrap(); + assert!(empty.is_empty()); + } +} diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 21d71dc1221..a58265f15e0 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -1170,6 +1170,7 @@ mod test { Cursor::new(&mut coalesced_output), 1024 * 1024, 8192, + Vec::new(), ); for batch in &small_batches { buf_writer.write(batch, &encode_time, &write_time).unwrap(); @@ -1186,6 +1187,7 @@ mod test { Cursor::new(&mut uncoalesced_output), 1024 * 1024, 1, + Vec::new(), ); for batch in &small_batches { buf_writer.write(batch, &encode_time, &write_time).unwrap(); @@ -1290,6 +1292,7 @@ mod test { Cursor::new(&mut output), 1024 * 1024, batch_size as usize, + Vec::new(), ); for batch in &inputs { buf_writer.write(batch, &encode_time, &write_time).unwrap(); diff --git a/native/shuffle/src/writers/buf_batch_writer.rs b/native/shuffle/src/writers/buf_batch_writer.rs index 55d88a4ba48..ea43ecb1e0b 100644 --- a/native/shuffle/src/writers/buf_batch_writer.rs +++ b/native/shuffle/src/writers/buf_batch_writer.rs @@ -47,16 +47,21 @@ pub(crate) struct BufBatchWriter, W: Write> { } impl, W: Write> BufBatchWriter { + /// `buffer` is the caller-owned byte buffer to serialize into. Passing a buffer recovered + /// from a previous writer's [`Self::into_buffer`] reuses its capacity instead of regrowing + /// a fresh allocation toward `buffer_max_size` for every writer. pub(crate) fn new( shuffle_block_writer: S, writer: W, buffer_max_size: usize, batch_size: usize, + mut buffer: Vec, ) -> Self { + buffer.clear(); Self { shuffle_block_writer, writer, - buffer: vec![], + buffer, buffer_max_size, compression_context: CompressionContext::default(), coalescer: None, @@ -64,6 +69,12 @@ impl, W: Write> BufBatchWriter { } } + /// Recover the byte buffer so its capacity can be recycled into the next writer. + /// The buffer is drained (empty) if the writer was flushed. + pub(crate) fn into_buffer(self) -> Vec { + self.buffer + } + pub(crate) fn write( &mut self, batch: &RecordBatch, @@ -160,3 +171,64 @@ impl, W: Write + Seek> BufBatchWriter { self.writer.stream_position().map_err(Into::into) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{read_ipc_compressed, CompressionCodec}; + use arrow::array::Int64Array; + use arrow::datatypes::{DataType, Field, Schema}; + use std::sync::Arc; + + fn test_batch(seed: i64) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let values: Vec = (0..100).map(|i| seed * 1_000 + i).collect(); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(values))]).unwrap() + } + + fn write_one_partition(seed: i64, buffer: Vec) -> (Vec, Vec) { + let batch = test_batch(seed); + let block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::Zstd(1)) + .unwrap(); + let mut output = Vec::new(); + let time = Time::default(); + let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192, buffer); + writer.write(&batch, &time, &time).unwrap(); + writer.flush(&time, &time).unwrap(); + let buffer = writer.into_buffer(); + (output, buffer) + } + + /// A buffer recycled across partitions must produce byte-identical output to fresh + /// per-partition buffers, come back empty, and keep its grown capacity. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call zstd's C FFI + fn recycled_buffer_matches_fresh_buffers_and_keeps_capacity() { + let fresh: Vec> = (0..3) + .map(|p| write_one_partition(p, Vec::new()).0) + .collect(); + + let mut scratch = Vec::new(); + let mut recycled = Vec::new(); + for p in 0..3 { + let (output, returned) = write_one_partition(p, scratch); + assert!( + returned.is_empty(), + "recycled buffer must come back drained" + ); + scratch = returned; + recycled.push(output); + } + + assert_eq!(fresh, recycled); + assert!( + scratch.capacity() > 0, + "capacity grown in one partition must survive into the next" + ); + for output in &recycled { + let decoded = read_ipc_compressed(&output[16..]).unwrap(); + assert_eq!(decoded.num_rows(), 100); + } + } +} diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 3a9a6484dba..cb20515638f 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -53,6 +53,11 @@ enum DataOutput { spill_writers: Vec, /// Runtime used to allocate the temporary spill files. runtime: Arc, + /// Byte buffer recycled through the short-lived per-partition `BufBatchWriter`s. + /// Partitions are written strictly one at a time, so a single buffer keeps its + /// grown capacity across the whole task instead of every partition regrowing a + /// fresh allocation toward the write buffer size. + recycled_buffer: Vec, }, } @@ -99,6 +104,7 @@ impl LocalPartitionWriter { output_file, write_buffer_size, batch_size, + Vec::new(), )) } else { let output_writer = BufWriter::with_capacity(write_buffer_size, output_file); @@ -116,6 +122,7 @@ impl LocalPartitionWriter { shuffle_block_writer, spill_writers, runtime, + recycled_buffer: Vec::new(), } }; Ok(Self { @@ -170,12 +177,13 @@ impl PartitionWriter for LocalPartitionWriter { DataOutput::Multi { spill_writers, runtime, + recycled_buffer, .. } => { // Multi-partition output buffers each partition's batches into its own // spill file. `finish_partition` later merges the spill files (and any // remaining in-memory batches) into the shuffle output in partition order. - spill_writers[pid].write(iter, runtime, metrics)?; + spill_writers[pid].write(iter, runtime, metrics, recycled_buffer)?; } } @@ -216,6 +224,7 @@ impl PartitionWriter for LocalPartitionWriter { output_writer, shuffle_block_writer, spill_writers, + recycled_buffer, .. } => { self.offsets[pid] = output_writer.stream_position()?; @@ -234,18 +243,21 @@ impl PartitionWriter for LocalPartitionWriter { } // Write in memory batches to output data file. Each partition uses its - // own writer so coalescing does not cross partition boundaries. + // own writer so coalescing does not cross partition boundaries, but the + // byte buffer is recycled so its capacity carries over to the next one. let mut buf_batch_writer = BufBatchWriter::new( shuffle_block_writer, output_writer, write_buffer_size, batch_size, + std::mem::take(recycled_buffer), ); for batch in iter.by_ref() { let batch = batch?; buf_batch_writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; } buf_batch_writer.flush(&metrics.encode_time, &metrics.write_time)?; + *recycled_buffer = buf_batch_writer.into_buffer(); } } Ok(()) diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index d8534b44564..6c8b6a004ae 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -50,11 +50,15 @@ impl SpillWriter { }) } + /// `recycled_buffer` is a scratch byte buffer shared by the sequential per-partition + /// spill writes; it is drained and handed back on return so one buffer's capacity + /// serves every partition instead of each write regrowing its own. pub(crate) fn write>>( &mut self, iter: &mut I, runtime: &RuntimeEnv, metrics: &ShufflePartitionerMetrics, + recycled_buffer: &mut Vec, ) -> datafusion::common::Result<()> { if let Some(batch) = iter.next() { self.ensure_spill_file_created(runtime)?; @@ -65,6 +69,7 @@ impl SpillWriter { &mut self.spill_file.as_mut().unwrap().file, self.write_buffer_size, self.batch_size, + std::mem::take(recycled_buffer), ); let initial_position = buf_batch_writer.writer_stream_position()?; buf_batch_writer.write(&batch?, &metrics.encode_time, &metrics.write_time)?; @@ -76,6 +81,7 @@ impl SpillWriter { let bytes_written = buf_batch_writer .writer_stream_position()? .saturating_sub(initial_position); + *recycled_buffer = buf_batch_writer.into_buffer(); usize::try_from(bytes_written).map_err(|_| { DataFusionError::Execution(format!( "Spill file byte count exceeds platform capacity: {bytes_written}" From 08ec7be48a93bb0d7cabf231dffbf067f1a6ae2d Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Tue, 1 Sep 2026 10:16:26 +0700 Subject: [PATCH 2/3] refactor: borrow shuffle write scratch per call and cap retained capacity The write scratch is now threaded through write and flush as a borrow instead of moving through the constructor, so one convention covers all task-scoped shuffle scratch, and an error mid-partition can no longer strand the buffer inside a dropped writer. The writer asserts a fresh scratch is drained and that every call passes the same buffer, since a swapped buffer would silently lose unflushed bytes. Flush shrinks the scratch back to the configured write-buffer size, bounding retained capacity that could previously reach the buffer size plus the largest block. The producer asserts the batch-ref slice covers every buffered batch, and the criterion suite gains a high-partition group so partition scaling stays visible to benchmark checks. --- native/shuffle/benches/shuffle_writer.rs | 29 +++ .../partitioned_batch_iterator.rs | 21 ++ native/shuffle/src/shuffle_writer.rs | 30 ++- .../shuffle/src/writers/buf_batch_writer.rs | 188 ++++++++++++++---- .../writers/local/local_partition_writer.rs | 54 +++-- native/shuffle/src/writers/local/spill.rs | 26 ++- 6 files changed, 275 insertions(+), 73 deletions(-) diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index 2e82f6fc532..cd3f9cdb69b 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -174,6 +174,35 @@ fn criterion_benchmark(c: &mut Criterion) { }, ); } + group.finish(); + + // High partition counts stress the per-partition write path (one short-lived + // buffered writer per partition), which low counts barely exercise; compression + // is disabled to isolate it. Few samples: each iteration is a full end-to-end + // write across thousands of partitions. + let mut high_partition_group = c.benchmark_group("shuffle_writer_high_partition"); + high_partition_group.sample_size(10); + for num_partitions in [200usize, 2000, 8000] { + high_partition_group.bench_function( + format!("shuffle_writer: end to end (partitions={num_partitions}, compression=None)"), + |b| { + let ctx = SessionContext::new(); + let exec = create_shuffle_writer_exec( + CompressionCodec::None, + CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], num_partitions), + 8192, + 10, + ); + b.iter(|| { + let task_ctx = ctx.task_ctx(); + let stream = exec.execute(0, task_ctx).unwrap(); + let rt = Runtime::new().unwrap(); + rt.block_on(collect(stream)).unwrap(); + }); + }, + ); + } + high_partition_group.finish(); } fn create_shuffle_writer_exec( diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs index 1ac066b8d6b..424fa827d52 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -55,6 +55,13 @@ impl PartitionedBatchesProducer { partition_id: usize, interleave_time: &'a Time, ) -> PartitionedBatchIterator<'a> { + // Partition indices index into `buffered_batches`; a refs slice built from a + // different producer would silently interleave wrong rows. + debug_assert_eq!( + refs.len(), + self.buffered_batches.len(), + "refs slice must cover every buffered batch" + ); PartitionedBatchIterator::new( &self.partition_indices[partition_id], refs, @@ -209,4 +216,18 @@ mod tests { .unwrap(); assert!(empty.is_empty()); } + + /// A refs slice that does not cover every buffered batch (e.g. built from a different + /// producer) must fail fast in debug builds instead of interleaving wrong rows. + #[cfg(debug_assertions)] + #[test] + #[should_panic(expected = "refs slice must cover every buffered batch")] + fn produce_rejects_mismatched_refs() { + let buffered = batches(); + let producer = PartitionedBatchesProducer::new(buffered, vec![vec![(0, 0), (2, 1)]], 4); + let refs = producer.batch_refs(); + let truncated = &refs[..refs.len() - 1]; + let time = Time::default(); + let _ = producer.produce(truncated, 0, &time); + } } diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index a58265f15e0..c9b91b344ce 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -1170,12 +1170,16 @@ mod test { Cursor::new(&mut coalesced_output), 1024 * 1024, 8192, - Vec::new(), ); + let mut scratch = Vec::new(); for batch in &small_batches { - buf_writer.write(batch, &encode_time, &write_time).unwrap(); + buf_writer + .write(batch, &mut scratch, &encode_time, &write_time) + .unwrap(); } - buf_writer.flush(&encode_time, &write_time).unwrap(); + buf_writer + .flush(&mut scratch, &encode_time, &write_time) + .unwrap(); } // Write without coalescing (batch_size=1) @@ -1187,12 +1191,16 @@ mod test { Cursor::new(&mut uncoalesced_output), 1024 * 1024, 1, - Vec::new(), ); + let mut scratch = Vec::new(); for batch in &small_batches { - buf_writer.write(batch, &encode_time, &write_time).unwrap(); + buf_writer + .write(batch, &mut scratch, &encode_time, &write_time) + .unwrap(); } - buf_writer.flush(&encode_time, &write_time).unwrap(); + buf_writer + .flush(&mut scratch, &encode_time, &write_time) + .unwrap(); } // Coalesced output should be smaller due to fewer IPC schema blocks @@ -1292,12 +1300,16 @@ mod test { Cursor::new(&mut output), 1024 * 1024, batch_size as usize, - Vec::new(), ); + let mut scratch = Vec::new(); for batch in &inputs { - buf_writer.write(batch, &encode_time, &write_time).unwrap(); + buf_writer + .write(batch, &mut scratch, &encode_time, &write_time) + .unwrap(); } - buf_writer.flush(&encode_time, &write_time).unwrap(); + buf_writer + .flush(&mut scratch, &encode_time, &write_time) + .unwrap(); } let blocks = read_all_ipc_batches(&output); diff --git a/native/shuffle/src/writers/buf_batch_writer.rs b/native/shuffle/src/writers/buf_batch_writer.rs index ea43ecb1e0b..a894adaead6 100644 --- a/native/shuffle/src/writers/buf_batch_writer.rs +++ b/native/shuffle/src/writers/buf_batch_writer.rs @@ -24,8 +24,13 @@ use std::borrow::Borrow; use std::io::{Cursor, Seek, SeekFrom, Write}; /// Write batches to writer while using a buffer to avoid frequent system calls. -/// The record batches were first written by ShuffleBlockWriter into an internal buffer. -/// Once the buffer exceeds the max size, the buffer will be flushed to the writer. +/// The record batches are first written by ShuffleBlockWriter into a caller-provided +/// scratch buffer. Once the scratch exceeds the max size, it is flushed to the writer. +/// +/// The scratch buffer is borrowed per call rather than owned: task-scoped scratch is +/// threaded through every `write`/`flush`, so one buffer serves all the short-lived +/// writers of a task, and an error mid-partition cannot strand the buffer inside a +/// dropped writer and silently end recycling. /// /// Small batches are coalesced using Arrow's [`BatchCoalescer`] before serialization, reducing /// per-block IPC schema overhead. Output batches hold at least `batch_size` rows, apart from the @@ -36,7 +41,6 @@ use std::io::{Cursor, Seek, SeekFrom, Write}; pub(crate) struct BufBatchWriter, W: Write> { shuffle_block_writer: S, writer: W, - buffer: Vec, buffer_max_size: usize, compression_context: CompressionContext, /// Coalesces small batches into target_batch_size before serialization. @@ -44,43 +48,70 @@ pub(crate) struct BufBatchWriter, W: Write> { coalescer: Option, /// Target batch size for coalescing batch_size: usize, + /// Address of the scratch `Vec` seen on first use; every later call must pass the same + /// one, or unflushed bytes in the other buffer would be silently abandoned. + #[cfg(debug_assertions)] + scratch_addr: Option, } impl, W: Write> BufBatchWriter { - /// `buffer` is the caller-owned byte buffer to serialize into. Passing a buffer recovered - /// from a previous writer's [`Self::into_buffer`] reuses its capacity instead of regrowing - /// a fresh allocation toward `buffer_max_size` for every writer. pub(crate) fn new( shuffle_block_writer: S, writer: W, buffer_max_size: usize, batch_size: usize, - mut buffer: Vec, ) -> Self { - buffer.clear(); Self { shuffle_block_writer, writer, - buffer, buffer_max_size, compression_context: CompressionContext::default(), coalescer: None, batch_size, + #[cfg(debug_assertions)] + scratch_addr: None, } } - /// Recover the byte buffer so its capacity can be recycled into the next writer. - /// The buffer is drained (empty) if the writer was flushed. - pub(crate) fn into_buffer(self) -> Vec { - self.buffer + /// A fresh writer must start from a drained scratch (stale bytes from a previous owner + /// would be silently prepended to its first block), and every later call must pass the + /// same scratch (bytes left unflushed in a swapped-out buffer would be silently lost). + /// Identity is the `Vec`'s own address, which is stable for the caller-owned field the + /// writer is used with, unlike the data pointer that moves on regrowth. + #[allow(unused_variables)] + #[allow(clippy::ptr_arg)] // identity check needs the Vec's own address, not a slice view + fn check_scratch(&mut self, scratch: &Vec) { + #[cfg(debug_assertions)] + { + let addr = scratch as *const Vec as usize; + match self.scratch_addr { + None => { + debug_assert!( + scratch.is_empty(), + "fresh BufBatchWriter handed a non-empty scratch buffer ({} bytes)", + scratch.len() + ); + self.scratch_addr = Some(addr); + } + Some(previous) => debug_assert_eq!( + previous, addr, + "BufBatchWriter must receive the same scratch buffer on every call" + ), + } + } } + /// `scratch` is the caller-owned byte buffer to serialize into; threading the same + /// buffer through every call reuses its capacity instead of regrowing a fresh + /// allocation toward `buffer_max_size` for every writer. pub(crate) fn write( &mut self, batch: &RecordBatch, + scratch: &mut Vec, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { + self.check_scratch(scratch); let batch_size = self.batch_size; let coalescer = self.coalescer.get_or_insert_with(|| { // Enable BatchCoalescer's zero-copy passthrough for batches that are already big @@ -107,19 +138,20 @@ impl, W: Write> BufBatchWriter { let mut bytes_written = 0; for batch in &completed { - bytes_written += self.write_batch_to_buffer(batch, encode_time, write_time)?; + bytes_written += self.write_batch_to_buffer(batch, scratch, encode_time, write_time)?; } Ok(bytes_written) } - /// Serialize a single batch into the byte buffer, flushing to the writer if needed. + /// Serialize a single batch into the scratch buffer, flushing to the writer if needed. fn write_batch_to_buffer( &mut self, batch: &RecordBatch, + scratch: &mut Vec, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { - let mut cursor = Cursor::new(&mut self.buffer); + let mut cursor = Cursor::new(&mut *scratch); cursor.seek(SeekFrom::End(0))?; let bytes_written = self.shuffle_block_writer.borrow().write_batch( batch, @@ -130,18 +162,21 @@ impl, W: Write> BufBatchWriter { let pos = cursor.position(); if pos >= self.buffer_max_size as u64 { let mut write_timer = write_time.timer(); - self.writer.write_all(&self.buffer)?; + self.writer.write_all(scratch)?; write_timer.stop(); - self.buffer.clear(); + scratch.clear(); } Ok(bytes_written) } + /// Flushes buffered rows and bytes; `scratch` is left drained for the next writer. pub(crate) fn flush( &mut self, + scratch: &mut Vec, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result<()> { + self.check_scratch(scratch); // Finish any remaining buffered rows in the coalescer let mut remaining = Vec::new(); if let Some(coalescer) = &mut self.coalescer { @@ -151,17 +186,20 @@ impl, W: Write> BufBatchWriter { } } for batch in &remaining { - self.write_batch_to_buffer(batch, encode_time, write_time)?; + self.write_batch_to_buffer(batch, scratch, encode_time, write_time)?; } - // Flush the byte buffer to the underlying writer + // Flush the scratch buffer to the underlying writer let mut write_timer = write_time.timer(); - if !self.buffer.is_empty() { - self.writer.write_all(&self.buffer)?; + if !scratch.is_empty() { + self.writer.write_all(scratch)?; } self.writer.flush()?; write_timer.stop(); - self.buffer.clear(); + scratch.clear(); + // The scratch's high-water mark can reach `buffer_max_size` plus the largest block + // that crossed the threshold; keep only the configured buffer size across reuses. + scratch.shrink_to(self.buffer_max_size); Ok(()) } } @@ -186,38 +224,36 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(values))]).unwrap() } - fn write_one_partition(seed: i64, buffer: Vec) -> (Vec, Vec) { + fn write_one_partition(seed: i64, scratch: &mut Vec) -> Vec { let batch = test_batch(seed); let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::Zstd(1)) .unwrap(); let mut output = Vec::new(); let time = Time::default(); - let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192, buffer); - writer.write(&batch, &time, &time).unwrap(); - writer.flush(&time, &time).unwrap(); - let buffer = writer.into_buffer(); - (output, buffer) + let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192); + writer.write(&batch, scratch, &time, &time).unwrap(); + writer.flush(scratch, &time, &time).unwrap(); + output } - /// A buffer recycled across partitions must produce byte-identical output to fresh - /// per-partition buffers, come back empty, and keep its grown capacity. + /// A scratch buffer recycled across partitions must produce byte-identical output to + /// fresh per-partition buffers, come back drained, and keep its grown capacity. #[test] #[cfg_attr(miri, ignore)] // miri can't call zstd's C FFI - fn recycled_buffer_matches_fresh_buffers_and_keeps_capacity() { + fn recycled_scratch_matches_fresh_buffers_and_keeps_capacity() { let fresh: Vec> = (0..3) - .map(|p| write_one_partition(p, Vec::new()).0) + .map(|p| write_one_partition(p, &mut Vec::new())) .collect(); let mut scratch = Vec::new(); let mut recycled = Vec::new(); for p in 0..3 { - let (output, returned) = write_one_partition(p, scratch); + let output = write_one_partition(p, &mut scratch); assert!( - returned.is_empty(), - "recycled buffer must come back drained" + scratch.is_empty(), + "recycled scratch must come back drained" ); - scratch = returned; recycled.push(output); } @@ -231,4 +267,82 @@ mod tests { assert_eq!(decoded.num_rows(), 100); } } + + /// Handing a non-empty scratch to a fresh writer would silently prepend stale bytes + /// to the first block; debug builds must catch it. + #[cfg(debug_assertions)] + #[test] + #[should_panic(expected = "non-empty scratch")] + fn fresh_writer_rejects_dirty_scratch() { + let batch = test_batch(0); + let block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut output = Vec::new(); + let time = Time::default(); + let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192); + let mut dirty = vec![0xAB, 0xCD]; + let _ = writer.write(&batch, &mut dirty, &time, &time); + } + + /// Swapping in a different scratch mid-writer would silently abandon any bytes still + /// buffered in the first one; the identity check has to catch it in debug builds. + #[test] + #[cfg(debug_assertions)] + #[should_panic(expected = "same scratch buffer")] + fn writer_rejects_swapped_scratch() { + let batch = test_batch(0); + let block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut output = Vec::new(); + let time = Time::default(); + let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192); + let mut first = Vec::new(); + writer.write(&batch, &mut first, &time, &time).unwrap(); + let mut second = Vec::new(); + let _ = writer.write(&batch, &mut second, &time, &time); + } + + /// A block that crosses `buffer_max_size` grows the scratch past the cap; `flush` + /// must shrink retained capacity back to the configured buffer size, while a + /// normally-sized run keeps its (sub-cap) capacity untouched. + #[test] + fn flush_caps_retained_scratch_capacity() { + let batch = test_batch(0); // 100 rows of Int64: block is far larger than 64 bytes + let buffer_max_size = 64usize; + // batch_size below the row count so the batch bypasses the coalescer and is + // serialized into the scratch during `write`. + let batch_size = 10usize; + let block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut output = Vec::new(); + let time = Time::default(); + let mut scratch = Vec::new(); + let mut writer = + BufBatchWriter::new(block_writer, &mut output, buffer_max_size, batch_size); + writer.write(&batch, &mut scratch, &time, &time).unwrap(); + assert!( + scratch.capacity() > buffer_max_size, + "oversized block must have grown the scratch past the cap" + ); + writer.flush(&mut scratch, &time, &time).unwrap(); + assert!(scratch.is_empty()); + assert!( + scratch.capacity() <= buffer_max_size, + "retained capacity {} exceeds cap {}", + scratch.capacity(), + buffer_max_size + ); + + // With a roomy cap the grown capacity is retained (shrink_to never grows the + // target below the cap, so no over-shrinking). + let large_cap = 1 << 20; + let block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut output = Vec::new(); + let mut scratch = Vec::new(); + let mut writer = BufBatchWriter::new(block_writer, &mut output, large_cap, 8192); + writer.write(&batch, &mut scratch, &time, &time).unwrap(); + writer.flush(&mut scratch, &time, &time).unwrap(); + assert!(scratch.capacity() > 0 && scratch.capacity() <= large_cap); + } } diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index cb20515638f..400f3e754b9 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -42,7 +42,12 @@ use std::sync::Arc; /// boundaries. They hold the raw output writer and block writer directly. enum DataOutput { /// Single-partition output: one long-lived writer streams all batches. - Single(BufBatchWriter), + Single { + writer: BufBatchWriter, + /// Task-scoped scratch byte buffer threaded through every call on the + /// long-lived writer, which borrows rather than owns its serialization buffer. + scratch: Vec, + }, /// Multi-partition output: batches are staged per partition and merged into /// `output_writer` one partition at a time during `finish_partition`. Multi { @@ -99,13 +104,15 @@ impl LocalPartitionWriter { .map_err(|e| DataFusionError::Execution(format!("shuffle write error: {e:?}")))?; let data_output = if num_output_partitions == 1 { - DataOutput::Single(BufBatchWriter::new( - shuffle_block_writer, - output_file, - write_buffer_size, - batch_size, - Vec::new(), - )) + DataOutput::Single { + writer: BufBatchWriter::new( + shuffle_block_writer, + output_file, + write_buffer_size, + batch_size, + ), + scratch: Vec::new(), + } } else { let output_writer = BufWriter::with_capacity(write_buffer_size, output_file); let spill_writers = (0..num_output_partitions) @@ -140,7 +147,7 @@ impl LocalPartitionWriter { pub(crate) fn get_spill_writers(&self) -> &Vec { match &self.data_output { DataOutput::Multi { spill_writers, .. } => spill_writers, - DataOutput::Single(_) => panic!("single-partition output has no spill writers"), + DataOutput::Single { .. } => panic!("single-partition output has no spill writers"), } } } @@ -156,7 +163,7 @@ impl PartitionWriter for LocalPartitionWriter { I: Iterator>, { match &mut self.data_output { - DataOutput::Single(writer) => { + DataOutput::Single { writer, scratch } => { if pid != 0 { return Err(DataFusionError::Execution( "LocalPartitionWriter single-partition output only supports partition 0." @@ -171,7 +178,7 @@ impl PartitionWriter for LocalPartitionWriter { // `finish_all`. for batch in iter.by_ref() { let batch = batch?; - writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + writer.write(&batch, scratch, &metrics.encode_time, &metrics.write_time)?; } } DataOutput::Multi { @@ -210,14 +217,14 @@ impl PartitionWriter for LocalPartitionWriter { let batch_size = self.batch_size; match &mut self.data_output { - DataOutput::Single(writer) => { + DataOutput::Single { writer, scratch } => { // Single-partition data was already streamed via `write`, starting at // offset 0 (already recorded in `self.offsets[0]`). Stream any trailing // batches (normally none) without flushing; the long-lived writer is // flushed once in `finish_all`. for batch in iter.by_ref() { let batch = batch?; - writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + writer.write(&batch, scratch, &metrics.encode_time, &metrics.write_time)?; } } DataOutput::Multi { @@ -244,20 +251,27 @@ impl PartitionWriter for LocalPartitionWriter { // Write in memory batches to output data file. Each partition uses its // own writer so coalescing does not cross partition boundaries, but the - // byte buffer is recycled so its capacity carries over to the next one. + // scratch buffer is shared so its capacity carries over to the next one. let mut buf_batch_writer = BufBatchWriter::new( shuffle_block_writer, output_writer, write_buffer_size, batch_size, - std::mem::take(recycled_buffer), ); for batch in iter.by_ref() { let batch = batch?; - buf_batch_writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.write( + &batch, + recycled_buffer, + &metrics.encode_time, + &metrics.write_time, + )?; } - buf_batch_writer.flush(&metrics.encode_time, &metrics.write_time)?; - *recycled_buffer = buf_batch_writer.into_buffer(); + buf_batch_writer.flush( + recycled_buffer, + &metrics.encode_time, + &metrics.write_time, + )?; } } Ok(()) @@ -270,8 +284,8 @@ impl PartitionWriter for LocalPartitionWriter { // Flush the data output and capture the final position. For the // single-partition writer this also finalizes the last coalesced batch. let final_offset = match &mut self.data_output { - DataOutput::Single(writer) => { - writer.flush(&metrics.encode_time, &metrics.write_time)?; + DataOutput::Single { writer, scratch } => { + writer.flush(scratch, &metrics.encode_time, &metrics.write_time)?; writer.writer_stream_position()? } DataOutput::Multi { output_writer, .. } => { diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index 6c8b6a004ae..0859d198a5b 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -51,8 +51,8 @@ impl SpillWriter { } /// `recycled_buffer` is a scratch byte buffer shared by the sequential per-partition - /// spill writes; it is drained and handed back on return so one buffer's capacity - /// serves every partition instead of each write regrowing its own. + /// spill writes; it is left drained on return so one buffer's capacity serves every + /// partition instead of each write regrowing its own. pub(crate) fn write>>( &mut self, iter: &mut I, @@ -69,19 +69,31 @@ impl SpillWriter { &mut self.spill_file.as_mut().unwrap().file, self.write_buffer_size, self.batch_size, - std::mem::take(recycled_buffer), ); let initial_position = buf_batch_writer.writer_stream_position()?; - buf_batch_writer.write(&batch?, &metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.write( + &batch?, + recycled_buffer, + &metrics.encode_time, + &metrics.write_time, + )?; for batch in iter.by_ref() { let batch = batch?; - buf_batch_writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.write( + &batch, + recycled_buffer, + &metrics.encode_time, + &metrics.write_time, + )?; } - buf_batch_writer.flush(&metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.flush( + recycled_buffer, + &metrics.encode_time, + &metrics.write_time, + )?; let bytes_written = buf_batch_writer .writer_stream_position()? .saturating_sub(initial_position); - *recycled_buffer = buf_batch_writer.into_buffer(); usize::try_from(bytes_written).map_err(|_| { DataFusionError::Execution(format!( "Spill file byte count exceeds platform capacity: {bytes_written}" From dd51509f1745c68a492f42c82f33908cff9f22c6 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Wed, 2 Sep 2026 17:36:09 +0700 Subject: [PATCH 3/3] fix: drain the write scratch when a partition errors An error mid-partition left encoded bytes in the recycled scratch buffer, and the next partition to reuse it would flush those bytes into its own range. Both call sites now clear the scratch before propagating the error. Also rewrites the capacity test to serialize for real and assert the exact capacity is preserved across a flush. --- .../shuffle/src/writers/buf_batch_writer.rs | 19 ++++- .../writers/local/local_partition_writer.rs | 81 ++++++++++++++++--- native/shuffle/src/writers/local/spill.rs | 52 +++++++++++- 3 files changed, 134 insertions(+), 18 deletions(-) diff --git a/native/shuffle/src/writers/buf_batch_writer.rs b/native/shuffle/src/writers/buf_batch_writer.rs index a894adaead6..47551d73d84 100644 --- a/native/shuffle/src/writers/buf_batch_writer.rs +++ b/native/shuffle/src/writers/buf_batch_writer.rs @@ -333,16 +333,27 @@ mod tests { buffer_max_size ); - // With a roomy cap the grown capacity is retained (shrink_to never grows the - // target below the cap, so no over-shrinking). + // With a roomy cap the grown capacity is retained: `write` serializes into the + // scratch (batch_size below the row count again), and `flush` must leave the + // sub-cap capacity exactly unchanged rather than shrinking it. let large_cap = 1 << 20; let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); let mut output = Vec::new(); let mut scratch = Vec::new(); - let mut writer = BufBatchWriter::new(block_writer, &mut output, large_cap, 8192); + let mut writer = BufBatchWriter::new(block_writer, &mut output, large_cap, batch_size); writer.write(&batch, &mut scratch, &time, &time).unwrap(); + let cap_after_write = scratch.capacity(); + assert!( + cap_after_write > 0 && cap_after_write <= large_cap, + "write must have serialized the batch into the scratch" + ); writer.flush(&mut scratch, &time, &time).unwrap(); - assert!(scratch.capacity() > 0 && scratch.capacity() <= large_cap); + assert!(scratch.is_empty()); + assert_eq!( + scratch.capacity(), + cap_after_write, + "flush must not shrink a scratch already under the cap" + ); } } diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 400f3e754b9..a7dae2b82ba 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -258,20 +258,25 @@ impl PartitionWriter for LocalPartitionWriter { write_buffer_size, batch_size, ); - for batch in iter.by_ref() { - let batch = batch?; - buf_batch_writer.write( - &batch, + let result: datafusion::common::Result<()> = (|| { + for batch in iter.by_ref() { + let batch = batch?; + buf_batch_writer.write( + &batch, + recycled_buffer, + &metrics.encode_time, + &metrics.write_time, + )?; + } + buf_batch_writer.flush( recycled_buffer, &metrics.encode_time, &metrics.write_time, - )?; - } - buf_batch_writer.flush( - recycled_buffer, - &metrics.encode_time, - &metrics.write_time, - )?; + ) + })(); + // An errored partition must hand back a drained buffer, or its bytes + // leak into the next partition's block. + result.inspect_err(|_| recycled_buffer.clear())?; } } Ok(()) @@ -320,3 +325,57 @@ impl PartitionWriter for LocalPartitionWriter { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::CompressionCodec; + use arrow::array::Int64Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + + /// A partition whose batch iterator fails after a batch was already encoded must hand + /// back a drained scratch; leftover bytes would land in the next partition's block. + #[test] + fn finish_partition_error_drains_recycled_buffer() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from_iter_values(0..100))], + ) + .unwrap(); + let block_writer = + ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::None).unwrap(); + let dir = tempfile::tempdir().unwrap(); + let mut writer = LocalPartitionWriter::try_new( + dir.path().join("data.out").to_str().unwrap().to_string(), + dir.path().join("index.out").to_str().unwrap().to_string(), + block_writer, + 2, + // batch_size below the row count so the write serializes into the scratch. + 10, + 1 << 20, + Arc::new(RuntimeEnv::default()), + ) + .unwrap(); + + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut iter = vec![ + Ok(batch), + Err(DataFusionError::Execution("injected failure".to_string())), + ] + .into_iter(); + + assert!(writer.finish_partition(0, &mut iter, &metrics).is_err()); + match &writer.data_output { + DataOutput::Multi { + recycled_buffer, .. + } => assert!( + recycled_buffer.is_empty(), + "errored partition left {} bytes in the recycled buffer", + recycled_buffer.len() + ), + DataOutput::Single { .. } => unreachable!("two partitions use the multi output"), + } + } +} diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index 0859d198a5b..771b4431e17 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -63,7 +63,7 @@ impl SpillWriter { if let Some(batch) = iter.next() { self.ensure_spill_file_created(runtime)?; - let total_bytes_written = { + let result = (|| { let mut buf_batch_writer = BufBatchWriter::new( &mut self.shuffle_block_writer, &mut self.spill_file.as_mut().unwrap().file, @@ -98,8 +98,11 @@ impl SpillWriter { DataFusionError::Execution(format!( "Spill file byte count exceeds platform capacity: {bytes_written}" )) - })? - }; + }) + })(); + // An errored spill must hand back a drained buffer, or its bytes leak into + // the next partition's block. + let total_bytes_written = result.inspect_err(|_| recycled_buffer.clear())?; metrics.spilled_bytes.add(total_bytes_written); } Ok(()) @@ -141,3 +144,46 @@ impl SpillWriter { self.spill_file.is_some() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::CompressionCodec; + use arrow::array::Int64Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use std::sync::Arc; + + /// A spill whose batch iterator fails after a batch was already encoded must hand + /// back a drained scratch; leftover bytes would land in the next partition's block. + #[test] + fn write_error_drains_recycled_buffer() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from_iter_values(0..100))], + ) + .unwrap(); + let block_writer = + ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::None).unwrap(); + // batch_size below the row count so the first write serializes into the scratch. + let mut spill = SpillWriter::try_new(block_writer, 1 << 20, 10).unwrap(); + let runtime = RuntimeEnv::default(); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut recycled = Vec::new(); + let mut iter = vec![ + Ok(batch), + Err(DataFusionError::Execution("injected failure".to_string())), + ] + .into_iter(); + + assert!(spill + .write(&mut iter, &runtime, &metrics, &mut recycled) + .is_err()); + assert!( + recycled.is_empty(), + "errored spill left {} bytes in the recycled buffer", + recycled.len() + ); + } +}