perf: reuse per-partition scratch in the shuffle write path

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(

native/shuffle/src/partitioners/multi_partition.rs

@@ -520,12 +520,17 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
         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,
                     )
                 })
@@ -579,15 +584,17 @@ impl<T: PartitionWriter> 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,
                 )?;
             }

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,39 +73,40 @@ 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,
 }
 
 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::<Vec<_>>();
-        let current_indices = indices
-            .iter()
-            .map(|(i_batch, i_row)| (*i_batch as usize, *i_row as usize))
-            .collect::<Vec<_>>();
         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<RecordBatch> {
+        let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]));
+        (0..3)
+            .map(|b| {
+                let values: Vec<i32> = (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<RecordBatch> = producer
+            .produce(&refs, 0, &time)
+            .collect::<datafusion::common::Result<_>>()
+            .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<RecordBatch> = 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<RecordBatch> = producer
+            .produce(&refs, 1, &time)
+            .collect::<datafusion::common::Result<_>>()
+            .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);
+    }
+}

native/shuffle/src/shuffle_writer.rs

@@ -1222,10 +1222,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)
@@ -1238,10 +1243,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
@@ -1342,10 +1352,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);