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 5644b1085d7..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}; @@ -112,10 +113,12 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { 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, + /// 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 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!( @@ -216,7 +219,8 @@ impl MultiPartitionShuffleRepartitioner { scratch, batch_size, reservation, - max_buffer_bytes, + buffered_bytes: 0, + memory_config, tracing_enabled, pinned_buffers: HashSet::new(), }) @@ -457,13 +461,17 @@ 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 + .memory_config .max_buffer_bytes - .is_some_and(|limit| self.reservation.size() >= limit) + .is_some_and(|limit| self.buffered_bytes >= limit) { self.spill()?; } @@ -471,8 +479,43 @@ impl MultiPartitionShuffleRepartitioner { 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 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..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, )?), }; @@ -277,6 +306,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 +315,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] @@ -414,7 +445,7 @@ mod test { runtime_env, 1024, false, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -486,7 +517,7 @@ mod test { runtime_env, batch_size, false, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -504,6 +535,146 @@ 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, + 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 { + 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, + memory_config, + ) + .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. + 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. @@ -541,7 +712,10 @@ mod test { runtime_env, 1024, false, - max_buffer_bytes, + ShuffleWriterMemoryConfig { + max_buffer_bytes, + ..Default::default() + }, ) .unwrap(); @@ -599,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(); @@ -724,7 +901,7 @@ mod test { "/tmp/index.out".to_string(), false, 1024 * 1024, // write_buffer_size: 1MB default - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -784,7 +961,7 @@ mod test { index_file.clone(), false, 1024 * 1024, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -1159,7 +1336,7 @@ mod test { index_file.to_str().unwrap().to_string(), false, 1024 * 1024, - None, + ShuffleWriterMemoryConfig::default(), ) .unwrap(); @@ -1248,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..a7ea23d0d1a 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 " + + "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) =>