perf: grow the native shuffle writer's reservation in steps

native/core/src/execution/planner.rs

@@ -35,7 +35,7 @@ use crate::execution::{
     planner::expression_registry::ExpressionRegistry,
     planner::operator_registry::OperatorRegistry,
     serde::to_arrow_datatype,
-    shuffle::{SchemaAlignExec, ShuffleWriterExec},
+    shuffle::{SchemaAlignExec, ShuffleWriterExec, ShuffleWriterMemoryConfig},
 };
 use crate::jvm_bridge::{jni_call, JVMClasses};
 use arrow::compute::CastOptions;
@@ -1809,6 +1809,12 @@ impl PhysicalPlanner {
                 // only ever sees a real limit or none at all.
                 let max_buffer_bytes =
                     (writer.max_buffer_bytes > 0).then_some(writer.max_buffer_bytes as usize);
+                let memory_config = ShuffleWriterMemoryConfig {
+                    max_buffer_bytes,
+                    // Zero on the wire means grow the reservation by exactly what each batch
+                    // needs, which is how the writer behaved before it grew in steps.
+                    reservation_step_bytes: writer.reservation_step_bytes as usize,
+                };
                 let shuffle_writer = Arc::new(ShuffleWriterExec::try_new(
                     writer_input,
                     partitioning,
@@ -1817,7 +1823,7 @@ impl PhysicalPlanner {
                     writer.output_index_file.clone(),
                     writer.tracing_enabled,
                     write_buffer_size,
-                    max_buffer_bytes,
+                    memory_config,
                 )?);
 
                 Ok((

native/proto/src/proto/operator.proto

@@ -455,6 +455,10 @@ message ShuffleWriter {
   // Maximum number of bytes that the writer buffers in memory before spilling to disk.
   // Zero means no limit, in which case spilling is driven only by memory pool pressure.
   uint64 max_buffer_bytes = 10;
+  // Granularity for growing the writer's memory reservation. The writer requests memory from the
+  // pool in multiples of this, so a run of small batches costs one request instead of one each.
+  // Zero grows the reservation by exactly what each batch needs.
+  uint64 reservation_step_bytes = 11;
 }
 
 message ParquetWriter {

native/shuffle/benches/shuffle_writer.rs

@@ -31,6 +31,7 @@ use datafusion::{
 };
 use datafusion_comet_shuffle::{
     CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleWriterExec,
+    ShuffleWriterMemoryConfig,
 };
 use itertools::Itertools;
 use std::io::Cursor;
@@ -192,7 +193,7 @@ fn create_shuffle_writer_exec(
         "/tmp/index.out".to_string(),
         false,
         1024 * 1024,
-        None,
+        ShuffleWriterMemoryConfig::default(),
     )
     .unwrap()
 }

native/shuffle/src/bin/shuffle_bench.rs

@@ -45,7 +45,9 @@ use datafusion::physical_plan::common::collect;
 use datafusion::physical_plan::metrics::{MetricValue, MetricsSet};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
-use datafusion_comet_shuffle::{CometPartitioning, CompressionCodec, ShuffleWriterExec};
+use datafusion_comet_shuffle::{
+    CometPartitioning, CompressionCodec, ShuffleWriterExec, ShuffleWriterMemoryConfig,
+};
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use std::fs;
 use std::path::{Path, PathBuf};
@@ -486,7 +488,10 @@ async fn execute_shuffle_write(
         index_file,
         false,
         write_buffer_size,
-        max_buffer_bytes,
+        ShuffleWriterMemoryConfig {
+            max_buffer_bytes,
+            ..Default::default()
+        },
     )
     .expect("Failed to create ShuffleWriterExec");
 

native/shuffle/src/lib.rs

@@ -28,5 +28,7 @@ pub(crate) mod writers;
 pub use comet_partitioning::CometPartitioning;
 pub use ipc::read_ipc_compressed;
 pub use schema_align::SchemaAlignExec;
-pub use shuffle_writer::ShuffleWriterExec;
+pub use shuffle_writer::{
+    ShuffleWriterExec, ShuffleWriterMemoryConfig, DEFAULT_RESERVATION_STEP_BYTES,
+};
 pub use writers::{CompressionCodec, ShuffleBlockWriter};

native/shuffle/src/partitioners/multi_partition.rs

@@ -18,6 +18,7 @@
 use crate::metrics::ShufflePartitionerMetrics;
 use crate::partitioners::partitioned_batch_iterator::PartitionedBatchesProducer;
 use crate::partitioners::ShufflePartitioner;
+use crate::shuffle_writer::ShuffleWriterMemoryConfig;
 use crate::writers::PartitionWriter;
 use crate::{comet_partitioning, CometPartitioning};
 use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch};
@@ -112,10 +113,12 @@ pub(crate) struct MultiPartitionShuffleRepartitioner<T: PartitionWriter> {
     batch_size: usize,
     /// Reservation for repartitioning
     reservation: MemoryReservation,
-    /// Spill once the reservation reaches this many bytes, independently of whether the memory
-    /// pool still has capacity. `None` disables the limit, leaving pool pressure as the only
-    /// spill trigger.
-    max_buffer_bytes: Option<usize>,
+    /// Exact bytes this writer currently holds resident, which is what the `max_buffer_bytes`
+    /// limit and the `used()` figure are measured against. `reservation` is grown in coarse steps,
+    /// so `reservation.size()` is this value plus up to one `reservation_step()`.
+    buffered_bytes: usize,
+    /// Spill limit and reservation granularity for this writer
+    memory_config: ShuffleWriterMemoryConfig,
     tracing_enabled: bool,
     /// Start addresses (as `usize`, since raw pointers are not `Send`) of the backing buffers
     /// currently pinned by `buffered_batches`, so the spill reservation charges each distinct
@@ -178,7 +181,7 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
         runtime: Arc<RuntimeEnv>,
         batch_size: usize,
         tracing_enabled: bool,
-        max_buffer_bytes: Option<usize>,
+        memory_config: ShuffleWriterMemoryConfig,
     ) -> datafusion::common::Result<Self> {
         let num_output_partitions = partitioning.partition_count();
         assert_ne!(
@@ -216,7 +219,8 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
             scratch,
             batch_size,
             reservation,
-            max_buffer_bytes,
+            buffered_bytes: 0,
+            memory_config,
             tracing_enabled,
             pinned_buffers: HashSet::new(),
         })
@@ -457,22 +461,61 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
             mem_growth += after_size.saturating_sub(before_size);
         }
 
-        // `try_grow` is evaluated first so the reservation accounts for this batch either way.
-        // Checking after buffering lets the writer overshoot the limit by at most one batch,
-        // which is how the memory-pressure trigger already behaves.
-        if self.reservation.try_grow(mem_growth).is_err()
+        self.buffered_bytes += mem_growth;
+
+        // Growing after buffering lets the writer overshoot its limit by at most one batch, which
+        // is how the memory-pressure trigger already behaves. The limit is checked against
+        // `buffered_bytes` rather than `reservation.size()` because the latter is rounded up to a
+        // step boundary and would otherwise trip the limit early.
+        if self.grow_reservation()
             || self
+                .memory_config
                 .max_buffer_bytes
-                .is_some_and(|limit| self.reservation.size() >= limit)
+                .is_some_and(|limit| self.buffered_bytes >= limit)
         {
             self.spill()?;
         }
 
         Ok(())
     }
 
+    /// Bytes to grow [`Self::reservation`] by at a time, from
+    /// `spark.comet.shuffle.native.reservationStepBytes`. Capped well below `max_buffer_bytes` so
+    /// that a writer with a small limit does not reserve a multiple of its own budget. Zero — from
+    /// either the config or the cap — means grow by exactly what the batch needs.
+    fn reservation_step(&self) -> usize {
+        let step = self.memory_config.reservation_step_bytes;
+        match self.memory_config.max_buffer_bytes {
+            Some(limit) => step.min(limit / 8),
+            None => step,
+        }
+    }
+
+    /// Top the reservation back up if `buffered_bytes` has outgrown it, rounding the request up to
+    /// [`Self::reservation_step`] so that a run of small batches costs one acquisition rather than
+    /// one per batch.
+    ///
+    /// Returns whether the pool denied the memory, in which case the caller must spill.
+    fn grow_reservation(&mut self) -> bool {
+        let deficit = self.buffered_bytes.saturating_sub(self.reservation.size());
+        if deficit == 0 {
+            return false;
+        }
+        if self
+            .reservation
+            .try_grow(deficit.max(self.reservation_step()))
+            .is_ok()
+        {
+            return false;
+        }
+        // A rounded-up request can be denied where the exact deficit would still fit, so retry at
+        // the exact size before spilling. This keeps the spill trigger where it was when every
+        // batch reserved exactly what it needed.
+        self.reservation.try_grow(deficit).is_err()
+    }
+
     fn used(&self) -> usize {
-        self.reservation.size()
+        self.buffered_bytes
     }
 
     fn spilled_bytes(&self) -> usize {
@@ -526,6 +569,7 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
             }
 
             self.reservation.free();
+            self.buffered_bytes = 0;
             self.pinned_buffers.clear();
             self.metrics.spill_count.add(1);
             Ok(())