This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 1981091602 feat: add spark.comet.shuffle.maxBufferBytes to cap native
shuffle writer memory (#4989)
1981091602 is described below
commit 198109160221e41cc97ecd7e415971006fa3b637
Author: Andy Grove <[email protected]>
AuthorDate: Wed Jul 22 06:42:40 2026 -0600
feat: add spark.comet.shuffle.maxBufferBytes to cap native shuffle writer
memory (#4989)
---
docs/source/contributor-guide/native_shuffle.md | 4 +-
native/core/src/execution/planner.rs | 5 +
native/proto/src/proto/operator.proto | 3 +
native/shuffle/benches/shuffle_writer.rs | 1 +
native/shuffle/src/bin/shuffle_bench.rs | 9 ++
native/shuffle/src/partitioners/multi_partition.rs | 15 +-
native/shuffle/src/shuffle_writer.rs | 177 ++++++++++++++++++++-
.../main/scala/org/apache/comet/CometConf.scala | 13 ++
.../shuffle/CometNativeShuffleWriter.scala | 1 +
.../comet/exec/CometNativeShuffleSuite.scala | 34 ++++
10 files changed, 253 insertions(+), 9 deletions(-)
diff --git a/docs/source/contributor-guide/native_shuffle.md
b/docs/source/contributor-guide/native_shuffle.md
index c46b8b45a8..4d6421de5e 100644
--- a/docs/source/contributor-guide/native_shuffle.md
+++ b/docs/source/contributor-guide/native_shuffle.md
@@ -214,7 +214,9 @@ sizes.
Native shuffle uses DataFusion's memory management with spilling support:
- **Memory pool**: Tracks memory usage across the shuffle operation.
-- **Spill threshold**: When buffered data exceeds the threshold, partitions
spill to disk.
+- **Spill triggers**: Partitions spill to disk when the memory pool denies an
allocation, or
+ when the buffered bytes reach `spark.comet.shuffle.maxBufferBytes`. That
config defaults to
+ 0, which disables the fixed limit and leaves memory pressure as the only
trigger.
- **Per-partition spilling**: Each partition has its own spill file. Multiple
spills for a
partition are concatenated when writing the final output.
- **Scratch space**: Reusable buffers for partition ID computation to reduce
allocations.
diff --git a/native/core/src/execution/planner.rs
b/native/core/src/execution/planner.rs
index 33e7df9387..339bfff156 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -1633,6 +1633,10 @@ impl PhysicalPlanner {
}?;
let write_buffer_size = writer.write_buffer_size as usize;
+ // Zero on the wire means the limit is disabled; normalize it
here so the writer
+ // 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 shuffle_writer = Arc::new(ShuffleWriterExec::try_new(
writer_input,
partitioning,
@@ -1641,6 +1645,7 @@ impl PhysicalPlanner {
writer.output_index_file.clone(),
writer.tracing_enabled,
write_buffer_size,
+ max_buffer_bytes,
)?);
Ok((
diff --git a/native/proto/src/proto/operator.proto
b/native/proto/src/proto/operator.proto
index 00158644d1..d00979f0a0 100644
--- a/native/proto/src/proto/operator.proto
+++ b/native/proto/src/proto/operator.proto
@@ -340,6 +340,9 @@ message ShuffleWriter {
// to absorb DataFusion-vs-Spark type drift. Empty when the child is a
placeholder Scan;
// that path already has a cast point upstream.
repeated SparkStructField expected_output_schema = 9;
+ // 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;
}
message ParquetWriter {
diff --git a/native/shuffle/benches/shuffle_writer.rs
b/native/shuffle/benches/shuffle_writer.rs
index 27abd919fa..ea6ad83fdf 100644
--- a/native/shuffle/benches/shuffle_writer.rs
+++ b/native/shuffle/benches/shuffle_writer.rs
@@ -153,6 +153,7 @@ fn create_shuffle_writer_exec(
"/tmp/index.out".to_string(),
false,
1024 * 1024,
+ None,
)
.unwrap()
}
diff --git a/native/shuffle/src/bin/shuffle_bench.rs
b/native/shuffle/src/bin/shuffle_bench.rs
index cc1c960848..cd43d41dbc 100644
--- a/native/shuffle/src/bin/shuffle_bench.rs
+++ b/native/shuffle/src/bin/shuffle_bench.rs
@@ -106,6 +106,10 @@ struct Args {
#[arg(long, default_value_t = 1048576)]
write_buffer_size: usize,
+ /// Maximum bytes buffered in memory before spilling (unset = spill on
memory pressure only)
+ #[arg(long)]
+ max_buffer_bytes: Option<usize>,
+
/// Limit rows processed per iteration (0 = no limit)
#[arg(long, default_value_t = 0)]
limit: usize,
@@ -413,6 +417,7 @@ fn run_shuffle_write(
args.batch_size,
args.memory_limit,
args.write_buffer_size,
+ args.max_buffer_bytes,
args.limit,
data_file.to_string(),
index_file.to_string(),
@@ -436,6 +441,7 @@ async fn execute_shuffle_write(
batch_size: usize,
memory_limit: Option<usize>,
write_buffer_size: usize,
+ max_buffer_bytes: Option<usize>,
limit: usize,
data_file: String,
index_file: String,
@@ -480,6 +486,7 @@ async fn execute_shuffle_write(
index_file,
false,
write_buffer_size,
+ max_buffer_bytes,
)
.expect("Failed to create ShuffleWriterExec");
@@ -543,6 +550,7 @@ fn run_concurrent_shuffle_writes(
let batch_size = args.batch_size;
let memory_limit = args.memory_limit;
let write_buffer_size = args.write_buffer_size;
+ let max_buffer_bytes = args.max_buffer_bytes;
let limit = args.limit;
handles.push(tokio::spawn(async move {
@@ -553,6 +561,7 @@ fn run_concurrent_shuffle_writes(
batch_size,
memory_limit,
write_buffer_size,
+ max_buffer_bytes,
limit,
data_file,
index_file,
diff --git a/native/shuffle/src/partitioners/multi_partition.rs
b/native/shuffle/src/partitioners/multi_partition.rs
index be2d977e83..5644b1085d 100644
--- a/native/shuffle/src/partitioners/multi_partition.rs
+++ b/native/shuffle/src/partitioners/multi_partition.rs
@@ -112,6 +112,10 @@ 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>,
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
@@ -174,6 +178,7 @@ impl<T: PartitionWriter>
MultiPartitionShuffleRepartitioner<T> {
runtime: Arc<RuntimeEnv>,
batch_size: usize,
tracing_enabled: bool,
+ max_buffer_bytes: Option<usize>,
) -> datafusion::common::Result<Self> {
let num_output_partitions = partitioning.partition_count();
assert_ne!(
@@ -211,6 +216,7 @@ impl<T: PartitionWriter>
MultiPartitionShuffleRepartitioner<T> {
scratch,
batch_size,
reservation,
+ max_buffer_bytes,
tracing_enabled,
pinned_buffers: HashSet::new(),
})
@@ -451,7 +457,14 @@ impl<T: PartitionWriter>
MultiPartitionShuffleRepartitioner<T> {
mem_growth += after_size.saturating_sub(before_size);
}
- if self.reservation.try_grow(mem_growth).is_err() {
+ // `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
+ .max_buffer_bytes
+ .is_some_and(|limit| self.reservation.size() >= limit)
+ {
self.spill()?;
}
diff --git a/native/shuffle/src/shuffle_writer.rs
b/native/shuffle/src/shuffle_writer.rs
index de6989e52a..8d1336dc1c 100644
--- a/native/shuffle/src/shuffle_writer.rs
+++ b/native/shuffle/src/shuffle_writer.rs
@@ -67,6 +67,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<usize>,
}
impl ShuffleWriterExec {
@@ -80,6 +82,7 @@ impl ShuffleWriterExec {
output_index_file: String,
tracing_enabled: bool,
write_buffer_size: usize,
+ max_buffer_bytes: Option<usize>,
) -> Result<Self> {
let cache = Arc::new(PlanProperties::new(
EquivalenceProperties::new(Arc::clone(&input.schema())),
@@ -98,6 +101,7 @@ impl ShuffleWriterExec {
codec,
tracing_enabled,
write_buffer_size,
+ max_buffer_bytes,
})
}
}
@@ -153,6 +157,7 @@ impl ExecutionPlan for ShuffleWriterExec {
self.output_index_file.clone(),
self.tracing_enabled,
self.write_buffer_size,
+ self.max_buffer_bytes,
)?)),
_ => panic!("ShuffleWriterExec wrong number of children"),
}
@@ -182,6 +187,7 @@ impl ExecutionPlan for ShuffleWriterExec {
self.codec.clone(),
self.tracing_enabled,
self.write_buffer_size,
+ self.max_buffer_bytes,
))
.try_flatten(),
)))
@@ -200,6 +206,7 @@ async fn external_shuffle(
codec: CompressionCodec,
tracing_enabled: bool,
write_buffer_size: usize,
+ max_buffer_bytes: Option<usize>,
) -> Result<SendableRecordBatchStream> {
let schema = input.schema();
@@ -237,6 +244,7 @@ async fn external_shuffle(
context.runtime_env(),
context.session_config().batch_size(),
tracing_enabled,
+ max_buffer_bytes,
)?),
};
@@ -369,6 +377,7 @@ mod test {
runtime_env,
1024,
false,
+ None,
)
.unwrap();
@@ -440,6 +449,7 @@ mod test {
runtime_env,
batch_size,
false,
+ None,
)
.unwrap();
@@ -457,6 +467,148 @@ mod test {
);
}
+ /// 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.
+ async fn spill_count_with_max_buffer_bytes(
+ max_buffer_bytes: Option<usize>,
+ num_batches: usize,
+ ) -> usize {
+ let batch = create_batch(1000);
+ let num_partitions = 4;
+ // Far larger than anything these batches can reserve, so pool
pressure never triggers
+ // a spill and any spill observed must come from the max_buffer_bytes
limit.
+ let runtime_env = create_runtime(1024 * 1024 * 1024);
+ let metrics_set = ExecutionPlanMetricsSet::new();
+ let metrics = ShufflePartitionerMetrics::new(&metrics_set, 0);
+ let spill_count = metrics.spill_count.clone();
+ let dir = tempfile::tempdir().unwrap();
+ let shuffle_block_writer =
+ ShuffleBlockWriter::try_new(batch.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,
+ max_buffer_bytes,
+ )
+ .unwrap();
+
+ for _ in 0..num_batches {
+ repartitioner.insert_batch(batch.clone()).await.unwrap();
+ }
+ repartitioner.shuffle_write().unwrap();
+
+ spill_count.value()
+ }
+
+ #[tokio::test]
+ #[cfg_attr(miri, ignore)] // miri can't call foreign function
`ZSTD_createCCtx`
+ async fn max_buffer_bytes_triggers_spill_without_memory_pressure() {
+ // A limit far below what 20 batches buffer must spill, even though
the pool never
+ // denies an allocation. The paired zero case pins that the spills
come from the new
+ // limit and not from the pre-existing memory-pressure trigger.
+ let spills = spill_count_with_max_buffer_bytes(Some(8 * 1024),
20).await;
+ assert!(
+ spills > 0,
+ "a max_buffer_bytes limit below the buffered size must trigger
spilling, got {spills}"
+ );
+ }
+
+ #[tokio::test]
+ #[cfg_attr(miri, ignore)] // miri can't call foreign function
`ZSTD_createCCtx`
+ async fn max_buffer_bytes_none_leaves_spilling_to_memory_pressure() {
+ let spills = spill_count_with_max_buffer_bytes(None, 20).await;
+ assert_eq!(
+ spills, 0,
+ "an unset max_buffer_bytes disables the limit, and this pool never
denies an allocation"
+ );
+ }
+
+ /// Run a shuffle end-to-end through `ShuffleWriterExec` with the given
`max_buffer_bytes`
+ /// and return the shuffled rows in output-file order.
+ fn shuffle_output_with_max_buffer_bytes(
+ dir: &std::path::Path,
+ tag: &str,
+ max_buffer_bytes: Option<usize>,
+ ) -> Vec<String> {
+ let batch = create_batch(1000);
+ let batches = (0..20).map(|_| batch.clone()).collect::<Vec<_>>();
+ let data_file = dir.join(format!("{tag}_data.out"));
+ let index_file = dir.join(format!("{tag}_index.out"));
+
+ let exec = ShuffleWriterExec::try_new(
+ Arc::new(DataSourceExec::new(Arc::new(
+ MemorySourceConfig::try_new(std::slice::from_ref(&batches),
batch.schema(), None)
+ .unwrap(),
+ ))),
+ CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16),
+ CompressionCodec::Zstd(1),
+ data_file.to_str().unwrap().to_string(),
+ index_file.to_str().unwrap().to_string(),
+ false,
+ 1024 * 1024,
+ max_buffer_bytes,
+ )
+ .unwrap();
+
+ // Large enough that the pool never denies an allocation, so only
max_buffer_bytes
+ // decides whether this run spills.
+ let runtime_env = create_runtime(1024 * 1024 * 1024);
+ let ctx = SessionContext::new_with_config_rt(SessionConfig::new(),
runtime_env);
+ let stream = exec.execute(0, ctx.task_ctx()).unwrap();
+ Runtime::new().unwrap().block_on(collect(stream)).unwrap();
+
+ read_all_ipc_batches(&std::fs::read(&data_file).unwrap())
+ .iter()
+ .flat_map(|batch| {
+ let column = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ column
+ .iter()
+ .map(|v| v.unwrap().to_string())
+ .collect::<Vec<_>>()
+ })
+ .collect()
+ }
+
+ #[test]
+ #[cfg_attr(miri, ignore)] // miri can't call foreign function
`ZSTD_createCCtx`
+ fn max_buffer_bytes_preserves_output() {
+ // End-to-end through ShuffleWriterExec: a small limit forces repeated
spills, and the
+ // shuffled rows must come out in the same order as an unlimited
writer produces.
+ let dir = tempfile::tempdir().unwrap();
+ let unlimited = shuffle_output_with_max_buffer_bytes(dir.path(),
"unlimited", None);
+ let limited = shuffle_output_with_max_buffer_bytes(dir.path(),
"limited", Some(8 * 1024));
+
+ assert_eq!(
+ unlimited.len(),
+ 20 * 1000,
+ "every input row must be written"
+ );
+ assert_eq!(
+ limited, unlimited,
+ "spilling on the max_buffer_bytes limit must not change the
shuffled output"
+ );
+ }
+
fn create_runtime(memory_limit: usize) -> Arc<RuntimeEnv> {
Arc::new(
RuntimeEnvBuilder::new()
@@ -535,6 +687,7 @@ mod test {
"/tmp/index.out".to_string(),
false,
1024 * 1024, // write_buffer_size: 1MB default
+ None,
)
.unwrap();
@@ -594,6 +747,7 @@ mod test {
index_file.clone(),
false,
1024 * 1024,
+ None,
)
.unwrap();
@@ -741,11 +895,11 @@ mod test {
);
}
- /// Read all IPC blocks from a byte buffer written by
BufBatchWriter/ShuffleBlockWriter,
- /// returning the total number of rows.
- fn read_all_ipc_blocks(data: &[u8]) -> usize {
+ /// Decode every IPC block in a byte buffer written by
BufBatchWriter/ShuffleBlockWriter,
+ /// in file order.
+ fn read_all_ipc_batches(data: &[u8]) -> Vec<RecordBatch> {
let mut offset = 0;
- let mut total_rows = 0;
+ let mut batches = vec![];
while offset < data.len() {
// First 8 bytes are the IPC length (little-endian u64)
let ipc_length =
@@ -756,11 +910,18 @@ mod test {
// read_ipc_compressed expects data starting after the 16-byte
header
// (i.e., after length + field_count), at the codec tag
let ipc_data = &data[block_start + 8..block_end];
- let batch = read_ipc_compressed(ipc_data).unwrap();
- total_rows += batch.num_rows();
+ batches.push(read_ipc_compressed(ipc_data).unwrap());
offset = block_end;
}
- total_rows
+ batches
+ }
+
+ /// Total number of rows across all IPC blocks in a byte buffer.
+ fn read_all_ipc_blocks(data: &[u8]) -> usize {
+ read_all_ipc_batches(data)
+ .iter()
+ .map(|b| b.num_rows())
+ .sum()
}
#[test]
@@ -798,6 +959,7 @@ mod test {
index_file.to_str().unwrap().to_string(),
false,
1024 * 1024,
+ None,
)
.unwrap();
@@ -886,6 +1048,7 @@ mod test {
index_file.to_str().unwrap().to_string(),
false,
1024 * 1024,
+ None,
)
.unwrap();
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala
b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 41c99d2724..1b71486f09 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -552,6 +552,19 @@ object CometConf extends ShimCometConf {
.doubleConf
.createWithDefault(1.0)
+ val COMET_SHUFFLE_MAX_BUFFER_BYTES: ConfigEntry[Long] =
+ conf("spark.comet.shuffle.maxBufferBytes")
+ .category(CATEGORY_SHUFFLE)
+ .doc(
+ "Maximum number of bytes that the native shuffle writer will buffer in
memory " +
+ "before spilling to disk. The default of 0 disables this limit, in
which case " +
+ "spilling is triggered only when the memory pool denies an
allocation. Setting a " +
+ "limit caps the writer's memory footprint independently of the size
of the memory " +
+ "pool, at the cost of more frequent spilling.")
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(v => v >= 0, "Must not be negative")
+ .createWithDefault(0)
+
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 a486e2e861..2f1b241174 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,7 @@ class CometNativeShuffleWriter[K, V](
CometConf.COMET_EXEC_SHUFFLE_COMPRESSION_ZSTD_LEVEL.get)
shuffleWriterBuilder.setWriteBufferSize(
CometConf.COMET_SHUFFLE_WRITE_BUFFER_SIZE.get().min(Int.MaxValue).toInt)
+
shuffleWriterBuilder.setMaxBufferBytes(CometConf.COMET_SHUFFLE_MAX_BUFFER_BYTES.get())
outputPartitioning match {
case p if isSinglePartitioning(p) =>
diff --git
a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
index 7c70c373fe..75c86e3568 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
@@ -446,6 +446,40 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
}
}
+ test("native shuffle: maxBufferBytes triggers spilling") {
+ withParquetTable((0 until 20000).map(i => (i, (i + 1).toLong, s"str$i")),
"tbl") {
+ def spillCountWithMaxBufferBytes(maxBufferBytes: String): Long = {
+ var spillCount = 0L
+ withSQLConf(CometConf.COMET_SHUFFLE_MAX_BUFFER_BYTES.key ->
maxBufferBytes) {
+ val shuffled = sql("SELECT * FROM tbl").repartition(10, $"_1")
+ checkShuffleAnswer(shuffled, 1)
+
+ // Materialize the shuffled data. The metrics live on this plan's
exchange, so the
+ // query has to run before they are read; checkShuffleAnswer
executes copies built
+ // from the logical plan and leaves this one's accumulators
untouched.
+ shuffled.collect()
+ // AdaptiveSparkPlanHelper's collectFirst, not TreeNode's:
AdaptiveSparkPlanExec is a
+ // leaf node, so a plain tree walk stops above the exchange when AQE
is on.
+ spillCount = collectFirst(shuffled.queryExecution.executedPlan) {
+ case e: CometShuffleExchangeExec => e.metrics("spill_count").value
+ }.get
+ }
+ spillCount
+ }
+
+ // A limit far below the size of this input forces the native writer to
spill, without
+ // relying on the memory pool denying an allocation. Compared against
the default of 0
+ // rather than asserted to be an absolute zero, since the default path
is still free to
+ // spill under memory pressure and how much of that happens depends on
the pool size.
+ val limited = spillCountWithMaxBufferBytes("64k")
+ val unlimited = spillCountWithMaxBufferBytes("0")
+ assert(
+ limited > unlimited,
+ "a 64k maxBufferBytes limit must spill more than the default of 0 " +
+ s"(got $limited vs $unlimited)")
+ }
+ }
+
test("native shuffle: round robin partitioning") {
withSQLConf(
CometConf.COMET_EXEC_SHUFFLE_WITH_ROUND_ROBIN_PARTITIONING_ENABLED.key
-> "true") {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
