diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 8d668de7725..ef19e55625f 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -1200,7 +1200,7 @@ mod test { #[test] #[cfg_attr(miri, ignore)] fn test_batch_coalescing_reduces_size() { - use crate::writers::BufBatchWriter; + use crate::writers::{BufBatchWriter, ShuffleScratch}; use arrow::array::Int32Array; // Create a wide schema to amplify per-block schema overhead @@ -1238,7 +1238,7 @@ mod test { 1024 * 1024, 8192, ); - let mut scratch = Vec::new(); + let mut scratch = ShuffleScratch::default(); for batch in &small_batches { buf_writer .write(batch, &mut scratch, &encode_time, &write_time) @@ -1259,7 +1259,7 @@ mod test { 1024 * 1024, 1, ); - let mut scratch = Vec::new(); + let mut scratch = ShuffleScratch::default(); for batch in &small_batches { buf_writer .write(batch, &mut scratch, &encode_time, &write_time) @@ -1326,7 +1326,7 @@ mod test { #[test] #[cfg_attr(miri, ignore)] fn test_full_batches_bypass_coalescer() { - use crate::writers::BufBatchWriter; + use crate::writers::{BufBatchWriter, ShuffleScratch}; use arrow::array::Int32Array; let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); @@ -1368,7 +1368,7 @@ mod test { 1024 * 1024, batch_size as usize, ); - let mut scratch = Vec::new(); + let mut scratch = ShuffleScratch::default(); for batch in &inputs { buf_writer .write(batch, &mut scratch, &encode_time, &write_time) diff --git a/native/shuffle/src/writers/buf_batch_writer.rs b/native/shuffle/src/writers/buf_batch_writer.rs index 1719cc105f9..b1e1f59215b 100644 --- a/native/shuffle/src/writers/buf_batch_writer.rs +++ b/native/shuffle/src/writers/buf_batch_writer.rs @@ -23,54 +23,113 @@ use datafusion::physical_plan::metrics::Time; use std::borrow::Borrow; use std::io::{Cursor, Seek, SeekFrom, Write}; +/// Task-scoped serialization state threaded through every [`BufBatchWriter`] of a task. +/// +/// A shuffle task creates one short-lived `BufBatchWriter` per output partition per spill or +/// finish cycle, so anything the writer owned would be rebuilt `partitions x cycles` times. +/// Both members here are cheap to keep and were previously rebuilt per writer: +/// +/// * `buffer` is the byte buffer blocks are serialized into before being handed to the +/// underlying writer in `buffer_max_size` chunks. It is capped back to `buffer_max_size` +/// after every drain, so it never retains more than one block past that size. +/// * `ipc_context` holds arrow-ipc's flatbuffer builder for record-batch metadata. With +/// arrow's default `reserve_scratch = false`, which is what this uses, the context does not +/// retain the block body between encodes (each block body is a fresh `Vec` that is dropped +/// after the write), so sharing it costs one small metadata builder per task, not a block. +/// +/// The struct is borrowed per call rather than owned, so an error mid-partition cannot strand it +/// inside a dropped writer and silently end recycling. +#[derive(Default)] +pub(crate) struct ShuffleScratch { + pub(crate) buffer: Vec, + pub(crate) ipc_context: IpcWriteContext, +} + +impl ShuffleScratch { + /// Drops any bytes left in the buffer. Used after an error so a failed partition's bytes + /// cannot leak into the next partition's block; the IPC context holds no block data. + pub(crate) fn clear(&mut self) { + self.buffer.clear(); + } +} + /// Write batches to writer while using a buffer to avoid frequent system calls. /// 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. +/// The scratch is borrowed per call rather than owned: task-scoped [`ShuffleScratch`] is +/// threaded through every `write`/`flush`, so one buffer and one IPC context serve all the +/// short-lived writers of a task. +/// +/// A writer either coalesces or passes batches through, chosen at construction: /// -/// 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 -/// remainder emitted on flush. The coalescer is lazily initialized on the first write and -/// configured (via `biggest_coalesce_batch_size`) to pass batches that are already at least -/// `batch_size` rows straight through, verbatim and without copying them, so an oversized input -/// batch is written as a single oversized block. +/// * [`Self::new`]: 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 remainder emitted on flush. The coalescer is lazily +/// initialized on the first write and configured (via `biggest_coalesce_batch_size`) to pass +/// batches that are already at least `batch_size` rows straight through, verbatim and without +/// copying them, so an oversized input batch is written as a single oversized block. This is +/// the mode for the long-lived single-partition writer, whose inputs can genuinely be small. +/// * [`Self::new_passthrough`]: every batch is serialized as its own block, verbatim. This is +/// the mode for the per-partition writers of a multi-partition shuffle, whose input is a +/// `PartitionedBatchIterator` that already emits maximal `batch_size` chunks plus one tail: +/// there is never a second batch for the tail to coalesce with, so coalescing there would +/// only copy the tail's rows into builders and re-emit the same block. pub(crate) struct BufBatchWriter, W: Write> { shuffle_block_writer: S, writer: W, buffer_max_size: usize, - compression_context: IpcWriteContext, /// Coalesces small batches into target_batch_size before serialization. - /// Lazily initialized on first write to capture the schema. + /// Lazily initialized on first write to capture the schema. Never set in passthrough mode. 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. + /// Target batch size for coalescing; `None` selects passthrough mode. + coalesce_batch_size: Option, + /// Address of the scratch 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, - /// Running total of bytes serialized through this writer, used to report spilled bytes when - /// the underlying writer does not implement [`Seek`] (e.g. a `Box`). + /// Running total of bytes serialized through this writer, used to report spilled bytes and + /// to track output offsets without a `Seek` on the underlying writer. total_bytes_written: u64, } impl, W: Write> BufBatchWriter { + /// A coalescing writer; see the type-level docs for when to use which mode. pub(crate) fn new( shuffle_block_writer: S, writer: W, buffer_max_size: usize, batch_size: usize, + ) -> Self { + Self::with_mode( + shuffle_block_writer, + writer, + buffer_max_size, + Some(batch_size), + ) + } + + /// A passthrough writer that serializes every batch as its own block, verbatim. + pub(crate) fn new_passthrough( + shuffle_block_writer: S, + writer: W, + buffer_max_size: usize, + ) -> Self { + Self::with_mode(shuffle_block_writer, writer, buffer_max_size, None) + } + + fn with_mode( + shuffle_block_writer: S, + writer: W, + buffer_max_size: usize, + coalesce_batch_size: Option, ) -> Self { Self { shuffle_block_writer, writer, buffer_max_size, - compression_context: IpcWriteContext::default(), coalescer: None, - batch_size, + coalesce_batch_size, #[cfg(debug_assertions)] scratch_addr: None, total_bytes_written: 0, @@ -80,20 +139,19 @@ impl, W: Write> BufBatchWriter { /// 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 + /// Identity is the scratch'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) { + fn check_scratch(&mut self, scratch: &ShuffleScratch) { #[cfg(debug_assertions)] { - let addr = scratch as *const Vec as usize; + let addr = scratch as *const ShuffleScratch as usize; match self.scratch_addr { None => { debug_assert!( - scratch.is_empty(), + scratch.buffer.is_empty(), "fresh BufBatchWriter handed a non-empty scratch buffer ({} bytes)", - scratch.len() + scratch.buffer.len() ); self.scratch_addr = Some(addr); } @@ -105,18 +163,20 @@ impl, W: Write> BufBatchWriter { } } - /// `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. + /// `scratch` is the caller-owned serialization state to encode into; threading the same + /// one 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, + scratch: &mut ShuffleScratch, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { self.check_scratch(scratch); - let batch_size = self.batch_size; + let Some(batch_size) = self.coalesce_batch_size else { + return self.write_batch_to_buffer(batch, scratch, encode_time, write_time); + }; let coalescer = self.coalescer.get_or_insert_with(|| { // Enable BatchCoalescer's zero-copy passthrough for batches that are already big // enough, so we don't `copy_rows` the whole batch into the in-progress builders just @@ -151,33 +211,43 @@ impl, W: Write> BufBatchWriter { fn write_batch_to_buffer( &mut self, batch: &RecordBatch, - scratch: &mut Vec, + scratch: &mut ShuffleScratch, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { - let mut cursor = Cursor::new(&mut *scratch); + let ShuffleScratch { + buffer, + ipc_context, + } = scratch; + let mut cursor = Cursor::new(&mut *buffer); cursor.seek(SeekFrom::End(0))?; let bytes_written = self.shuffle_block_writer.borrow().write_batch( batch, &mut cursor, - &mut self.compression_context, + ipc_context, encode_time, )?; let pos = cursor.position(); if pos >= self.buffer_max_size as u64 { let mut write_timer = write_time.timer(); - self.writer.write_all(scratch)?; + self.writer.write_all(buffer)?; write_timer.stop(); - scratch.clear(); + buffer.clear(); } self.total_bytes_written += bytes_written as u64; Ok(bytes_written) } - /// Flushes buffered rows and bytes; `scratch` is left drained for the next writer. - pub(crate) fn flush( + /// Writes any rows still held by the coalescer and any bytes still in `scratch` to the + /// underlying writer, without flushing that writer. `scratch` is left drained for the next + /// writer. + /// + /// This is what a per-partition writer over a shared buffered output wants: the output's + /// own buffer keeps accumulating across partitions and is flushed once at the end, instead + /// of every partition ending in its own write syscall. + pub(crate) fn drain( &mut self, - scratch: &mut Vec, + scratch: &mut ShuffleScratch, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result<()> { @@ -194,23 +264,38 @@ impl, W: Write> BufBatchWriter { self.write_batch_to_buffer(batch, scratch, encode_time, write_time)?; } - // Flush the scratch buffer to the underlying writer + // Hand the scratch buffer to the underlying writer let mut write_timer = write_time.timer(); - if !scratch.is_empty() { - self.writer.write_all(scratch)?; + if !scratch.buffer.is_empty() { + self.writer.write_all(&scratch.buffer)?; } - self.writer.flush()?; write_timer.stop(); - scratch.clear(); + scratch.buffer.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); + scratch.buffer.shrink_to(self.buffer_max_size); + Ok(()) + } + + /// [`Self::drain`], then flushes the underlying writer. + pub(crate) fn flush( + &mut self, + scratch: &mut ShuffleScratch, + encode_time: &Time, + write_time: &Time, + ) -> datafusion::common::Result<()> { + self.drain(scratch, encode_time, write_time)?; + let mut write_timer = write_time.timer(); + self.writer.flush()?; + write_timer.stop(); Ok(()) } /// Total number of bytes serialized through this writer since it was created. Unlike /// [`Self::writer_stream_position`], this does not require the underlying writer to implement - /// [`Seek`], so it is used to report spilled bytes when writing to a `Box`. + /// [`Seek`], so it is used to report spilled bytes when writing to a `Box` + /// and to track output offsets without flushing a buffered output. After [`Self::drain`] or + /// [`Self::flush`] every counted byte has been handed to the underlying writer. pub(crate) fn bytes_written(&self) -> u64 { self.total_bytes_written } @@ -236,7 +321,7 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(values))]).unwrap() } - fn write_one_partition(seed: i64, scratch: &mut Vec) -> Vec { + fn write_one_partition(seed: i64, scratch: &mut ShuffleScratch) -> Vec { let batch = test_batch(seed); let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::Zstd(1)) @@ -255,15 +340,15 @@ mod tests { #[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())) + .map(|p| write_one_partition(p, &mut ShuffleScratch::default())) .collect(); - let mut scratch = Vec::new(); + let mut scratch = ShuffleScratch::default(); let mut recycled = Vec::new(); for p in 0..3 { let output = write_one_partition(p, &mut scratch); assert!( - scratch.is_empty(), + scratch.buffer.is_empty(), "recycled scratch must come back drained" ); recycled.push(output); @@ -271,7 +356,7 @@ mod tests { assert_eq!(fresh, recycled); assert!( - scratch.capacity() > 0, + scratch.buffer.capacity() > 0, "capacity grown in one partition must survive into the next" ); for output in &recycled { @@ -280,6 +365,62 @@ mod tests { } } + /// A passthrough writer must emit every input batch as its own block, in order, with no + /// coalescing of a small tail into a following batch and no rows held back until `drain`. + #[test] + fn passthrough_writes_each_batch_as_its_own_block() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let make_batch = |start: i64, rows: i64| { + let values: Vec = (start..start + rows).collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + ) + .unwrap() + }; + // The shape a PartitionedBatchIterator produces: maximal chunks then one tail, but + // also a tail followed by more batches, which a coalescer would merge. + let inputs = [ + make_batch(0, 100), + make_batch(100, 30), + make_batch(130, 100), + make_batch(230, 7), + ]; + let block_writer = + ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::None).unwrap(); + let mut output = Vec::new(); + let time = Time::default(); + let mut scratch = ShuffleScratch::default(); + let mut writer = BufBatchWriter::new_passthrough(block_writer, &mut output, 1 << 20); + for batch in &inputs { + writer.write(batch, &mut scratch, &time, &time).unwrap(); + } + writer.drain(&mut scratch, &time, &time).unwrap(); + assert!( + scratch.buffer.is_empty(), + "drain must leave the scratch empty" + ); + assert_eq!(writer.bytes_written() as usize, output.len()); + + let mut block_rows = Vec::new(); + let mut next = 0i64; + let mut pos = 0; + while pos < output.len() { + let len = u64::from_le_bytes(output[pos..pos + 8].try_into().unwrap()) as usize; + let block = read_ipc_compressed(&output[pos + 16..pos + 8 + len]).unwrap(); + let values = block + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), next, "rows must stay in input order"); + next += block.num_rows() as i64; + block_rows.push(block.num_rows()); + pos += 8 + len; + } + assert_eq!(block_rows, vec![100, 30, 100, 7]); + } + /// 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)] @@ -292,7 +433,10 @@ mod tests { 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 mut dirty = ShuffleScratch { + buffer: vec![0xAB, 0xCD], + ..Default::default() + }; let _ = writer.write(&batch, &mut dirty, &time, &time); } @@ -308,9 +452,9 @@ mod tests { 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(); + let mut first = ShuffleScratch::default(); writer.write(&batch, &mut first, &time, &time).unwrap(); - let mut second = Vec::new(); + let mut second = ShuffleScratch::default(); let _ = writer.write(&batch, &mut second, &time, &time); } @@ -328,20 +472,20 @@ mod tests { 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 scratch = ShuffleScratch::default(); 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, + scratch.buffer.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.buffer.is_empty()); assert!( - scratch.capacity() <= buffer_max_size, + scratch.buffer.capacity() <= buffer_max_size, "retained capacity {} exceeds cap {}", - scratch.capacity(), + scratch.buffer.capacity(), buffer_max_size ); @@ -352,18 +496,18 @@ mod tests { 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 scratch = ShuffleScratch::default(); 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(); + let cap_after_write = scratch.buffer.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!(scratch.buffer.is_empty()); assert_eq!( - scratch.capacity(), + scratch.buffer.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 e22e339f949..9d17daf9344 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -18,7 +18,7 @@ use crate::metrics::ShufflePartitionerMetrics; use crate::writers::local::spill::SpillWriter; use crate::writers::partition_writer::PartitionWriter; -use crate::writers::BufBatchWriter; +use crate::writers::{BufBatchWriter, ShuffleScratch}; use crate::ShuffleBlockWriter; use arrow::array::RecordBatch; use datafusion::common::DataFusionError; @@ -38,16 +38,18 @@ use std::sync::Arc; /// [`PartitionWriter::finish_all`]. /// * Multi-partition shuffles finalize one partition at a time in /// [`PartitionWriter::finish_partition`], each with its own short-lived -/// `BufBatchWriter`, so coalescing intentionally does not cross partition -/// boundaries. They hold the raw output writer and block writer directly. +/// passthrough `BufBatchWriter`: the batches for a partition arrive as maximal +/// `batch_size` chunks plus one tail, so there is nothing to coalesce and each +/// chunk is serialized verbatim as one block. They hold the raw output writer +/// and block writer directly. #[allow(clippy::large_enum_variant)] enum DataOutput { /// Single-partition output: one long-lived writer streams all batches. 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, + /// Task-scoped serialization scratch threaded through every call on the + /// long-lived writer, which borrows rather than owns it. + scratch: ShuffleScratch, }, /// Multi-partition output: batches are staged per partition and merged into /// `output_writer` one partition at a time during `finish_partition`. @@ -59,11 +61,17 @@ 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, + /// Serialization scratch recycled through the short-lived per-partition + /// `BufBatchWriter`s. Partitions are written strictly one at a time, so a single + /// byte buffer keeps its grown capacity and a single IPC context keeps its + /// encoder state across the whole task, instead of every partition regrowing + /// fresh ones. + scratch: ShuffleScratch, + /// Bytes handed to `output_writer` so far. Partition offsets are derived from + /// this running total rather than from `stream_position()`, which would flush + /// the output buffer once per partition; the buffer is flushed once, in + /// `finish_all`, and the total is checked against the file position there. + bytes_written: u64, }, } @@ -79,7 +87,6 @@ pub(crate) struct LocalPartitionWriter { /// with the total length so partition sizes are simple offset differences. /// Has `num_output_partitions + 1` elements. offsets: Vec, - batch_size: usize, write_buffer_size: usize, num_output_partitions: usize, /// Id of the last partition passed to `finish_partition`, used to assert @@ -112,32 +119,26 @@ impl LocalPartitionWriter { write_buffer_size, batch_size, ), - scratch: Vec::new(), + scratch: ShuffleScratch::default(), } } else { let output_writer = BufWriter::with_capacity(write_buffer_size, output_file); let spill_writers = (0..num_output_partitions) - .map(|_| { - SpillWriter::try_new( - shuffle_block_writer.clone(), - write_buffer_size, - batch_size, - ) - }) + .map(|_| SpillWriter::try_new(shuffle_block_writer.clone(), write_buffer_size)) .collect::>>()?; DataOutput::Multi { output_writer, shuffle_block_writer, spill_writers, runtime, - recycled_buffer: Vec::new(), + scratch: ShuffleScratch::default(), + bytes_written: 0, } }; Ok(Self { output_index_file, data_output, offsets: vec![0u64; num_output_partitions + 1], - batch_size, write_buffer_size, num_output_partitions, last_finish_pid: -1, @@ -185,13 +186,13 @@ impl PartitionWriter for LocalPartitionWriter { DataOutput::Multi { spill_writers, runtime, - recycled_buffer, + scratch, .. } => { // 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, recycled_buffer)?; + spill_writers[pid].write(iter, runtime, metrics, scratch)?; } } @@ -215,7 +216,6 @@ impl PartitionWriter for LocalPartitionWriter { self.last_finish_pid = pid as i32; let write_buffer_size = self.write_buffer_size; - let batch_size = self.batch_size; match &mut self.data_output { DataOutput::Single { writer, scratch } => { @@ -232,10 +232,14 @@ impl PartitionWriter for LocalPartitionWriter { output_writer, shuffle_block_writer, spill_writers, - recycled_buffer, + scratch, + bytes_written, .. } => { - self.offsets[pid] = output_writer.stream_position()?; + // The offset is the running byte total, not `stream_position()`: asking a + // `BufWriter` for its position flushes it, which would turn every partition + // into its own write syscall and defeat the output buffer. + self.offsets[pid] = *bytes_written; // if we wrote a spill file for this partition then copy the // contents into the shuffle file @@ -245,39 +249,39 @@ impl PartitionWriter for LocalPartitionWriter { // can use copy_file_range/sendfile for zero-copy on Linux. let mut spill_file = File::open(spill_path)?; let mut write_timer = metrics.write_time.timer(); - std::io::copy(&mut spill_file, output_writer)?; + *bytes_written += std::io::copy(&mut spill_file, output_writer)?; write_timer.stop(); } } // Write in memory batches to output data file. Each partition uses its - // 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( + // own short-lived writer, but the scratch (byte buffer and IPC context) is + // shared so its capacity carries over to the next one. The batches arrive + // as maximal chunks plus one tail, so they pass through as their own blocks + // rather than being copied through a coalescer. `drain` rather than `flush` + // hands the bytes to `output_writer` without flushing it; the output is + // flushed once in `finish_all`. + let mut buf_batch_writer = BufBatchWriter::new_passthrough( shuffle_block_writer, - output_writer, + &mut *output_writer, write_buffer_size, - batch_size, ); let result: datafusion::common::Result<()> = (|| { for batch in iter.by_ref() { let batch = batch?; buf_batch_writer.write( &batch, - recycled_buffer, + scratch, &metrics.encode_time, &metrics.write_time, )?; } - buf_batch_writer.flush( - recycled_buffer, - &metrics.encode_time, - &metrics.write_time, - ) + buf_batch_writer.drain(scratch, &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())?; + result.inspect_err(|_| scratch.clear())?; + *bytes_written += buf_batch_writer.bytes_written(); } } Ok(()) @@ -294,11 +298,23 @@ impl PartitionWriter for LocalPartitionWriter { writer.flush(scratch, &metrics.encode_time, &metrics.write_time)?; writer.writer_stream_position()? } - DataOutput::Multi { output_writer, .. } => { + DataOutput::Multi { + output_writer, + bytes_written, + .. + } => { let mut write_timer = metrics.write_time.timer(); output_writer.flush()?; let pos = output_writer.stream_position()?; write_timer.stop(); + // The offsets were derived from the running total; the file must agree, + // or the index would point readers at the wrong bytes. + if pos != *bytes_written { + return Err(DataFusionError::Execution(format!( + "shuffle write error: data file holds {pos} bytes but the partition \ + offsets account for {bytes_written}" + ))); + } pos } }; @@ -383,17 +399,77 @@ mod tests { assert!(writer.finish_partition(0, &mut iter, &metrics).is_err()); match &writer.data_output { - DataOutput::Multi { - recycled_buffer, .. - } => assert!( - recycled_buffer.is_empty(), + DataOutput::Multi { scratch, .. } => assert!( + scratch.buffer.is_empty(), "errored partition left {} bytes in the recycled buffer", - recycled_buffer.len() + scratch.buffer.len() ), DataOutput::Single { .. } => unreachable!("two partitions use the multi output"), } } + /// Partition offsets are tracked arithmetically instead of read back from the output + /// file, so the index must still describe the data file exactly: every offset lands on a + /// block boundary and each partition's blocks decode to the rows written for it, with a + /// spilled prefix and in-memory tail both accounted for. + #[test] + fn offsets_match_data_file_with_spilled_and_in_memory_batches() { + let batch = test_batch(); + let dir = tempfile::tempdir().unwrap(); + let mut writer = partition_writer(&batch, &dir, Arc::new(RuntimeEnv::default())); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + + // Partition 0 spills one batch, then finishes with two more in memory. + // Partition 1 has no spill and finishes with one batch. + writer + .write(0, &mut vec![Ok(batch.clone())].into_iter(), &metrics) + .unwrap(); + writer + .finish_partition( + 0, + &mut vec![Ok(batch.clone()), Ok(batch.clone())].into_iter(), + &metrics, + ) + .unwrap(); + writer + .finish_partition(1, &mut vec![Ok(batch.clone())].into_iter(), &metrics) + .unwrap(); + writer.finish_all(&metrics).unwrap(); + + let data = std::fs::read(dir.path().join("data.out")).unwrap(); + let index = std::fs::read(dir.path().join("index.out")).unwrap(); + let offsets: Vec = index + .as_chunks::<8>() + .0 + .iter() + .map(|c| i64::from_le_bytes(*c) as usize) + .collect(); + assert_eq!(offsets.len(), 3); + assert_eq!(offsets[0], 0); + assert_eq!(offsets[2], data.len()); + + // Decode every block within a partition's byte range; a wrong offset would land + // mid-block and fail to parse. + let rows_in = |range: std::ops::Range| { + let mut pos = range.start; + let mut rows = 0; + while pos < range.end { + let len = u64::from_le_bytes(data[pos..pos + 8].try_into().unwrap()) as usize; + assert!( + pos + 8 + len <= range.end, + "block crosses a partition boundary" + ); + rows += crate::read_ipc_compressed(&data[pos + 16..pos + 8 + len]) + .unwrap() + .num_rows(); + pos += 8 + len; + } + rows + }; + assert_eq!(rows_in(offsets[0]..offsets[1]), 300); + assert_eq!(rows_in(offsets[1]..offsets[2]), 100); + } + /// Spilled bytes the writer cannot reach must fail the task. Skipping the copy the way /// an unspilled partition is skipped would leave the offsets claiming bytes that were /// never written, so the reader would decode the next partition's block as this one's. diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index 77fe009046d..144c413dd11 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -16,7 +16,7 @@ // under the License. use crate::metrics::ShufflePartitionerMetrics; -use crate::writers::BufBatchWriter; +use crate::writers::{BufBatchWriter, ShuffleScratch}; use crate::ShuffleBlockWriter; use arrow::record_batch::RecordBatch; use datafusion::common::DataFusionError; @@ -33,7 +33,6 @@ struct ActiveSpillFile { pub(crate) struct SpillWriter { shuffle_block_writer: ShuffleBlockWriter, write_buffer_size: usize, - batch_size: usize, spill_file: Option, } @@ -41,35 +40,35 @@ impl SpillWriter { pub(crate) fn try_new( shuffle_block_writer: ShuffleBlockWriter, write_buffer_size: usize, - batch_size: usize, ) -> datafusion::common::Result { Ok(Self { shuffle_block_writer, write_buffer_size, - batch_size, spill_file: None, }) } - /// `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. + /// `recycled_buffer` is the serialization scratch shared by the sequential per-partition + /// spill writes; it is left drained on return so one buffer's capacity and one IPC context + /// serve every partition instead of each write regrowing its own. + /// + /// `iter` comes from a `PartitionedBatchIterator`, which already emits maximal `batch_size` + /// chunks plus one tail, so the batches are written through verbatim rather than coalesced. pub(crate) fn write>>( &mut self, iter: &mut I, runtime: &RuntimeEnv, metrics: &ShufflePartitionerMetrics, - recycled_buffer: &mut Vec, + recycled_buffer: &mut ShuffleScratch, ) -> datafusion::common::Result<()> { if let Some(batch) = iter.next() { self.ensure_spill_file_created(runtime)?; let result = (|| { - let mut buf_batch_writer = BufBatchWriter::new( + let mut buf_batch_writer = BufBatchWriter::new_passthrough( &mut self.shuffle_block_writer, &mut self.spill_file.as_mut().unwrap().writer, self.write_buffer_size, - self.batch_size, ); buf_batch_writer.write( &batch?, @@ -251,11 +250,11 @@ mod tests { .unwrap() } - fn spill_writer(batch: &RecordBatch, batch_size: usize) -> SpillWriter { + fn spill_writer(batch: &RecordBatch) -> SpillWriter { let block_writer = ShuffleBlockWriter::try_new(batch.schema_ref().as_ref(), CompressionCodec::None) .unwrap(); - SpillWriter::try_new(block_writer, 1 << 20, batch_size).unwrap() + SpillWriter::try_new(block_writer, 1 << 20).unwrap() } /// A spill whose batch iterator fails after a batch was already encoded must hand @@ -263,11 +262,10 @@ mod tests { #[test] fn write_error_drains_recycled_buffer() { let batch = test_batch(); - // batch_size below the row count so the first write serializes into the scratch. - let mut spill = spill_writer(&batch, 10); + let mut spill = spill_writer(&batch); let runtime = RuntimeEnv::default(); let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); - let mut recycled = Vec::new(); + let mut recycled = ShuffleScratch::default(); let mut iter = vec![ Ok(batch), Err(DataFusionError::Execution("injected failure".to_string())), @@ -278,9 +276,9 @@ mod tests { .write(&mut iter, &runtime, &metrics, &mut recycled) .is_err()); assert!( - recycled.is_empty(), + recycled.buffer.is_empty(), "errored spill left {} bytes in the recycled buffer", - recycled.len() + recycled.buffer.len() ); } @@ -288,7 +286,7 @@ mod tests { #[test] fn path_is_none_when_nothing_spilled() { let batch = test_batch(); - let spill = spill_writer(&batch, 10); + let spill = spill_writer(&batch); assert!(!spill.has_spill_file()); assert_eq!(spill.path().unwrap(), None); } @@ -298,10 +296,10 @@ mod tests { #[test] fn path_errors_when_backend_has_no_local_path() { let batch = test_batch(); - let mut spill = spill_writer(&batch, 10); + let mut spill = spill_writer(&batch); let runtime = pathless_backend::runtime(); let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); - let mut recycled = Vec::new(); + let mut recycled = ShuffleScratch::default(); let mut iter = vec![Ok(batch)].into_iter(); spill diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index fb3af2c991a..82f0517752d 100644 --- a/native/shuffle/src/writers/mod.rs +++ b/native/shuffle/src/writers/mod.rs @@ -22,7 +22,7 @@ mod partition_writer; mod rss; mod shuffle_block_writer; -pub(crate) use buf_batch_writer::BufBatchWriter; +pub(crate) use buf_batch_writer::{BufBatchWriter, ShuffleScratch}; pub(crate) use checksum::Checksum; pub(crate) use local::local_partition_writer::LocalPartitionWriter; pub(crate) use partition_writer::PartitionWriter;