From 77a4308237dc1ea6fc5a4ee6e1f8e535abad493e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 17 Aug 2026 13:34:17 -0600 Subject: [PATCH 1/3] perf: grow the native shuffle writer's reservation in steps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `buffer_partitioned_batch_may_spill` reserved exactly what each batch newly pinned, which for an 8192-row `i32` column is 32 KB. Against Comet's unified memory pool every one of those is a JNI round-trip into Spark's memory manager, so a 201-partition shuffle of 10M rows spent 256 acquisitions per task, 96% of them for 32 KB, to reserve 24 MB. Track the exact resident figure in `buffered_bytes` and grow the reservation in 1 MB steps on top of it. On the same benchmark that is 13 acquisitions per task instead of 256, and 50 µs per task inside the JNI call instead of 327 µs. The cost is up to one step of reservation held beyond what the writer is using, per active writer. Two details keep the spill behaviour where it was: - the `max_buffer_bytes` limit is now compared against `buffered_bytes`, since `reservation.size()` is rounded up and would trip the limit early - a rounded-up request that the pool denies is retried at the exact deficit before spilling, so a batch that would have fit before still fits Part of #5383. --- .../src/partitioners/multi_partition.rs | 63 +++++++++- native/shuffle/src/shuffle_writer.rs | 118 ++++++++++++++++++ 2 files changed, 175 insertions(+), 6 deletions(-) diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 5644b1085d7..b7a0ef71b7e 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -98,6 +98,16 @@ impl ScratchSpace { } } +/// Granularity for growing [`MultiPartitionShuffleRepartitioner::reservation`]. +/// +/// The writer's per-batch growth is small — one Arrow buffer per column, 32 KB for an 8192-row +/// `i32` column — and against Comet's unified memory pool every `try_grow` is a JNI round-trip +/// into Spark's memory manager. On a 201-partition shuffle of 10M rows that was 256 acquisitions +/// per task, 96% of them for 32 KB, at ~1.3 µs each (see +/// ). Growing in coarse steps took that +/// to 13 acquisitions per task, trading up to one step of over-reservation per active writer. +const RESERVATION_STEP_BYTES: usize = 1024 * 1024; + /// A partitioner that uses a hash function to partition data into multiple partitions pub(crate) struct MultiPartitionShuffleRepartitioner { buffered_batches: Vec, @@ -112,6 +122,10 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { batch_size: usize, /// Reservation for repartitioning reservation: MemoryReservation, + /// 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 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. @@ -216,6 +230,7 @@ impl MultiPartitionShuffleRepartitioner { scratch, batch_size, reservation, + buffered_bytes: 0, max_buffer_bytes, tracing_enabled, pinned_buffers: HashSet::new(), @@ -457,13 +472,16 @@ impl MultiPartitionShuffleRepartitioner { 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 .max_buffer_bytes - .is_some_and(|limit| self.reservation.size() >= limit) + .is_some_and(|limit| self.buffered_bytes >= limit) { self.spill()?; } @@ -471,8 +489,40 @@ impl MultiPartitionShuffleRepartitioner { Ok(()) } + /// Bytes to grow [`Self::reservation`] by at a time. Capped well below `max_buffer_bytes` so + /// that a writer with a small limit does not reserve a multiple of its own budget. + fn reservation_step(&self) -> usize { + match self.max_buffer_bytes { + Some(limit) => RESERVATION_STEP_BYTES.min(limit / 8).max(1), + None => RESERVATION_STEP_BYTES, + } + } + + /// 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 +576,7 @@ impl MultiPartitionShuffleRepartitioner { } self.reservation.free(); + self.buffered_bytes = 0; self.pinned_buffers.clear(); self.metrics.spill_count.add(1); Ok(()) diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 630e0e33431..58ac87ad066 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -277,6 +277,7 @@ mod test { use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::config::SessionConfig; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, MemoryReservation}; use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion::physical_expr::expressions::{col, Column}; use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; @@ -285,6 +286,7 @@ mod test { use datafusion::prelude::SessionContext; use itertools::Itertools; use std::io::Cursor; + use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::runtime::Runtime; #[test] @@ -504,6 +506,122 @@ mod test { ); } + /// A `MemoryPool` that counts the `try_grow` calls reaching it, standing in for Comet's + /// unified pool where each call is a JNI round-trip into Spark's memory manager. + #[derive(Debug)] + struct CountingMemoryPool { + inner: GreedyMemoryPool, + grow_calls: Arc, + } + + impl std::fmt::Display for CountingMemoryPool { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "CountingMemoryPool({})", self.inner) + } + } + + impl MemoryPool for CountingMemoryPool { + fn name(&self) -> &str { + "CountingMemoryPool" + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.grow_calls.fetch_add(1, Ordering::Relaxed); + self.inner.grow(reservation, additional) + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink) + } + + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> datafusion::common::Result<()> { + self.grow_calls.fetch_add(1, Ordering::Relaxed); + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + } + + /// Buffer `num_batches` distinct batches of `batch_rows` rows each through a + /// `MultiPartitionShuffleRepartitioner` and return how many `try_grow` calls reached the pool. + /// + /// The batches are built separately rather than cloned so that each one pins its own buffers; + /// `count_new_buffers` dedups by address, so re-inserting one batch would charge nothing after + /// the first insert and would not exercise the per-batch growth this is measuring. + async fn grow_calls_for_batches(batch_rows: usize, num_batches: usize) -> usize { + let grow_calls = Arc::new(AtomicUsize::new(0)); + // Far larger than anything these batches can reserve, so no request is ever denied. + let pool = CountingMemoryPool { + inner: GreedyMemoryPool::new(1024 * 1024 * 1024), + grow_calls: Arc::clone(&grow_calls), + }; + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(pool)) + .build() + .unwrap(), + ); + + let batches: Vec = + (0..num_batches).map(|_| create_batch(batch_rows)).collect(); + let num_partitions = 4; + let metrics_set = ExecutionPlanMetricsSet::new(); + let metrics = ShufflePartitionerMetrics::new(&metrics_set, 0); + let dir = tempfile::tempdir().unwrap(); + let shuffle_block_writer = + ShuffleBlockWriter::try_new(batches[0].schema().as_ref(), CompressionCodec::Lz4Frame) + .unwrap(); + let local_partition_writer = LocalPartitionWriter::try_new( + dir.path().join("data.out").to_str().unwrap().to_string(), + dir.path().join("index.out").to_str().unwrap().to_string(), + shuffle_block_writer, + num_partitions, + 1024, + 1024 * 1024, + Arc::clone(&runtime_env), + ) + .unwrap(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + local_partition_writer, + CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], num_partitions), + metrics, + runtime_env, + 1024, + false, + None, + ) + .unwrap(); + + // The partition writer reserves as well, so count only what the repartitioner adds. + let before = grow_calls.load(Ordering::Relaxed); + for batch in batches { + repartitioner.insert_batch(batch).await.unwrap(); + } + let calls = grow_calls.load(Ordering::Relaxed) - before; + repartitioner.shuffle_write().unwrap(); + calls + } + + #[tokio::test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + async fn small_batches_share_one_reservation_step() { + // 40 batches of 100 rows buffer far less than one 1 MB step between them, so the writer + // must reach the pool a handful of times rather than once per batch. Growing by the exact + // amount each time costs 40+ acquisitions here. + let calls = grow_calls_for_batches(100, 40).await; + assert!( + calls <= 4, + "40 small batches must not need one acquisition each, got {calls}" + ); + } + /// Buffer `num_batches` batches through a `MultiPartitionShuffleRepartitioner` configured /// with `max_buffer_bytes`, against a memory pool large enough that `try_grow` never fails, /// and return how many times it spilled. From 0d5450bad1d4b70a7fa1bf9d3e1d9b2b3d9f6d67 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 17 Aug 2026 13:56:05 -0600 Subject: [PATCH 2/3] feat: make the shuffle writer's reservation step configurable Adds `spark.comet.shuffle.native.reservationStepBytes`, defaulting to the 1 MB that was hardcoded, and threads it to the writer alongside the existing `maxBufferBytes` limit. Zero grows the reservation by exactly what each batch needs, which is how the writer behaved before it grew in steps. The two knobs now travel together in a `ShuffleWriterMemoryConfig` rather than as more positional `usize` and `Option` arguments, since `ShuffleWriterExec::try_new` and `external_shuffle` already carry a `write_buffer_size` that is easy to confuse with them. --- native/core/src/execution/planner.rs | 10 +- native/proto/src/proto/operator.proto | 4 + native/shuffle/benches/shuffle_writer.rs | 3 +- native/shuffle/src/bin/shuffle_bench.rs | 9 +- native/shuffle/src/lib.rs | 4 +- .../src/partitioners/multi_partition.rs | 35 +++--- native/shuffle/src/shuffle_writer.rs | 101 ++++++++++++++---- .../scala/org/apache/comet/CometConf.scala | 16 +++ .../shuffle/CometNativeShuffleWriter.scala | 2 + 9 files changed, 136 insertions(+), 48 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index c179c3b57c5..52ac5e63eb7 100644 --- a/native/core/src/execution/planner.rs +++ b/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(( diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ced87262f32..ac951f90647 100644 --- a/native/proto/src/proto/operator.proto +++ b/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 { diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index c934226b623..830c0731d0f 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/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() } diff --git a/native/shuffle/src/bin/shuffle_bench.rs b/native/shuffle/src/bin/shuffle_bench.rs index cd43d41dbc1..cfe13ab9d14 100644 --- a/native/shuffle/src/bin/shuffle_bench.rs +++ b/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"); diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 2263ae0dac0..05d2ded8bb8 100644 --- a/native/shuffle/src/lib.rs +++ b/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}; diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index b7a0ef71b7e..4fb8341e66c 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/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}; @@ -98,16 +99,6 @@ impl ScratchSpace { } } -/// Granularity for growing [`MultiPartitionShuffleRepartitioner::reservation`]. -/// -/// The writer's per-batch growth is small — one Arrow buffer per column, 32 KB for an 8192-row -/// `i32` column — and against Comet's unified memory pool every `try_grow` is a JNI round-trip -/// into Spark's memory manager. On a 201-partition shuffle of 10M rows that was 256 acquisitions -/// per task, 96% of them for 32 KB, at ~1.3 µs each (see -/// ). Growing in coarse steps took that -/// to 13 acquisitions per task, trading up to one step of over-reservation per active writer. -const RESERVATION_STEP_BYTES: usize = 1024 * 1024; - /// A partitioner that uses a hash function to partition data into multiple partitions pub(crate) struct MultiPartitionShuffleRepartitioner { buffered_batches: Vec, @@ -126,10 +117,8 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { /// 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 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, + /// 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 @@ -192,7 +181,7 @@ impl MultiPartitionShuffleRepartitioner { runtime: Arc, batch_size: usize, tracing_enabled: bool, - max_buffer_bytes: Option, + memory_config: ShuffleWriterMemoryConfig, ) -> datafusion::common::Result { let num_output_partitions = partitioning.partition_count(); assert_ne!( @@ -231,7 +220,7 @@ impl MultiPartitionShuffleRepartitioner { batch_size, reservation, buffered_bytes: 0, - max_buffer_bytes, + memory_config, tracing_enabled, pinned_buffers: HashSet::new(), }) @@ -480,6 +469,7 @@ impl MultiPartitionShuffleRepartitioner { // step boundary and would otherwise trip the limit early. if self.grow_reservation() || self + .memory_config .max_buffer_bytes .is_some_and(|limit| self.buffered_bytes >= limit) { @@ -489,12 +479,15 @@ impl MultiPartitionShuffleRepartitioner { Ok(()) } - /// Bytes to grow [`Self::reservation`] by at a time. Capped well below `max_buffer_bytes` so - /// that a writer with a small limit does not reserve a multiple of its own budget. + /// 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 { - match self.max_buffer_bytes { - Some(limit) => RESERVATION_STEP_BYTES.min(limit / 8).max(1), - None => RESERVATION_STEP_BYTES, + let step = self.memory_config.reservation_step_bytes; + match self.memory_config.max_buffer_bytes { + Some(limit) => step.min(limit / 8), + None => step, } } diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 58ac87ad066..26f58b245c5 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -46,6 +46,35 @@ use std::{ sync::Arc, }; +/// Default for [`ShuffleWriterMemoryConfig::reservation_step_bytes`], used when the writer is +/// constructed outside the planner (benchmarks and tests). Production values come from +/// `spark.comet.shuffle.native.reservationStepBytes`. +pub const DEFAULT_RESERVATION_STEP_BYTES: usize = 1024 * 1024; + +/// How the shuffle writer manages the memory it buffers batches in. +#[derive(Debug, Clone, Copy)] +pub struct ShuffleWriterMemoryConfig { + /// Spill once the writer buffers this many bytes, independently of whether the memory pool + /// still has capacity. `None` disables the limit, leaving pool pressure as the only trigger. + pub max_buffer_bytes: Option, + /// Granularity for growing the writer's memory reservation. The writer's per-batch growth is + /// small, and against Comet's unified memory pool each request is a JNI round-trip into + /// Spark's memory manager, so memory is reserved in multiples of this and a run of small + /// batches costs one request instead of one each. Zero reserves exactly what each batch needs. + /// Capped at one eighth of `max_buffer_bytes` when that limit is set, so a writer with a small + /// limit does not reserve a multiple of its own budget. + pub reservation_step_bytes: usize, +} + +impl Default for ShuffleWriterMemoryConfig { + fn default() -> Self { + Self { + max_buffer_bytes: None, + reservation_step_bytes: DEFAULT_RESERVATION_STEP_BYTES, + } + } +} + /// The shuffle writer operator maps each input partition to M output partitions based on a /// partitioning scheme. No guarantees are made about the order of the resulting partitions. #[derive(Debug)] @@ -67,8 +96,8 @@ pub struct ShuffleWriterExec { tracing_enabled: bool, /// Size of the write buffer in bytes write_buffer_size: usize, - /// Maximum bytes buffered in memory before spilling; `None` disables the limit - max_buffer_bytes: Option, + /// How the writer manages the memory it buffers batches in + memory_config: ShuffleWriterMemoryConfig, } impl ShuffleWriterExec { @@ -82,7 +111,7 @@ impl ShuffleWriterExec { output_index_file: String, tracing_enabled: bool, write_buffer_size: usize, - max_buffer_bytes: Option, + memory_config: ShuffleWriterMemoryConfig, ) -> Result { let cache = Arc::new(PlanProperties::new( EquivalenceProperties::new(Arc::clone(&input.schema())), @@ -101,7 +130,7 @@ impl ShuffleWriterExec { codec, tracing_enabled, write_buffer_size, - max_buffer_bytes, + memory_config, }) } } @@ -157,7 +186,7 @@ impl ExecutionPlan for ShuffleWriterExec { self.output_index_file.clone(), self.tracing_enabled, self.write_buffer_size, - self.max_buffer_bytes, + self.memory_config, )?)), _ => panic!("ShuffleWriterExec wrong number of children"), } @@ -187,7 +216,7 @@ impl ExecutionPlan for ShuffleWriterExec { self.codec.clone(), self.tracing_enabled, self.write_buffer_size, - self.max_buffer_bytes, + self.memory_config, )) .try_flatten(), ))) @@ -206,7 +235,7 @@ async fn external_shuffle( codec: CompressionCodec, tracing_enabled: bool, write_buffer_size: usize, - max_buffer_bytes: Option, + memory_config: ShuffleWriterMemoryConfig, ) -> Result { let schema = input.schema(); @@ -243,7 +272,7 @@ async fn external_shuffle( context.runtime_env(), context.session_config().batch_size(), tracing_enabled, - max_buffer_bytes, + memory_config, )?), }; @@ -416,7 +445,7 @@ mod test { runtime_env, 1024, false, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -488,7 +517,7 @@ mod test { runtime_env, batch_size, false, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -554,7 +583,11 @@ mod test { /// The batches are built separately rather than cloned so that each one pins its own buffers; /// `count_new_buffers` dedups by address, so re-inserting one batch would charge nothing after /// the first insert and would not exercise the per-batch growth this is measuring. - async fn grow_calls_for_batches(batch_rows: usize, num_batches: usize) -> usize { + async fn grow_calls_for_batches( + batch_rows: usize, + num_batches: usize, + memory_config: ShuffleWriterMemoryConfig, + ) -> usize { let grow_calls = Arc::new(AtomicUsize::new(0)); // Far larger than anything these batches can reserve, so no request is ever denied. let pool = CountingMemoryPool { @@ -595,7 +628,7 @@ mod test { runtime_env, 1024, false, - None, + memory_config, ) .unwrap(); @@ -613,15 +646,35 @@ mod test { #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` async fn small_batches_share_one_reservation_step() { // 40 batches of 100 rows buffer far less than one 1 MB step between them, so the writer - // must reach the pool a handful of times rather than once per batch. Growing by the exact - // amount each time costs 40+ acquisitions here. - let calls = grow_calls_for_batches(100, 40).await; + // must reach the pool a handful of times rather than once per batch. + let calls = grow_calls_for_batches(100, 40, ShuffleWriterMemoryConfig::default()).await; assert!( calls <= 4, "40 small batches must not need one acquisition each, got {calls}" ); } + #[tokio::test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + async fn zero_reservation_step_reserves_per_batch() { + // A zero step disables the rounding, which is the behaviour the writer had before it grew + // in steps: one acquisition per batch. Paired with the test above, this pins that the + // reduction comes from the step and not from something else in the insert path. + let calls = grow_calls_for_batches( + 100, + 40, + ShuffleWriterMemoryConfig { + reservation_step_bytes: 0, + ..Default::default() + }, + ) + .await; + assert_eq!( + calls, 40, + "a zero step must reserve exactly what each batch needs" + ); + } + /// Buffer `num_batches` batches through a `MultiPartitionShuffleRepartitioner` configured /// with `max_buffer_bytes`, against a memory pool large enough that `try_grow` never fails, /// and return how many times it spilled. @@ -659,7 +712,10 @@ mod test { runtime_env, 1024, false, - max_buffer_bytes, + ShuffleWriterMemoryConfig { + max_buffer_bytes, + ..Default::default() + }, ) .unwrap(); @@ -717,7 +773,10 @@ mod test { index_file.to_str().unwrap().to_string(), false, 1024 * 1024, - max_buffer_bytes, + ShuffleWriterMemoryConfig { + max_buffer_bytes, + ..Default::default() + }, ) .unwrap(); @@ -842,7 +901,7 @@ mod test { "/tmp/index.out".to_string(), false, 1024 * 1024, // write_buffer_size: 1MB default - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -902,7 +961,7 @@ mod test { index_file.clone(), false, 1024 * 1024, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -1277,7 +1336,7 @@ mod test { index_file.to_str().unwrap().to_string(), false, 1024 * 1024, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -1366,7 +1425,7 @@ mod test { index_file.to_str().unwrap().to_string(), false, 1024 * 1024, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6e7bb4271bb..9f7d7c36f69 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -595,6 +595,22 @@ object CometConf extends ShimCometConf { .checkValue(v => v >= 0, "Must not be negative") .createWithDefault(0) + val COMET_SHUFFLE_NATIVE_RESERVATION_STEP_BYTES: ConfigEntry[Long] = + conf("spark.comet.shuffle.native.reservationStepBytes") + .category(CATEGORY_SHUFFLE) + .doc( + "Granularity with which the native shuffle writer grows its memory reservation. The " + + "writer's per-batch growth is small, and in off-heap mode each request to the memory " + + "pool is a JNI call into Spark's memory manager, so memory is reserved in multiples " + + "of this value and a run of small batches costs one request instead of one each. The " + + "cost is up to this much reservation held beyond what the writer is using, per " + + s"concurrent task. Zero grows the reservation by exactly what each batch needs. When " + + s"${COMET_SHUFFLE_NATIVE_MAX_BUFFER_BYTES.key} is set, the step is capped at one " + + "eighth of that limit.") + .bytesConf(ByteUnit.BYTE) + .checkValue(v => v >= 0, "Must not be negative") + .createWithDefault(1024 * 1024) + val COMET_DEBUG_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.debug.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index 5b1145e9e07..fb04817c47f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -235,6 +235,8 @@ class CometNativeShuffleWriter[K, V]( shuffleWriterBuilder.setWriteBufferSize( CometConf.COMET_SHUFFLE_NATIVE_WRITE_BUFFER_SIZE.get().min(Int.MaxValue).toInt) shuffleWriterBuilder.setMaxBufferBytes(CometConf.COMET_SHUFFLE_NATIVE_MAX_BUFFER_BYTES.get()) + shuffleWriterBuilder.setReservationStepBytes( + CometConf.COMET_SHUFFLE_NATIVE_RESERVATION_STEP_BYTES.get()) outputPartitioning match { case p if isSinglePartitioning(p) => From b886d00e7f7921d7409c8898854a2b40ffdd3415 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 18 Aug 2026 09:03:50 -0600 Subject: [PATCH 3/3] fix: drop redundant string interpolator in the new config's doc The first line of the doc string had an `s` prefix but no interpolation, which scalafix's RedundantSyntax rule rejects (Lint Scala and Lint Java both run it). --- spark/src/main/scala/org/apache/comet/CometConf.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 9f7d7c36f69..a7ea23d0d1a 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -604,7 +604,7 @@ object CometConf extends ShimCometConf { "pool is a JNI call into Spark's memory manager, so memory is reserved in multiples " + "of this value and a run of small batches costs one request instead of one each. The " + "cost is up to this much reservation held beyond what the writer is using, per " + - s"concurrent task. Zero grows the reservation by exactly what each batch needs. When " + + "concurrent task. Zero grows the reservation by exactly what each batch needs. When " + s"${COMET_SHUFFLE_NATIVE_MAX_BUFFER_BYTES.key} is set, the step is capped at one " + "eighth of that limit.") .bytesConf(ByteUnit.BYTE)