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/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..424fa827d52 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -42,14 +42,29 @@ 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> { + // 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], - &self.buffered_batches, + refs, self.batch_size, interleave_time, ) @@ -58,9 +73,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 +87,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 +122,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 +143,91 @@ 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()); + } + + /// 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 21d71dc1221..c9b91b344ce 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -1171,10 +1171,15 @@ mod test { 1024 * 1024, 8192, ); + 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,10 +1192,15 @@ mod test { 1024 * 1024, 1, ); + 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 @@ -1291,10 +1301,15 @@ mod test { 1024 * 1024, batch_size as usize, ); + 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 55d88a4ba48..47551d73d84 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,6 +48,10 @@ 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 { @@ -56,20 +64,54 @@ impl, W: Write> BufBatchWriter { Self { shuffle_block_writer, writer, - buffer: vec![], buffer_max_size, compression_context: CompressionContext::default(), coalescer: None, batch_size, + #[cfg(debug_assertions)] + scratch_addr: None, + } + } + + /// 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 @@ -96,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, @@ -119,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 { @@ -140,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(()) } } @@ -160,3 +209,151 @@ 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, 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); + writer.write(&batch, scratch, &time, &time).unwrap(); + writer.flush(scratch, &time, &time).unwrap(); + output + } + + /// 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_scratch_matches_fresh_buffers_and_keeps_capacity() { + let fresh: Vec> = (0..3) + .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 = write_one_partition(p, &mut scratch); + assert!( + scratch.is_empty(), + "recycled scratch must come back drained" + ); + 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); + } + } + + /// 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: `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, 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.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 3a9a6484dba..a7dae2b82ba 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 { @@ -53,6 +58,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, }, } @@ -94,12 +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, - )) + 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) @@ -116,6 +129,7 @@ impl LocalPartitionWriter { shuffle_block_writer, spill_writers, runtime, + recycled_buffer: Vec::new(), } }; Ok(Self { @@ -133,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"), } } } @@ -149,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." @@ -164,18 +178,19 @@ 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 { 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)?; } } @@ -202,20 +217,21 @@ 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 { output_writer, shuffle_block_writer, spill_writers, + recycled_buffer, .. } => { self.offsets[pid] = output_writer.stream_position()?; @@ -234,18 +250,33 @@ 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 + // 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, ); - 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)?; + 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, + ) + })(); + // 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(()) @@ -258,8 +289,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, .. } => { @@ -294,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 d8534b44564..771b4431e17 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -50,16 +50,20 @@ impl SpillWriter { }) } + /// `recycled_buffer` is a scratch byte buffer shared by the sequential per-partition + /// 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, runtime: &RuntimeEnv, metrics: &ShufflePartitionerMetrics, + recycled_buffer: &mut Vec, ) -> datafusion::common::Result<()> { 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, @@ -67,12 +71,26 @@ impl SpillWriter { self.batch_size, ); 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); @@ -80,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(()) @@ -123,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() + ); + } +}