Use single spill file for multiple partitions in native shuffle

Author: Shekharrajak

native/shuffle/src/partitioners/multi_partition.rs

@@ -20,7 +20,7 @@ use crate::partitioners::partitioned_batch_iterator::{
     PartitionedBatchIterator, PartitionedBatchesProducer,
 };
 use crate::partitioners::ShufflePartitioner;
-use crate::writers::{BufBatchWriter, PartitionWriter};
+use crate::writers::{BufBatchWriter, PartitionSpillRange, PartitionWriter, SpillInfo};
 use crate::{comet_partitioning, CometPartitioning, CompressionCodec, ShuffleBlockWriter};
 use arrow::array::{ArrayRef, RecordBatch};
 use arrow::datatypes::SchemaRef;
@@ -125,6 +125,9 @@ pub(crate) struct MultiPartitionShuffleRepartitioner {
     tracing_enabled: bool,
     /// Size of the write buffer in bytes
     write_buffer_size: usize,
+    /// Combined spill files. Each entry is a single file containing data from
+    /// multiple partitions, created during one spill event.
+    spill_infos: Vec<SpillInfo>,
 }
 
 impl MultiPartitionShuffleRepartitioner {
@@ -190,6 +193,7 @@ impl MultiPartitionShuffleRepartitioner {
             reservation,
             tracing_enabled,
             write_buffer_size,
+            spill_infos: vec![],
         })
     }
 
@@ -502,20 +506,53 @@ impl MultiPartitionShuffleRepartitioner {
         with_trace("shuffle_spill", self.tracing_enabled, || {
             let num_output_partitions = self.partition_writers.len();
             let mut partitioned_batches = self.partitioned_batches();
-            let mut spilled_bytes = 0;
+            let mut spilled_bytes: usize = 0;
+
+            // Create a single temporary file for this spill event
+            let temp_file = self
+                .runtime
+                .disk_manager
+                .create_tmp_file("shuffle writer spill")?;
+            let mut spill_file = OpenOptions::new()
+                .write(true)
+                .create(true)
+                .truncate(true)
+                .open(temp_file.path())
+                .map_err(|e| {
+                    DataFusionError::Execution(format!("Error occurred while spilling {e}"))
+                })?;
+
+            let mut partition_ranges = Vec::with_capacity(num_output_partitions);
 
             for partition_id in 0..num_output_partitions {
                 let partition_writer = &mut self.partition_writers[partition_id];
                 let mut iter = partitioned_batches.produce(partition_id);
-                spilled_bytes += partition_writer.spill(
+
+                let offset = spill_file.stream_position()?;
+                let bytes_written = partition_writer.write_to(
                     &mut iter,
-                    &self.runtime,
+                    &mut spill_file,
                     &self.metrics,
                     self.write_buffer_size,
                     self.batch_size,
                 )?;
+
+                if bytes_written > 0 {
+                    partition_ranges.push(Some(PartitionSpillRange {
+                        offset,
+                        length: bytes_written as u64,
+                    }));
+                    spilled_bytes += bytes_written;
+                } else {
+                    partition_ranges.push(None);
+                }
             }
 
+            spill_file.flush()?;
+
+            self.spill_infos
+                .push(SpillInfo::new(temp_file, partition_ranges));
+
             self.reservation.free();
             self.metrics.spill_count.add(1);
             self.metrics.spilled_bytes.add(spilled_bytes);
@@ -524,8 +561,8 @@ impl MultiPartitionShuffleRepartitioner {
     }
 
     #[cfg(test)]
-    pub(crate) fn partition_writers(&self) -> &[PartitionWriter] {
-        &self.partition_writers
+    pub(crate) fn spill_count_files(&self) -> usize {
+        self.spill_infos.len()
     }
 }
 
@@ -579,14 +616,12 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartitioner {
             for i in 0..num_output_partitions {
                 offsets[i] = output_data.stream_position()?;
 
-                // if we wrote a spill file for this partition then copy the
-                // contents into the shuffle file
-                if let Some(spill_path) = self.partition_writers[i].path() {
-                    // Use raw File handle (not BufReader) so that std::io::copy
-                    // can use copy_file_range/sendfile for zero-copy on Linux.
-                    let mut spill_file = File::open(spill_path)?;
+                // Copy spilled data for this partition from each spill file.
+                // Each SpillInfo is a single file containing data from all partitions
+                // ordered by partition ID, with byte ranges tracked per partition.
+                for spill_info in &self.spill_infos {
                     let mut write_timer = self.metrics.write_time.timer();
-                    std::io::copy(&mut spill_file, &mut output_data)?;
+                    spill_info.copy_partition_to(i, &mut output_data)?;
                     write_timer.stop();
                 }
 

native/shuffle/src/shuffle_writer.rs

@@ -367,25 +367,20 @@ mod test {
 
         repartitioner.insert_batch(batch.clone()).await.unwrap();
 
-        {
-            let partition_writers = repartitioner.partition_writers();
-            assert_eq!(partition_writers.len(), 2);
-
-            assert!(!partition_writers[0].has_spill_file());
-            assert!(!partition_writers[1].has_spill_file());
-        }
+        // before spill, no spill files should exist
+        assert_eq!(repartitioner.spill_count_files(), 0);
 
         repartitioner.spill().unwrap();
 
-        // after spill, there should be spill files
-        {
-            let partition_writers = repartitioner.partition_writers();
-            assert!(partition_writers[0].has_spill_file());
-            assert!(partition_writers[1].has_spill_file());
-        }
+        // after spill, exactly one combined spill file should exist (not one per partition)
+        assert_eq!(repartitioner.spill_count_files(), 1);
 
         // insert another batch after spilling
         repartitioner.insert_batch(batch.clone()).await.unwrap();
+
+        // spill again -- should create a second combined spill file
+        repartitioner.spill().unwrap();
+        assert_eq!(repartitioner.spill_count_files(), 2);
     }
 
     fn create_runtime(memory_limit: usize) -> Arc<RuntimeEnv> {

native/shuffle/src/writers/mod.rs

@@ -23,4 +23,4 @@ mod spill;
 pub(crate) use buf_batch_writer::BufBatchWriter;
 pub(crate) use checksum::Checksum;
 pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter};
-pub(crate) use spill::PartitionWriter;
+pub(crate) use spill::{PartitionSpillRange, PartitionWriter, SpillInfo};

native/shuffle/src/writers/spill.rs

@@ -19,24 +19,68 @@ use super::ShuffleBlockWriter;
 use crate::metrics::ShufflePartitionerMetrics;
 use crate::partitioners::PartitionedBatchIterator;
 use crate::writers::buf_batch_writer::BufBatchWriter;
-use datafusion::common::DataFusionError;
 use datafusion::execution::disk_manager::RefCountedTempFile;
-use datafusion::execution::runtime_env::RuntimeEnv;
-use std::fs::{File, OpenOptions};
+use std::fs::File;
+use std::io::{Read, Seek, SeekFrom, Write};
 
-/// A temporary disk file for spilling a partition's intermediate shuffle data.
-struct SpillFile {
-    temp_file: RefCountedTempFile,
-    file: File,
+/// The byte range of a single partition's data within a combined spill file.
+#[derive(Debug, Clone)]
+pub(crate) struct PartitionSpillRange {
+    pub offset: u64,
+    pub length: u64,
 }
 
-/// Manages encoding and optional disk spilling for a single shuffle partition.
+/// Represents a single spill file that contains data from multiple partitions.
+/// Data is written sequentially ordered by partition ID. Each partition's byte
+/// range is tracked in `partition_ranges` so it can be read back during merge.
+pub(crate) struct SpillInfo {
+    /// The temporary file handle -- kept alive to prevent cleanup until we are done.
+    _temp_file: RefCountedTempFile,
+    /// Path to the spill file on disk.
+    path: std::path::PathBuf,
+    /// Byte range for each partition. None means the partition had no data in this spill.
+    pub partition_ranges: Vec<Option<PartitionSpillRange>>,
+}
+
+impl SpillInfo {
+    pub(crate) fn new(
+        temp_file: RefCountedTempFile,
+        partition_ranges: Vec<Option<PartitionSpillRange>>,
+    ) -> Self {
+        let path = temp_file.path().to_path_buf();
+        Self {
+            _temp_file: temp_file,
+            path,
+            partition_ranges,
+        }
+    }
+
+    /// Copy the data for `partition_id` from this spill file into `output`.
+    /// Returns the number of bytes copied.
+    pub(crate) fn copy_partition_to(
+        &self,
+        partition_id: usize,
+        output: &mut impl Write,
+    ) -> datafusion::common::Result<u64> {
+        if let Some(ref range) = self.partition_ranges[partition_id] {
+            if range.length == 0 {
+                return Ok(0);
+            }
+            let mut spill_file = File::open(&self.path)?;
+            spill_file.seek(SeekFrom::Start(range.offset))?;
+            let mut limited = spill_file.take(range.length);
+            let copied = std::io::copy(&mut limited, output)?;
+            Ok(copied)
+        } else {
+            Ok(0)
+        }
+    }
+}
+
+/// Manages encoding for a single shuffle partition. Does not own any spill file --
+/// spill files are managed at the repartitioner level as combined SpillInfo objects.
 pub(crate) struct PartitionWriter {
-    /// Spill file for intermediate shuffle output for this partition. Each spill event
-    /// will append to this file and the contents will be copied to the shuffle file at
-    /// the end of processing.
-    spill_file: Option<SpillFile>,
-    /// Writer that performs encoding and compression
+    /// Writer that performs encoding and compression.
     shuffle_block_writer: ShuffleBlockWriter,
 }
 
@@ -45,51 +89,25 @@ impl PartitionWriter {
         shuffle_block_writer: ShuffleBlockWriter,
     ) -> datafusion::common::Result<Self> {
         Ok(Self {
-            spill_file: None,
             shuffle_block_writer,
         })
     }
 
-    fn ensure_spill_file_created(
-        &mut self,
-        runtime: &RuntimeEnv,
-    ) -> datafusion::common::Result<()> {
-        if self.spill_file.is_none() {
-            // Spill file is not yet created, create it
-            let spill_file = runtime
-                .disk_manager
-                .create_tmp_file("shuffle writer spill")?;
-            let spill_data = OpenOptions::new()
-                .write(true)
-                .create(true)
-                .truncate(true)
-                .open(spill_file.path())
-                .map_err(|e| {
-                    DataFusionError::Execution(format!("Error occurred while spilling {e}"))
-                })?;
-            self.spill_file = Some(SpillFile {
-                temp_file: spill_file,
-                file: spill_data,
-            });
-        }
-        Ok(())
-    }
-
-    pub(crate) fn spill(
+    /// Encode and write a partition's batches to the provided writer.
+    /// Returns the number of bytes written.
+    pub(crate) fn write_to<W: Write>(
         &mut self,
         iter: &mut PartitionedBatchIterator,
-        runtime: &RuntimeEnv,
+        writer: &mut W,
         metrics: &ShufflePartitionerMetrics,
         write_buffer_size: usize,
         batch_size: usize,
     ) -> datafusion::common::Result<usize> {
         if let Some(batch) = iter.next() {
-            self.ensure_spill_file_created(runtime)?;
-
             let total_bytes_written = {
                 let mut buf_batch_writer = BufBatchWriter::new(
                     &mut self.shuffle_block_writer,
-                    &mut self.spill_file.as_mut().unwrap().file,
+                    writer,
                     write_buffer_size,
                     batch_size,
                 );
@@ -106,21 +124,9 @@ impl PartitionWriter {
                 buf_batch_writer.flush(&metrics.encode_time, &metrics.write_time)?;
                 bytes_written
             };
-
             Ok(total_bytes_written)
         } else {
             Ok(0)
         }
     }
-
-    pub(crate) fn path(&self) -> Option<&std::path::Path> {
-        self.spill_file
-            .as_ref()
-            .map(|spill_file| spill_file.temp_file.path())
-    }
-
-    #[cfg(test)]
-    pub(crate) fn has_spill_file(&self) -> bool {
-        self.spill_file.is_some()
-    }
 }