From 1cbe9c2dea53c72546b1b24ba2402890f8d3fca4 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Wed, 9 Sep 2026 23:20:37 +0800 Subject: [PATCH 1/8] perf: return native shuffle partition offsets over JNI instead of an index file The native shuffle writer knew every partition offset by the time it finished a map task, but handed them to the JVM through a temporary file: LocalPartitionWriter::finish_all created an index file and wrote num_output_partitions + 1 little-endian i64 offsets into it, and CometNativeShuffleWriter read the whole file back with Files.readAllBytes, converted the offsets to lengths, deleted it, and passed the lengths to IndexShuffleBlockResolver.writeMetadataFileAndCommit, which writes Spark's real index file. The temp file existed only to move an array of longs across the JNI boundary, and cost every map task a create, write, read and unlink on top of the index file Spark commits anyway. The parse also allocated an intermediate array and a ByteBuffer per partition (item 5 of #5198), which goes away with the file. The offsets are now published in memory through a PartitionOffsets slot shared by the writer and its ShuffleWriterDestination, and read back over JNI by Native.getShufflePartitionOffsets. The index path no longer travels in the plan, so LocalPartitionWriter.output_index_file and the legacy ShuffleWriter.output_index_file are removed and their field numbers reserved. The offsets have to be read while the native plan is still alive. CometExecIterator closes itself when its stream reaches the end, and close releases the execution context that owns the writer, so reading after drainAndClose returned freed memory and produced garbage lengths. The iterator instead captures the offsets at end of stream, before close, when built with capturePartitionOffsets, which only the local destination sets: RSS reports its partition lengths through its pusher. Partition lengths are derived from effectivePartitionCount, the output partition count, not the numParts constructor argument, which is the input partition count. Co-Authored-By: Claude Opus 5 --- native/core/src/execution/jni_api.rs | 61 ++++++- native/core/src/execution/planner.rs | 117 ++----------- native/proto/src/proto/operator.proto | 13 +- native/shuffle/benches/shuffle_writer.rs | 1 - native/shuffle/src/bin/shuffle_bench.rs | 9 - native/shuffle/src/lib.rs | 2 +- native/shuffle/src/rss_execution_tests.rs | 21 ++- native/shuffle/src/shuffle_writer.rs | 155 ++++++++++-------- .../writers/local/local_partition_writer.rs | 41 +++-- .../org/apache/comet/CometExecIterator.scala | 23 ++- .../main/scala/org/apache/comet/Native.scala | 15 ++ .../shuffle/CometNativeShuffleWriter.scala | 63 ++++--- .../comet/exec/CometNativeShuffleSuite.scala | 12 +- ...ometCelebornNativeShuffleWriterSuite.scala | 3 +- 14 files changed, 277 insertions(+), 259 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 65a2d68ec18..5ef4fb44b0b 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -100,7 +100,7 @@ use tokio::sync::mpsc; use crate::execution::memory_pools::{create_memory_pool, parse_memory_pool_config}; use crate::execution::operators::{ScanExec, ShuffleScanExec}; use crate::execution::shuffle::{ - decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec, + decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec, ShuffleWriterExec, }; use crate::execution::spark_plan::SparkPlan; @@ -1155,6 +1155,65 @@ fn get_execution_context<'a>(id: i64) -> &'a mut ExecutionContext { } } +/// Returns the partition offsets published by a finished native shuffle write. +/// +/// The writer knows every offset by the time its plan completes, and the only consumer is the +/// Spark task driving that plan, so the offsets are handed back in memory rather than serialized +/// to a temporary index file and read back. Call after the plan has been fully drained; the +/// offsets are not published until the writer finishes. +/// +/// The returned array holds `num_output_partitions + 1` offsets, the last being the total data +/// file length, so partition lengths are successive differences. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_Native_getShufflePartitionOffsets( + e: EnvUnowned, + _class: JClass, + exec_context: jlong, +) -> jlongArray { + try_unwrap_or_throw(&e, |env| { + let context = get_execution_context(exec_context); + + let root_op = context.root_op.as_ref().ok_or_else(|| { + CometError::Internal( + "Cannot read shuffle partition offsets before the plan has been executed" + .to_string(), + ) + })?; + + // `ExecutionPlan` has `Any` as a supertrait but no `as_any` method of its own, so upcast + // the trait object before downcasting to the writer. + let writer = (root_op.native_plan.as_ref() as &dyn std::any::Any) + .downcast_ref::() + .ok_or_else(|| { + CometError::Internal( + "Shuffle partition offsets are only available on a native shuffle write plan" + .to_string(), + ) + })?; + + let offsets = writer + .partition_offsets() + .ok_or_else(|| { + CometError::Internal( + "Shuffle partition offsets are not published by a remote shuffle destination" + .to_string(), + ) + })? + .get() + .ok_or_else(|| { + CometError::Internal( + "Shuffle writer has not published its partition offsets; the plan was not \ + drained to completion" + .to_string(), + ) + })?; + + let long_array = env.new_long_array(offsets.len())?; + long_array.set_region(env, 0, offsets)?; + Ok(long_array.into_raw()) + }) +} + /// Used by Comet shuffle external sorter to write sorted records to disk. /// # Safety /// This function is inherently unsafe since it deals with raw pointers passed from JNI. diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index fe01bdcfeec..0c2938438da 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -47,7 +47,7 @@ use crate::execution::{ planner::expression_registry::ExpressionRegistry, planner::operator_registry::OperatorRegistry, serde::{to_arrow_datatype, to_arrow_field}, - shuffle::{SchemaAlignExec, ShuffleWriterDestination, ShuffleWriterExec}, + shuffle::{PartitionOffsets, SchemaAlignExec, ShuffleWriterDestination, ShuffleWriterExec}, }; use crate::jvm_bridge::{jni_call, JVMClasses, ShufflePartitionPusher}; use arrow::compute::CastOptions; @@ -4143,7 +4143,7 @@ fn shuffle_writer_destination( return Ok(ShuffleWriterDestination::Local { output_data_file: writer.output_data_file.clone(), - output_index_file: writer.output_index_file.clone(), + partition_offsets: Arc::new(PartitionOffsets::default()), }); }; @@ -4155,12 +4155,6 @@ fn shuffle_writer_destination( )); } - if local.output_index_file.is_empty() { - return Err(GeneralError( - "Local shuffle partition writer is missing its output index file".to_string(), - )); - } - if !writer.output_data_file.is_empty() && writer.output_data_file != local.output_data_file { @@ -4171,16 +4165,6 @@ fn shuffle_writer_destination( )); } - if !writer.output_index_file.is_empty() - && writer.output_index_file != local.output_index_file - { - return Err(GeneralError( - "Local shuffle partition writer output index file conflicts with the legacy \ - shuffle output index file" - .to_string(), - )); - } - if shuffle_partition_pusher.is_some() { return Err(GeneralError( "Local shuffle partition writer cannot use a remote shuffle callback" @@ -4190,11 +4174,11 @@ fn shuffle_writer_destination( Ok(ShuffleWriterDestination::Local { output_data_file: local.output_data_file.clone(), - output_index_file: local.output_index_file.clone(), + partition_offsets: Arc::new(PartitionOffsets::default()), }) } Some(spark_operator::partition_writer::Writer::Rss(_)) => { - if !writer.output_data_file.is_empty() || !writer.output_index_file.is_empty() { + if !writer.output_data_file.is_empty() { return Err(GeneralError( "RSS shuffle partition writer cannot have local output files".to_string(), )); @@ -5155,15 +5139,11 @@ mod tests { } } - fn local_shuffle_partition_writer( - output_data_file: &str, - output_index_file: &str, - ) -> spark_operator::PartitionWriter { + fn local_shuffle_partition_writer(output_data_file: &str) -> spark_operator::PartitionWriter { spark_operator::PartitionWriter { writer: Some(spark_operator::partition_writer::Writer::Local( spark_operator::LocalPartitionWriter { output_data_file: output_data_file.to_string(), - output_index_file: output_index_file.to_string(), }, )), } @@ -5180,15 +5160,15 @@ mod tests { fn assert_local_shuffle_destination( writer: &spark_operator::ShuffleWriter, expected_data_file: &str, - expected_index_file: &str, ) { match super::shuffle_writer_destination(writer, None).unwrap() { ShuffleWriterDestination::Local { output_data_file, - output_index_file, + partition_offsets, } => { assert_eq!(output_data_file, expected_data_file); - assert_eq!(output_index_file, expected_index_file); + // A fresh destination has not run a writer yet, so nothing is published. + assert!(partition_offsets.get().is_none()); } destination => panic!("expected a local shuffle destination, got {destination:?}"), } @@ -5198,49 +5178,38 @@ mod tests { fn shuffle_partition_writer_legacy_paths_remain_supported() { let writer = spark_operator::ShuffleWriter { output_data_file: "legacy.data".to_string(), - output_index_file: "legacy.index".to_string(), ..Default::default() }; - assert_local_shuffle_destination(&writer, "legacy.data", "legacy.index"); + assert_local_shuffle_destination(&writer, "legacy.data"); } #[test] fn shuffle_partition_writer_uses_nested_local_paths() { let writer = spark_operator::ShuffleWriter { - partition_writer: Some(local_shuffle_partition_writer( - "shuffle.data", - "shuffle.index", - )), + partition_writer: Some(local_shuffle_partition_writer("shuffle.data")), ..Default::default() }; - assert_local_shuffle_destination(&writer, "shuffle.data", "shuffle.index"); + assert_local_shuffle_destination(&writer, "shuffle.data"); } #[test] fn shuffle_partition_writer_accepts_matching_legacy_paths() { let writer = spark_operator::ShuffleWriter { output_data_file: "shuffle.data".to_string(), - output_index_file: "shuffle.index".to_string(), - partition_writer: Some(local_shuffle_partition_writer( - "shuffle.data", - "shuffle.index", - )), + partition_writer: Some(local_shuffle_partition_writer("shuffle.data")), ..Default::default() }; - assert_local_shuffle_destination(&writer, "shuffle.data", "shuffle.index"); + assert_local_shuffle_destination(&writer, "shuffle.data"); } #[test] fn shuffle_partition_writer_rejects_conflicting_legacy_data_path() { let writer = spark_operator::ShuffleWriter { output_data_file: "legacy.data".to_string(), - partition_writer: Some(local_shuffle_partition_writer( - "shuffle.data", - "shuffle.index", - )), + partition_writer: Some(local_shuffle_partition_writer("shuffle.data")), ..Default::default() }; @@ -5251,29 +5220,11 @@ mod tests { ); } - #[test] - fn shuffle_partition_writer_rejects_conflicting_legacy_index_path() { - let writer = spark_operator::ShuffleWriter { - output_index_file: "legacy.index".to_string(), - partition_writer: Some(local_shuffle_partition_writer( - "shuffle.data", - "shuffle.index", - )), - ..Default::default() - }; - - let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); - assert!( - error.to_string().contains("output index file conflicts"), - "unexpected error: {error}" - ); - } - #[test] fn shuffle_partition_writer_rejects_empty_local_data_path() { let writer = spark_operator::ShuffleWriter { output_data_file: "legacy.data".to_string(), - partition_writer: Some(local_shuffle_partition_writer("", "shuffle.index")), + partition_writer: Some(local_shuffle_partition_writer("")), ..Default::default() }; @@ -5284,21 +5235,6 @@ mod tests { ); } - #[test] - fn shuffle_partition_writer_rejects_empty_local_index_path() { - let writer = spark_operator::ShuffleWriter { - output_index_file: "legacy.index".to_string(), - partition_writer: Some(local_shuffle_partition_writer("shuffle.data", "")), - ..Default::default() - }; - - let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); - assert!( - error.to_string().contains("missing its output index file"), - "unexpected error: {error}" - ); - } - #[test] fn shuffle_partition_writer_rejects_missing_destination() { let writer = spark_operator::ShuffleWriter { @@ -5686,28 +5622,10 @@ mod tests { ); } - #[test] - fn shuffle_partition_writer_rejects_rss_with_legacy_index_path() { - let writer = spark_operator::ShuffleWriter { - output_index_file: "legacy.index".to_string(), - partition_writer: Some(rss_shuffle_partition_writer()), - ..Default::default() - }; - let callback: Arc = - Arc::new(RecordingShufflePartitionPusher::default()); - - let error = super::shuffle_writer_destination(&writer, Some(&callback)).unwrap_err(); - assert!( - error.to_string().contains("cannot have local output files"), - "unexpected error: {error}" - ); - } - #[test] fn shuffle_partition_writer_rejects_callback_for_legacy_local_destination() { let writer = spark_operator::ShuffleWriter { output_data_file: "legacy.data".to_string(), - output_index_file: "legacy.index".to_string(), ..Default::default() }; let callback: Arc = @@ -5725,10 +5643,7 @@ mod tests { #[test] fn shuffle_partition_writer_rejects_callback_for_explicit_local_destination() { let writer = spark_operator::ShuffleWriter { - partition_writer: Some(local_shuffle_partition_writer( - "shuffle.data", - "shuffle.index", - )), + partition_writer: Some(local_shuffle_partition_writer("shuffle.data")), ..Default::default() }; let callback: Arc = diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 75a6e06d998..35e51a811d4 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -806,10 +806,13 @@ message PartitionWriter { } } -// Local shuffle output consists of a data file and its partition-offset index. +// Local shuffle output consists of a data file. The partition offsets are returned to the JVM +// in memory rather than through an index file, so no index path travels in the plan. message LocalPartitionWriter { + reserved 2; + reserved "output_index_file"; + string output_data_file = 1; - string output_index_file = 2; } // Marker for remote shuffle output. The task-owned callback is bound outside @@ -818,10 +821,12 @@ message RssPartitionWriter {} message ShuffleWriter { spark.spark_partitioning.Partitioning partitioning = 1; + reserved 4; + reserved "output_index_file"; + // Retained for compatibility with native binaries that predate partition_writer. - // Local plans also carry these paths in partition_writer.local. + // Local plans also carry this path in partition_writer.local. string output_data_file = 3; - string output_index_file = 4; CompressionCodec codec = 5; int32 compression_level = 6; bool tracing_enabled = 7; diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index c9c088e2802..1e1e68565e6 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -221,7 +221,6 @@ fn create_shuffle_writer_exec( partitioning, compression_codec, "/tmp/data.out".to_string(), - "/tmp/index.out".to_string(), false, 1024 * 1024, None, diff --git a/native/shuffle/src/bin/shuffle_bench.rs b/native/shuffle/src/bin/shuffle_bench.rs index cd43d41dbc1..7630d39dd4e 100644 --- a/native/shuffle/src/bin/shuffle_bench.rs +++ b/native/shuffle/src/bin/shuffle_bench.rs @@ -126,7 +126,6 @@ fn main() { // Create output directory fs::create_dir_all(&args.output_dir).expect("Failed to create output directory"); let data_file = args.output_dir.join("data.out"); - let index_file = args.output_dir.join("index.out"); let (schema, total_rows) = read_parquet_metadata(&args.input, args.limit); @@ -189,7 +188,6 @@ fn main() { &hash_col_indices, &args, data_file.to_str().unwrap(), - index_file.to_str().unwrap(), ) }; let data_size = fs::metadata(&data_file).map(|m| m.len()).unwrap_or(0); @@ -253,7 +251,6 @@ fn main() { } let _ = fs::remove_file(&data_file); - let _ = fs::remove_file(&index_file); } fn print_shuffle_metrics(metrics: &MetricsSet, total_wall_time_secs: f64) { @@ -398,7 +395,6 @@ fn run_shuffle_write( hash_col_indices: &[usize], args: &Args, data_file: &str, - index_file: &str, ) -> (f64, Option, Option) { let partitioning = build_partitioning( &args.partitioning, @@ -420,7 +416,6 @@ fn run_shuffle_write( args.max_buffer_bytes, args.limit, data_file.to_string(), - index_file.to_string(), ) .await .unwrap(); @@ -444,7 +439,6 @@ async fn execute_shuffle_write( max_buffer_bytes: Option, limit: usize, data_file: String, - index_file: String, ) -> datafusion::common::Result<(MetricsSet, MetricsSet)> { let config = SessionConfig::new().with_batch_size(batch_size); let mut runtime_builder = RuntimeEnvBuilder::new(); @@ -483,7 +477,6 @@ async fn execute_shuffle_write( partitioning, codec, data_file, - index_file, false, write_buffer_size, max_buffer_bytes, @@ -537,7 +530,6 @@ fn run_concurrent_shuffle_writes( let task_dir = args.output_dir.join(format!("task_{task_id}")); fs::create_dir_all(&task_dir).expect("Failed to create task output directory"); let data_file = task_dir.join("data.out").to_str().unwrap().to_string(); - let index_file = task_dir.join("index.out").to_str().unwrap().to_string(); let input_str = input_path.to_str().unwrap().to_string(); let codec = codec.clone(); @@ -564,7 +556,6 @@ fn run_concurrent_shuffle_writes( max_buffer_bytes, limit, data_file, - index_file, ) .await .unwrap() diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 766634eb71e..d60d9b6eab9 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -34,5 +34,5 @@ pub use comet_partitioning::CometPartitioning; pub use ipc::{read_ipc_compressed, read_ipc_compressed_validated}; pub use remote_schema::{decode_remote_shuffle_batch, validate_remote_schema}; pub use schema_align::SchemaAlignExec; -pub use shuffle_writer::{ShuffleWriterDestination, ShuffleWriterExec}; +pub use shuffle_writer::{PartitionOffsets, ShuffleWriterDestination, ShuffleWriterExec}; pub use writers::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/rss_execution_tests.rs b/native/shuffle/src/rss_execution_tests.rs index 491562c5732..e9269edc9e5 100644 --- a/native/shuffle/src/rss_execution_tests.rs +++ b/native/shuffle/src/rss_execution_tests.rs @@ -16,8 +16,8 @@ // under the License. use crate::{ - read_ipc_compressed, CometPartitioning, CompressionCodec, ShuffleWriterDestination, - ShuffleWriterExec, + read_ipc_compressed, CometPartitioning, CompressionCodec, PartitionOffsets, + ShuffleWriterDestination, ShuffleWriterExec, }; use arrow::array::{Array, Int32Array, RecordBatch, RecordBatchOptions}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -422,18 +422,18 @@ fn rss_callback_survives_execution_plan_child_replacement() { } #[test] -fn explicit_local_destination_preserves_data_and_index_files() { +fn explicit_local_destination_writes_data_and_publishes_offsets() { let batch = int_batch(0, 8); let directory = tempfile::tempdir().unwrap(); let data_file = directory.path().join("shuffle.data"); - let index_file = directory.path().join("shuffle.index"); + let offsets = Arc::new(PartitionOffsets::default()); let execution = ShuffleWriterExec::try_new_with_destination( memory_input(vec![batch.clone()], batch.schema()), CometPartitioning::SinglePartition, CompressionCodec::None, ShuffleWriterDestination::Local { output_data_file: data_file.to_str().unwrap().to_string(), - output_index_file: index_file.to_str().unwrap().to_string(), + partition_offsets: Arc::clone(&offsets), }, false, 1024 * 1024, @@ -442,7 +442,14 @@ fn explicit_local_destination_preserves_data_and_index_files() { .unwrap(); run_execution(&execution).unwrap(); - let frame = std::fs::read(data_file).unwrap(); + let frame = std::fs::read(&data_file).unwrap(); assert_eq!(decode_frame(&frame).num_rows(), 8); - assert_eq!(std::fs::read(index_file).unwrap().len(), 16); + // A single-partition writer publishes two offsets, the partition start and the total + // length, in memory rather than through an index file. + let published = offsets + .get() + .expect("writer published its partition offsets"); + assert_eq!(published.len(), 2); + assert_eq!(published[0], 0); + assert_eq!(published[1] as usize, frame.len()); } diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 8d668de7725..b203717c884 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -45,18 +45,46 @@ use futures::{StreamExt, TryStreamExt}; use std::{ fmt, fmt::{Debug, Formatter}, - sync::Arc, + sync::{Arc, OnceLock}, }; +/// One-shot slot carrying a local shuffle task's partition offsets from the writer to the JVM. +/// +/// The writer knows every offset by the time it finishes, and the only consumer is the Spark task +/// driving the same plan, so the offsets are handed over in memory instead of being serialized to +/// a temporary index file and read back. `set` is called exactly once, from `finish_all`. +#[derive(Debug, Default)] +pub struct PartitionOffsets(OnceLock>); + +impl PartitionOffsets { + /// Publishes the finished task's offsets. Returns an error if called more than once, which + /// would mean two writers shared one slot and the JVM could read either one's offsets. + pub fn set(&self, offsets: Vec) -> Result<()> { + self.0.set(offsets).map_err(|_| { + DataFusionError::Execution( + "shuffle write error: partition offsets were already published".to_string(), + ) + }) + } + + /// The finished task's offsets, or `None` if the writer has not completed. + pub fn get(&self) -> Option<&[i64]> { + self.0.get().map(Vec::as_slice) + } +} + /// Storage destination for a native shuffle writer. #[derive(Clone)] pub enum ShuffleWriterDestination { - /// Writes partition data and offsets to local shuffle files. + /// Writes partition data to a local shuffle file and publishes the partition offsets in + /// memory. Local { /// Path of the local shuffle data file. output_data_file: String, - /// Path of the local shuffle index file. - output_index_file: String, + /// Receives the partition offsets once the writer finishes, so the JVM can read them + /// over JNI rather than through a temporary index file. Shared with the plan that owns + /// this destination, and set exactly once per task. + partition_offsets: Arc, }, /// Pushes complete encoded partition blocks to a task-owned callback. Rss { @@ -72,11 +100,11 @@ impl Debug for ShuffleWriterDestination { match self { Self::Local { output_data_file, - output_index_file, + partition_offsets, } => f .debug_struct("Local") .field("output_data_file", output_data_file) - .field("output_index_file", output_index_file) + .field("partition_offsets", &partition_offsets.get().is_some()) .finish(), Self::Rss { max_frame_size, .. } => f .debug_struct("Rss") @@ -110,14 +138,14 @@ pub struct ShuffleWriterExec { } impl ShuffleWriterExec { - /// Creates a shuffle writer that writes to local data and index files. + /// Creates a shuffle writer that writes to a local data file, publishing its partition + /// offsets through a fresh [`PartitionOffsets`] slot readable via [`Self::partition_offsets`]. #[allow(clippy::too_many_arguments)] pub fn try_new( input: Arc, partitioning: CometPartitioning, codec: CompressionCodec, output_data_file: String, - output_index_file: String, tracing_enabled: bool, write_buffer_size: usize, max_buffer_bytes: Option, @@ -128,7 +156,7 @@ impl ShuffleWriterExec { codec, ShuffleWriterDestination::Local { output_data_file, - output_index_file, + partition_offsets: Arc::new(PartitionOffsets::default()), }, tracing_enabled, write_buffer_size, @@ -136,6 +164,17 @@ impl ShuffleWriterExec { ) } + /// The slot carrying this task's partition offsets, for a local destination. `None` for a + /// remote destination, where the pusher reports partition lengths instead. + pub fn partition_offsets(&self) -> Option<&Arc> { + match &self.destination { + ShuffleWriterDestination::Local { + partition_offsets, .. + } => Some(partition_offsets), + ShuffleWriterDestination::Rss { .. } => None, + } + } + /// Creates a shuffle writer for a local or task-owned remote destination. pub fn try_new_with_destination( input: Arc, @@ -286,12 +325,12 @@ async fn external_shuffle( let mut repartitioner = match destination { ShuffleWriterDestination::Local { output_data_file, - output_index_file, + partition_offsets, } => { let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?; let writer = LocalPartitionWriter::try_new( output_data_file, - output_index_file, + partition_offsets, shuffle_block_writer, partitioning.partition_count(), context.session_config().batch_size(), @@ -550,7 +589,7 @@ mod test { .unwrap(); let local_partition_writer = LocalPartitionWriter::try_new( "/tmp/data.out".to_string(), - "/tmp/index.out".to_string(), + Arc::new(PartitionOffsets::default()), shuffle_block_writer, num_partitions, 1024, @@ -622,7 +661,7 @@ mod test { ShuffleBlockWriter::try_new(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(), + Arc::new(PartitionOffsets::default()), shuffle_block_writer, num_partitions, 1024, @@ -684,7 +723,7 @@ mod test { ShuffleBlockWriter::try_new(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(), + Arc::new(PartitionOffsets::default()), shuffle_block_writer, num_partitions, batch_size, @@ -827,7 +866,7 @@ mod test { .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(), + Arc::new(PartitionOffsets::default()), shuffle_block_writer, num_partitions, 1024, @@ -935,7 +974,6 @@ mod test { let batch = create_batch(1000); let batches = (0..20).map(|_| batch.clone()).collect::>(); 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( @@ -945,7 +983,6 @@ mod test { 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, @@ -1070,7 +1107,6 @@ mod test { partitioning, CompressionCodec::Zstd(1), "/tmp/data.out".to_string(), - "/tmp/index.out".to_string(), false, 1024 * 1024, // write_buffer_size: 1MB default None, @@ -1118,9 +1154,9 @@ mod test { let batches = (0..num_batches).map(|_| batch.clone()).collect::>(); // Run shuffle twice and compare results + let mut offsets_per_run: Vec> = Vec::new(); for run in 0..2 { let data_file = format!("/tmp/rr_data_{}.out", run); - let index_file = format!("/tmp/rr_index_{}.out", run); let partitions = std::slice::from_ref(&batches); let exec = ShuffleWriterExec::try_new( @@ -1130,7 +1166,6 @@ mod test { CometPartitioning::RoundRobin(num_partitions, 0), CompressionCodec::Zstd(1), data_file.clone(), - index_file.clone(), false, 1024 * 1024, None, @@ -1153,6 +1188,14 @@ mod test { while stream.next().await.is_some() {} }); + offsets_per_run.push( + exec.partition_offsets() + .expect("local destination publishes offsets") + .get() + .expect("writer published its partition offsets") + .to_vec(), + ); + if run == 1 { // Compare data files let mut data0 = Vec::new(); @@ -1170,20 +1213,10 @@ mod test { "Round robin shuffle data should be identical across runs" ); - // Compare index files - let mut index0 = Vec::new(); - fs::File::open("/tmp/rr_index_0.out") - .unwrap() - .read_to_end(&mut index0) - .unwrap(); - let mut index1 = Vec::new(); - fs::File::open("/tmp/rr_index_1.out") - .unwrap() - .read_to_end(&mut index1) - .unwrap(); + // Compare the published partition offsets assert_eq!( - index0, index1, - "Round robin shuffle index should be identical across runs" + offsets_per_run[0], offsets_per_run[1], + "Round robin shuffle partition offsets should be identical across runs" ); } } @@ -1431,13 +1464,11 @@ mod test { let dir = tempfile::tempdir().unwrap(); let data_file = dir.path().join("data.out").to_str().unwrap().to_string(); - let index_file = dir.path().join("index.out").to_str().unwrap().to_string(); - let block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::Lz4Frame).unwrap(); let writer = LocalPartitionWriter::try_new( data_file.clone(), - index_file, + Arc::new(PartitionOffsets::default()), block_writer, 1, // single partition batch_size, @@ -1511,7 +1542,6 @@ mod test { let dir = tempfile::tempdir().unwrap(); let data_file = dir.path().join("data.out"); - let index_file = dir.path().join("index.out"); let exec = ShuffleWriterExec::try_new( Arc::new(DataSourceExec::new(Arc::new( @@ -1520,7 +1550,6 @@ mod test { CometPartitioning::RoundRobin(num_partitions, 0), CompressionCodec::Zstd(1), data_file.to_str().unwrap().to_string(), - index_file.to_str().unwrap().to_string(), false, 1024 * 1024, None, @@ -1551,29 +1580,26 @@ mod test { "Row count should survive roundtrip" ); - // Verify index file structure: num_partitions + 1 offsets - let mut index_data = Vec::new(); - fs::File::open(&index_file) - .unwrap() - .read_to_end(&mut index_data) - .unwrap(); - let expected_index_size = (num_partitions + 1) * 8; - assert_eq!(index_data.len(), expected_index_size); + // Verify the published offsets: num_partitions + 1 of them + let offsets = exec + .partition_offsets() + .expect("local destination publishes offsets") + .get() + .expect("writer published its partition offsets") + .to_vec(); + assert_eq!(offsets.len(), num_partitions + 1); // First offset should be 0 - let first_offset = i64::from_le_bytes(index_data[0..8].try_into().unwrap()); - assert_eq!(first_offset, 0); + assert_eq!(offsets[0], 0); // Second offset should equal data file length (partition 0 holds all data) let data_len = data.len() as i64; - let second_offset = i64::from_le_bytes(index_data[8..16].try_into().unwrap()); - assert_eq!(second_offset, data_len); + assert_eq!(offsets[1], data_len); // All remaining offsets should equal data file length (empty partitions) - for i in 2..=num_partitions { - let offset = i64::from_le_bytes(index_data[i * 8..(i + 1) * 8].try_into().unwrap()); + for (i, offset) in offsets.iter().enumerate().skip(2) { assert_eq!( - offset, data_len, + *offset, data_len, "Partition {i} offset should equal data length" ); } @@ -1600,7 +1626,6 @@ mod test { let dir = tempfile::tempdir().unwrap(); let data_file = dir.path().join("data.out"); - let index_file = dir.path().join("index.out"); let exec = ShuffleWriterExec::try_new( Arc::new(DataSourceExec::new(Arc::new( @@ -1609,7 +1634,6 @@ mod test { CometPartitioning::RoundRobin(num_partitions, 0), CompressionCodec::Zstd(1), data_file.to_str().unwrap().to_string(), - index_file.to_str().unwrap().to_string(), false, 1024 * 1024, None, @@ -1632,17 +1656,16 @@ mod test { .unwrap(); assert!(data.is_empty(), "Data file should be empty with zero rows"); - // Index file should have all-zero offsets - let mut index_data = Vec::new(); - fs::File::open(&index_file) - .unwrap() - .read_to_end(&mut index_data) - .unwrap(); - let expected_index_size = (num_partitions + 1) * 8; - assert_eq!(index_data.len(), expected_index_size); - for i in 0..=num_partitions { - let offset = i64::from_le_bytes(index_data[i * 8..(i + 1) * 8].try_into().unwrap()); - assert_eq!(offset, 0, "All offsets should be 0 with zero rows"); + // Published offsets should be all zero + let offsets = exec + .partition_offsets() + .expect("local destination publishes offsets") + .get() + .expect("writer published its partition offsets") + .to_vec(); + assert_eq!(offsets.len(), num_partitions + 1); + for offset in &offsets { + assert_eq!(*offset, 0, "All offsets should be 0 with zero rows"); } } } diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index e22e339f949..403aa811ab6 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -19,7 +19,7 @@ use crate::metrics::ShufflePartitionerMetrics; use crate::writers::local::spill::SpillWriter; use crate::writers::partition_writer::PartitionWriter; use crate::writers::BufBatchWriter; -use crate::ShuffleBlockWriter; +use crate::{PartitionOffsets, ShuffleBlockWriter}; use arrow::array::RecordBatch; use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; @@ -73,7 +73,7 @@ enum DataOutput { /// byte offset where each partition begins. See [`DataOutput`] for how the /// single- and multi-partition modes differ. pub(crate) struct LocalPartitionWriter { - output_index_file: String, + partition_offsets: Arc, data_output: DataOutput, /// Start offset of each partition in the data file, plus a trailing entry /// with the total length so partition sizes are simple offset differences. @@ -90,7 +90,7 @@ pub(crate) struct LocalPartitionWriter { impl LocalPartitionWriter { pub(crate) fn try_new( output_data_file: String, - output_index_file: String, + partition_offsets: Arc, shuffle_block_writer: ShuffleBlockWriter, num_output_partitions: usize, batch_size: usize, @@ -134,7 +134,7 @@ impl LocalPartitionWriter { } }; Ok(Self { - output_index_file, + partition_offsets, data_output, offsets: vec![0u64; num_output_partitions + 1], batch_size, @@ -306,22 +306,21 @@ impl PartitionWriter for LocalPartitionWriter { // add one extra offset at last to ease partition length computation self.offsets[self.num_output_partitions] = final_offset; - let mut write_timer = metrics.write_time.timer(); - let mut output_index = BufWriter::new( - File::create(self.output_index_file.clone()) - .map_err(|e| DataFusionError::Execution(format!("shuffle write error: {e:?}")))?, - ); - - for offset in &self.offsets { - let offset_i64 = i64::try_from(*offset).map_err(|_| { - DataFusionError::Execution(format!( - "shuffle write error: offset overflow ({offset})" - )) - })?; - output_index.write_all(&offset_i64.to_le_bytes())?; - } - output_index.flush()?; - write_timer.stop(); + // The offsets go straight to the Spark task driving this plan, which reads them over + // JNI. Writing them to a temporary index file first would cost every map task a create, + // write, read and unlink on top of the index file Spark itself commits. + let offsets = self + .offsets + .iter() + .map(|offset| { + i64::try_from(*offset).map_err(|_| { + DataFusionError::Execution(format!( + "shuffle write error: offset overflow ({offset})" + )) + }) + }) + .collect::, _>>()?; + self.partition_offsets.set(offsets)?; Ok(()) } @@ -355,7 +354,7 @@ mod tests { .unwrap(); LocalPartitionWriter::try_new( dir.path().join("data.out").to_str().unwrap().to_string(), - dir.path().join("index.out").to_str().unwrap().to_string(), + Arc::new(PartitionOffsets::default()), block_writer, 2, // batch_size below the row count so the write serializes into the scratch. diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e2c132904d5..066a35682cc 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -77,7 +77,8 @@ class CometExecIterator( encryptedFilePaths: Seq[String] = Seq.empty, shuffleBlockIterators: Map[Int, CometShuffleBlockIterator] = Map.empty, taskFilePaths: Seq[String] = Seq.empty, - shufflePartitionPusher: Option[ShufflePartitionPusher] = None) + shufflePartitionPusher: Option[ShufflePartitionPusher] = None, + capturePartitionOffsets: Boolean = false) extends Iterator[ColumnarBatch] with Logging { @@ -179,6 +180,21 @@ class CometExecIterator( } } + /** + * Partition offsets published by a native shuffle write, captured when the plan reached the end + * of its output and before the native execution context was released. + * + * Only populated when the iterator was built with `capturePartitionOffsets`, and only after the + * iterator has been drained. `null` otherwise. + */ + private var capturedPartitionOffsets: Array[Long] = _ + + /** + * The partition offsets captured at end of stream, or `null` if the plan has not been drained + * or the iterator was not built to capture them. + */ + def shufflePartitionOffsets: Array[Long] = capturedPartitionOffsets + private var nextBatch: Option[ColumnarBatch] = None private var prevBatch: ColumnarBatch = null private var currentBatch: ColumnarBatch = null @@ -248,6 +264,11 @@ class CometExecIterator( logTrace(s"Task $taskAttemptId memory pool usage is ${cometTaskMemoryManager.getUsed} bytes") if (nextBatch.isEmpty) { + // The offsets live in the native execution context, which `close` releases, so they have + // to be read here while the plan is still alive. + if (capturePartitionOffsets && capturedPartitionOffsets == null) { + capturedPartitionOffsets = nativeLib.getShufflePartitionOffsets(plan) + } close() false } else { diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index af406632f15..c3dccca162c 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -116,6 +116,21 @@ class Native extends NativeBase { arrayAddrs: Array[Long], schemaAddrs: Array[Long]): Long + /** + * Returns the partition offsets published by a finished native shuffle write. + * + * The writer knows every offset once its plan completes, so they are handed back in memory + * rather than through a temporary index file. Call only after the plan has been fully drained, + * and only for a plan whose root is a native shuffle writer with a local destination. + * + * @param plan + * the address to native query plan. + * @return + * `numPartitions + 1` offsets, the last being the total data file length, so that partition + * lengths are successive differences. + */ + @native def getShufflePartitionOffsets(plan: Long): Array[Long] + /** * Release and drop the native query plan object and context object. * 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 21dd7686302..a1faadb2c38 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 @@ -19,7 +19,6 @@ package org.apache.spark.sql.comet.execution.shuffle -import java.nio.{ByteBuffer, ByteOrder} import java.nio.file.{Files, Paths} import java.util.concurrent.{ScheduledFuture, TimeUnit} import java.util.concurrent.atomic.AtomicBoolean @@ -69,8 +68,6 @@ class CometNativeShuffleWriter[K, V]( extends ShuffleWriter[K, V] with Logging { - private val OFFSET_LENGTH = 8 - var partitionLengths: Array[Long] = _ var mapStatus: MapStatus = _ private var stopped = false @@ -116,12 +113,7 @@ class CometNativeShuffleWriter[K, V]( val resolver = SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver] val dataFile = resolver.getDataFile(shuffleId, mapId) - val indexFile = resolver.getIndexFile(shuffleId, mapId) - Some( - LocalShuffleOutput( - resolver, - dataFile.getPath.replace(".data", ".data.tmp"), - indexFile.getPath.replace(".index", ".index.tmp"))) + Some(LocalShuffleOutput(resolver, dataFile.getPath.replace(".data", ".data.tmp"))) } else { None } @@ -142,8 +134,8 @@ class CometNativeShuffleWriter[K, V]( val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators val unifiedPlan = localOutput match { - case Some(output) => buildUnifiedPlan(output.dataFile, output.indexFile) - case None => buildUnifiedPlan("", "") + case Some(output) => buildUnifiedPlan(output.dataFile) + case None => buildUnifiedPlan("") } val ctx = spec.execContext val finalNativePlan = if (ctx.commonByKey.nonEmpty) { @@ -198,7 +190,10 @@ class CometNativeShuffleWriter[K, V]( ctx.broadcastedHadoopConfForEncryption, ctx.encryptedFilePaths, shuffleBlockIters, - shufflePartitionPusher = remoteDestination.map(_.callback)) + shufflePartitionPusher = remoteDestination.map(_.callback), + // Only a local destination publishes partition offsets; RSS reports lengths through its + // pusher instead. + capturePartitionOffsets = localOutput.isDefined) // Register subqueries against the iterator id so native callbacks resolve them to values. ctx.subqueries.foreach { sub => @@ -213,6 +208,8 @@ class CometNativeShuffleWriter[K, V]( } CometNativeShuffleWriter.drainAndClose(cometIter, () => cometIter.close()) + // Captured by the iterator at end of stream, before it released the native plan that owns it. + val partitionOffsets = cometIter.shufflePartitionOffsets remoteDestination match { case Some(destination) => @@ -243,22 +240,22 @@ class CometNativeShuffleWriter[K, V]( case None => val output = localOutput.get val tempDataFilePath = Paths.get(output.dataFile) - val tempIndexFilePath = Paths.get(output.indexFile) - - var offset = 0L - partitionLengths = Files - .readAllBytes(tempIndexFilePath) - .grouped(OFFSET_LENGTH) - .drop(1) - .map(indexBytes => { - val partitionOffset = - ByteBuffer.wrap(indexBytes).order(ByteOrder.LITTLE_ENDIAN).getLong - val partitionLength = partitionOffset - offset - offset = partitionOffset - partitionLength - }) - .toArray - Files.delete(tempIndexFilePath) + + // One offset per output partition plus a trailing total, so lengths are successive + // differences. `numParts` is the input partition count and is not this. + val numOutputPartitions = effectivePartitionCount + require( + partitionOffsets != null && partitionOffsets.length == numOutputPartitions + 1, + s"Native shuffle returned ${if (partitionOffsets == null) "no" + else partitionOffsets.length.toString} partition offsets " + + s"for $numOutputPartitions output partitions") + partitionLengths = new Array[Long](numOutputPartitions) + var partition = 0 + while (partition < numOutputPartitions) { + partitionLengths(partition) = + partitionOffsets(partition + 1) - partitionOffsets(partition) + partition += 1 + } metricsReporter.incBytesWritten(Files.size(tempDataFilePath)) output.resolver.writeMetadataFileAndCommit( @@ -290,7 +287,7 @@ class CometNativeShuffleWriter[K, V]( * Build the unified `ShuffleWriter(child = childNativeOp)` plan with the partitioning serde, * compression settings, and output file paths. */ - private[shuffle] def buildUnifiedPlan(dataFile: String, indexFile: String): Operator = { + private[shuffle] def buildUnifiedPlan(dataFile: String): Operator = { val shuffleWriterBuilder = OperatorOuterClass.ShuffleWriter.newBuilder() remoteDestination match { case Some(_) => @@ -300,9 +297,9 @@ class CometNativeShuffleWriter[K, V]( .setRss(OperatorOuterClass.RssPartitionWriter.getDefaultInstance) .build()) case None => - // Keep legacy paths for older native libraries while newer libraries use the destination. + // Keep the legacy path for older native libraries while newer libraries use the + // destination. Partition offsets come back over JNI, so no index path travels in the plan. shuffleWriterBuilder.setOutputDataFile(dataFile) - shuffleWriterBuilder.setOutputIndexFile(indexFile) shuffleWriterBuilder.setPartitionWriter( OperatorOuterClass.PartitionWriter .newBuilder() @@ -310,7 +307,6 @@ class CometNativeShuffleWriter[K, V]( OperatorOuterClass.LocalPartitionWriter .newBuilder() .setOutputDataFile(dataFile) - .setOutputIndexFile(indexFile) .build()) .build()) } @@ -482,8 +478,7 @@ class CometNativeShuffleWriter[K, V]( private final case class LocalShuffleOutput( resolver: IndexShuffleBlockResolver, - dataFile: String, - indexFile: String) + dataFile: String) } private[shuffle] object CometNativeShuffleWriter { 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 d5e768f886e..a65d86fd203 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -344,18 +344,15 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper assert(results.sameElements(Array((true, true, true, true)))) } - test("native shuffle plan preserves local partition writer and legacy output paths") { + test("native shuffle plan preserves local partition writer and legacy output path") { val dataFile = "/tmp/comet-shuffle.data" - val indexFile = "/tmp/comet-shuffle.index" val localWriter = OperatorOuterClass.LocalPartitionWriter .newBuilder() .setOutputDataFile(dataFile) - .setOutputIndexFile(indexFile) .build() val writer = OperatorOuterClass.ShuffleWriter .newBuilder() .setOutputDataFile(dataFile) - .setOutputIndexFile(indexFile) .setPartitionWriter( OperatorOuterClass.PartitionWriter.newBuilder().setLocal(localWriter).build()) .build() @@ -366,16 +363,13 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper assert(decoded.getPartitionWriter.hasLocal) assert(!decoded.getPartitionWriter.hasRss) assert(decoded.getPartitionWriter.getLocal.getOutputDataFile == dataFile) - assert(decoded.getPartitionWriter.getLocal.getOutputIndexFile == indexFile) assert(decoded.getOutputDataFile == dataFile) - assert(decoded.getOutputIndexFile == indexFile) } test("native shuffle plan preserves RSS partition writer and excludes local destination") { val localWriter = OperatorOuterClass.LocalPartitionWriter .newBuilder() .setOutputDataFile("/tmp/comet-shuffle.data") - .setOutputIndexFile("/tmp/comet-shuffle.index") .build() val partitionWriter = OperatorOuterClass.PartitionWriter .newBuilder() @@ -393,23 +387,19 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper assert(decoded.getPartitionWriter.hasRss) assert(!decoded.getPartitionWriter.hasLocal) assert(decoded.getOutputDataFile.isEmpty) - assert(decoded.getOutputIndexFile.isEmpty) } test("legacy native shuffle plans remain valid without a partition writer") { val dataFile = "/tmp/legacy-shuffle.data" - val indexFile = "/tmp/legacy-shuffle.index" val writer = OperatorOuterClass.ShuffleWriter .newBuilder() .setOutputDataFile(dataFile) - .setOutputIndexFile(indexFile) .build() val decoded = OperatorOuterClass.ShuffleWriter.parseFrom(writer.toByteArray) assert(!decoded.hasPartitionWriter) assert(decoded.getOutputDataFile == dataFile) - assert(decoded.getOutputIndexFile == indexFile) } // TODO: this test takes a long time to run, we should reduce the test time. diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala index a300bf8de9f..048b4c09a6e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala @@ -209,10 +209,9 @@ class CometCelebornNativeShuffleWriterSuite extends CometTestBase { validations += 1 true }) - val plan = writer.buildUnifiedPlan("", "").getShuffleWriter + val plan = writer.buildUnifiedPlan("").getShuffleWriter assert(plan.getPartitionWriter.hasRss) assert(plan.getOutputDataFile.isEmpty) - assert(plan.getOutputIndexFile.isEmpty) writer.write(inputs) val status = writer.stop(success = true).get From 6e1fb10d54e0f9fc82e18154a8e057c5083d8b2e Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 10 Sep 2026 02:17:00 +0800 Subject: [PATCH 2/8] review: address feedback on naming and layering - The shuffle crate does not know about JNI, so PartitionOffsets and the local writer describe the handover in terms of the caller driving the plan rather than the JVM reading over JNI. - ShuffleWriterExec::try_new says what it writes where: partition data to a local file, offsets in memory. - CometExecIterator had three lookalike names for one thing. The read is now a named method, readPartitionOffsetsBeforeClose, whose name and doc carry the constraint that made the placement surprising: the offsets live in the native execution context, close releases it, and hasNext closes as soon as the plan runs out of output, so the final hasNext is the last point they can be read. The field is partitionOffsets and the constructor flag is documented. Co-Authored-By: Claude Opus 5 --- native/shuffle/src/shuffle_writer.rs | 19 +++++----- .../writers/local/local_partition_writer.rs | 6 +-- .../org/apache/comet/CometExecIterator.scala | 37 ++++++++++++------- 3 files changed, 36 insertions(+), 26 deletions(-) diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index b203717c884..98ffc2848fb 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -48,17 +48,17 @@ use std::{ sync::{Arc, OnceLock}, }; -/// One-shot slot carrying a local shuffle task's partition offsets from the writer to the JVM. +/// One-shot slot carrying a local shuffle task's partition offsets out of the writer. /// -/// The writer knows every offset by the time it finishes, and the only consumer is the Spark task -/// driving the same plan, so the offsets are handed over in memory instead of being serialized to -/// a temporary index file and read back. `set` is called exactly once, from `finish_all`. +/// The writer knows every offset by the time it finishes, and the only consumer is whoever is +/// driving the plan, so the offsets are handed over in memory instead of being serialized to a +/// temporary index file and read back. `set` is called exactly once, from `finish_all`. #[derive(Debug, Default)] pub struct PartitionOffsets(OnceLock>); impl PartitionOffsets { /// Publishes the finished task's offsets. Returns an error if called more than once, which - /// would mean two writers shared one slot and the JVM could read either one's offsets. + /// would mean two writers shared one slot and a reader could observe either one's offsets. pub fn set(&self, offsets: Vec) -> Result<()> { self.0.set(offsets).map_err(|_| { DataFusionError::Execution( @@ -81,8 +81,8 @@ pub enum ShuffleWriterDestination { Local { /// Path of the local shuffle data file. output_data_file: String, - /// Receives the partition offsets once the writer finishes, so the JVM can read them - /// over JNI rather than through a temporary index file. Shared with the plan that owns + /// Receives the partition offsets once the writer finishes, so the caller can read them + /// in memory rather than through a temporary index file. Shared with the plan that owns /// this destination, and set exactly once per task. partition_offsets: Arc, }, @@ -138,8 +138,9 @@ pub struct ShuffleWriterExec { } impl ShuffleWriterExec { - /// Creates a shuffle writer that writes to a local data file, publishing its partition - /// offsets through a fresh [`PartitionOffsets`] slot readable via [`Self::partition_offsets`]. + /// Creates a shuffle writer that writes partition data to a local file and publishes its + /// partition offsets in memory, through a fresh [`PartitionOffsets`] slot that + /// [`Self::partition_offsets`] hands back to the caller. #[allow(clippy::too_many_arguments)] pub fn try_new( input: Arc, diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 403aa811ab6..2a8f791762c 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -306,9 +306,9 @@ impl PartitionWriter for LocalPartitionWriter { // add one extra offset at last to ease partition length computation self.offsets[self.num_output_partitions] = final_offset; - // The offsets go straight to the Spark task driving this plan, which reads them over - // JNI. Writing them to a temporary index file first would cost every map task a create, - // write, read and unlink on top of the index file Spark itself commits. + // The offsets go straight to whoever is driving this plan. Writing them to a temporary + // index file first would cost every task a create, write, read and unlink on top of the + // index the caller ultimately commits. let offsets = self .offsets .iter() diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 066a35682cc..4da95af18d0 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -64,6 +64,10 @@ import org.apache.comet.vector.NativeUtil * Paths to encrypted Parquet files that need key unwrapping. * @param shufflePartitionPusher * Optional task-owned callback that receives remote shuffle output. + * @param capturePartitionOffsets + * Whether to read the shuffle writer's partition offsets when the plan reaches the end of its + * output, for a plan rooted at a native shuffle writer with a local destination. Remote shuffle + * reports its partition lengths through its pusher instead, so it leaves this false. */ class CometExecIterator( val id: Long, @@ -180,20 +184,29 @@ class CometExecIterator( } } + /** Set once by [[readPartitionOffsetsBeforeClose]]; `null` until then. */ + private var partitionOffsets: Array[Long] = _ + /** - * Partition offsets published by a native shuffle write, captured when the plan reached the end - * of its output and before the native execution context was released. - * - * Only populated when the iterator was built with `capturePartitionOffsets`, and only after the - * iterator has been drained. `null` otherwise. + * Partition offsets from a native shuffle write, or `null` if this iterator was not built to + * collect them or has not yet reached the end of its output. */ - private var capturedPartitionOffsets: Array[Long] = _ + def shufflePartitionOffsets: Array[Long] = partitionOffsets /** - * The partition offsets captured at end of stream, or `null` if the plan has not been drained - * or the iterator was not built to capture them. + * Reads the shuffle writer's partition offsets out of the native plan, if this iterator was + * built to collect them. + * + * This has to run at end of stream rather than after iteration finishes. The offsets live in + * the native execution context; [[close]] releases that context, and [[hasNext]] calls + * [[close]] as soon as the plan runs out of output. So the final [[hasNext]] is the last point + * at which they can still be read. */ - def shufflePartitionOffsets: Array[Long] = capturedPartitionOffsets + private def readPartitionOffsetsBeforeClose(): Unit = { + if (capturePartitionOffsets && partitionOffsets == null) { + partitionOffsets = nativeLib.getShufflePartitionOffsets(plan) + } + } private var nextBatch: Option[ColumnarBatch] = None private var prevBatch: ColumnarBatch = null @@ -264,11 +277,7 @@ class CometExecIterator( logTrace(s"Task $taskAttemptId memory pool usage is ${cometTaskMemoryManager.getUsed} bytes") if (nextBatch.isEmpty) { - // The offsets live in the native execution context, which `close` releases, so they have - // to be read here while the plan is still alive. - if (capturePartitionOffsets && capturedPartitionOffsets == null) { - capturedPartitionOffsets = nativeLib.getShufflePartitionOffsets(plan) - } + readPartitionOffsetsBeforeClose() close() false } else { From f83eca44624b3699f64c8df682c09863102c218a Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 10 Sep 2026 02:43:58 +0800 Subject: [PATCH 3/8] fix: size shuffle partition lengths from the offsets the writer returned Two CI failures, both from assuming more than the writer guarantees. The proto crate's own tests still referenced ShuffleWriter.output_index_file and LocalPartitionWriter.output_index_file, so datafusion-comet-proto failed to compile its test target. I had only checked the shuffle and core crates locally rather than the whole workspace. The round-trip tests now assert that a new plan carries no index path, and that a plan still carrying the retired tag 4 decodes cleanly because the tag is reserved rather than reused. partitionLengths was sized by effectivePartitionCount, which is not what the writer produces. isSinglePartitioning serializes a range partitioning whose sampled bounds came out empty as SinglePartition, so native writes one partition while the declared output partitioning still reports several, and the require failed with "returned 2 partition offsets for 10 output partitions". The index file was always sized by what the writer produced, so deriving the length count from the returned offsets restores the previous behaviour exactly. Co-Authored-By: Claude Opus 5 --- native/proto/src/lib.rs | 13 ++++++------ .../shuffle/CometNativeShuffleWriter.scala | 20 ++++++++++--------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/native/proto/src/lib.rs b/native/proto/src/lib.rs index c814760c4b1..5f79d35b06e 100644 --- a/native/proto/src/lib.rs +++ b/native/proto/src/lib.rs @@ -68,15 +68,12 @@ mod tests { fn local_shuffle_writer() -> ShuffleWriter { let output_data_file = "/tmp/shuffle.data".to_string(); - let output_index_file = "/tmp/shuffle.index".to_string(); ShuffleWriter { output_data_file: output_data_file.clone(), - output_index_file: output_index_file.clone(), partition_writer: Some(PartitionWriter { writer: Some(partition_writer::Writer::Local(LocalPartitionWriter { output_data_file, - output_index_file, })), }), ..Default::default() @@ -89,14 +86,12 @@ mod tests { let decoded = ShuffleWriter::decode(encoded.as_slice()).unwrap(); assert_eq!(decoded.output_data_file, "/tmp/shuffle.data"); - assert_eq!(decoded.output_index_file, "/tmp/shuffle.index"); let Some(partition_writer::Writer::Local(local)) = decoded.partition_writer.and_then(|writer| writer.writer) else { panic!("expected a local shuffle partition writer"); }; assert_eq!(local.output_data_file, "/tmp/shuffle.data"); - assert_eq!(local.output_index_file, "/tmp/shuffle.index"); } #[test] @@ -121,9 +116,14 @@ mod tests { let decoded = LegacyShuffleWriter::decode(encoded.as_slice()).unwrap(); assert_eq!(decoded.output_data_file, "/tmp/shuffle.data"); - assert_eq!(decoded.output_index_file, "/tmp/shuffle.index"); + // Partition offsets are returned in memory now, so a new plan carries no index path and + // a reader still expecting tag 4 simply sees it unset. + assert!(decoded.output_index_file.is_empty()); } + /// A plan still carrying the retired index path decodes cleanly: tag 4 is reserved rather + /// than reused, so it is skipped as an unknown field instead of being misread as something + /// else. #[test] fn new_shuffle_writer_decodes_legacy_plan_without_destination() { let legacy = LegacyShuffleWriter { @@ -133,7 +133,6 @@ mod tests { let decoded = ShuffleWriter::decode(legacy.encode_to_vec().as_slice()).unwrap(); assert_eq!(decoded.output_data_file, "/tmp/legacy.data"); - assert_eq!(decoded.output_index_file, "/tmp/legacy.index"); assert!(decoded.partition_writer.is_none()); } } 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 a1faadb2c38..817de8f114d 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 @@ -241,17 +241,19 @@ class CometNativeShuffleWriter[K, V]( val output = localOutput.get val tempDataFilePath = Paths.get(output.dataFile) - // One offset per output partition plus a trailing total, so lengths are successive - // differences. `numParts` is the input partition count and is not this. - val numOutputPartitions = effectivePartitionCount + // The writer emits one offset per partition it actually wrote plus a trailing total, so + // lengths are successive differences and their count comes from the offsets themselves. + // It is not `effectivePartitionCount`: `isSinglePartitioning` serializes a range + // partitioning whose sampled bounds are empty as SinglePartition, so native writes one + // partition while the declared output partitioning still reports more. Sizing this from + // the offsets keeps the behaviour the index file gave, which was also sized by what the + // writer produced. require( - partitionOffsets != null && partitionOffsets.length == numOutputPartitions + 1, - s"Native shuffle returned ${if (partitionOffsets == null) "no" - else partitionOffsets.length.toString} partition offsets " + - s"for $numOutputPartitions output partitions") - partitionLengths = new Array[Long](numOutputPartitions) + partitionOffsets != null && partitionOffsets.length >= 1, + "Native shuffle returned no partition offsets") + partitionLengths = new Array[Long](partitionOffsets.length - 1) var partition = 0 - while (partition < numOutputPartitions) { + while (partition < partitionLengths.length) { partitionLengths(partition) = partitionOffsets(partition + 1) - partitionOffsets(partition) partition += 1 From 313fdc29ec1aa28ce3be05cadcfbe5d8f699ab20 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 10 Sep 2026 02:54:57 +0800 Subject: [PATCH 4/8] review: drop the reserved declarations for the retired index path Requested in review. The plan proto is built by the JVM and consumed by native in the same process from the same artifact, so no old plan ever meets a new reader and there is nothing for a reserved tag to protect against. Decoding is unaffected either way: an undeclared tag is skipped as an unknown field, and reserved only stops protoc from later reusing the number. The proto round-trip test covering a plan that still carries the retired tag 4 keeps passing, and its comment no longer credits reserved for that. Co-Authored-By: Claude Opus 5 --- native/proto/src/lib.rs | 5 ++--- native/proto/src/proto/operator.proto | 6 ------ 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/native/proto/src/lib.rs b/native/proto/src/lib.rs index 5f79d35b06e..f9f1f120b11 100644 --- a/native/proto/src/lib.rs +++ b/native/proto/src/lib.rs @@ -121,9 +121,8 @@ mod tests { assert!(decoded.output_index_file.is_empty()); } - /// A plan still carrying the retired index path decodes cleanly: tag 4 is reserved rather - /// than reused, so it is skipped as an unknown field instead of being misread as something - /// else. + /// A plan still carrying the retired index path decodes cleanly, since tag 4 is no longer + /// declared and is skipped as an unknown field. #[test] fn new_shuffle_writer_decodes_legacy_plan_without_destination() { let legacy = LegacyShuffleWriter { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 35e51a811d4..0f3101a5c12 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -809,9 +809,6 @@ message PartitionWriter { // Local shuffle output consists of a data file. The partition offsets are returned to the JVM // in memory rather than through an index file, so no index path travels in the plan. message LocalPartitionWriter { - reserved 2; - reserved "output_index_file"; - string output_data_file = 1; } @@ -821,9 +818,6 @@ message RssPartitionWriter {} message ShuffleWriter { spark.spark_partitioning.Partitioning partitioning = 1; - reserved 4; - reserved "output_index_file"; - // Retained for compatibility with native binaries that predate partition_writer. // Local plans also carry this path in partition_writer.local. string output_data_file = 3; From 7e4c5e4139561f013480a39d263a936f541f5284 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 10 Sep 2026 09:41:31 +0800 Subject: [PATCH 5/8] review: trim comments and close the gap left by the retired field number Comments cut back to what they need to say, per review: the JNI entry point, the PartitionOffsets type and its set/get, the destination field, try_new, the partition_offsets accessor, the zero-offset test, the finish_all note, and the two comments in CometNativeShuffleWriter. ShuffleWriter field numbers 5 through 11 shift down to 4 through 10, closing the gap the retired output_index_file left. Both sides of the plan are generated from this file and ship together, so no encoded plan outlives the change. That does mean tag 4 now belongs to codec, and the LegacyShuffleWriter test struct claimed it for a string. Decoding a plan carrying it would be a wire type mismatch rather than a skipped unknown field, so the struct drops that field. The test still covers a legacy plan decoding without a partition writer. Co-Authored-By: Claude Opus 5 --- native/core/src/execution/jni_api.rs | 7 +----- native/proto/src/lib.rs | 8 ------- native/proto/src/proto/operator.proto | 17 +++++++------- native/shuffle/src/shuffle_writer.rs | 22 ++++++------------- .../writers/local/local_partition_writer.rs | 3 --- .../shuffle/CometNativeShuffleWriter.scala | 9 +------- 6 files changed, 17 insertions(+), 49 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 5ef4fb44b0b..06aac12ac77 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1157,13 +1157,8 @@ fn get_execution_context<'a>(id: i64) -> &'a mut ExecutionContext { /// Returns the partition offsets published by a finished native shuffle write. /// -/// The writer knows every offset by the time its plan completes, and the only consumer is the -/// Spark task driving that plan, so the offsets are handed back in memory rather than serialized -/// to a temporary index file and read back. Call after the plan has been fully drained; the -/// offsets are not published until the writer finishes. -/// /// The returned array holds `num_output_partitions + 1` offsets, the last being the total data -/// file length, so partition lengths are successive differences. +/// file length. #[no_mangle] pub extern "system" fn Java_org_apache_comet_Native_getShufflePartitionOffsets( e: EnvUnowned, diff --git a/native/proto/src/lib.rs b/native/proto/src/lib.rs index f9f1f120b11..dd5ccfbe05b 100644 --- a/native/proto/src/lib.rs +++ b/native/proto/src/lib.rs @@ -62,8 +62,6 @@ mod tests { struct LegacyShuffleWriter { #[prost(string, tag = "3")] output_data_file: String, - #[prost(string, tag = "4")] - output_index_file: String, } fn local_shuffle_writer() -> ShuffleWriter { @@ -116,18 +114,12 @@ mod tests { let decoded = LegacyShuffleWriter::decode(encoded.as_slice()).unwrap(); assert_eq!(decoded.output_data_file, "/tmp/shuffle.data"); - // Partition offsets are returned in memory now, so a new plan carries no index path and - // a reader still expecting tag 4 simply sees it unset. - assert!(decoded.output_index_file.is_empty()); } - /// A plan still carrying the retired index path decodes cleanly, since tag 4 is no longer - /// declared and is skipped as an unknown field. #[test] fn new_shuffle_writer_decodes_legacy_plan_without_destination() { let legacy = LegacyShuffleWriter { output_data_file: "/tmp/legacy.data".to_string(), - output_index_file: "/tmp/legacy.index".to_string(), }; let decoded = ShuffleWriter::decode(legacy.encode_to_vec().as_slice()).unwrap(); diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 0f3101a5c12..66db291b388 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -806,8 +806,7 @@ message PartitionWriter { } } -// Local shuffle output consists of a data file. The partition offsets are returned to the JVM -// in memory rather than through an index file, so no index path travels in the plan. +// Local shuffle output consists of a data file. The partition offsets are returned to the JVM via JNI message LocalPartitionWriter { string output_data_file = 1; } @@ -821,23 +820,23 @@ message ShuffleWriter { // Retained for compatibility with native binaries that predate partition_writer. // Local plans also carry this path in partition_writer.local. string output_data_file = 3; - CompressionCodec codec = 5; - int32 compression_level = 6; - bool tracing_enabled = 7; + CompressionCodec codec = 4; + int32 compression_level = 5; + bool tracing_enabled = 6; // Size of the write buffer in bytes used when writing shuffle data to disk. // Larger values may improve write performance but use more memory. - int32 write_buffer_size = 8; + int32 write_buffer_size = 7; // Spark-declared output schema of the writer's child. When the child is an inlined native // subtree, the native planner casts the child's actual output to this schema before // serializing to shuffle blocks, since there is no FFI boundary or ScanExec between them // 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; + repeated SparkStructField expected_output_schema = 8; // 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; + uint64 max_buffer_bytes = 9; // Explicit output destination. When absent, use the legacy output file fields. - PartitionWriter partition_writer = 11; + PartitionWriter partition_writer = 10; } message ParquetWriter { diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 98ffc2848fb..d6aaf0ae0e9 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -49,16 +49,11 @@ use std::{ }; /// One-shot slot carrying a local shuffle task's partition offsets out of the writer. -/// -/// The writer knows every offset by the time it finishes, and the only consumer is whoever is -/// driving the plan, so the offsets are handed over in memory instead of being serialized to a -/// temporary index file and read back. `set` is called exactly once, from `finish_all`. #[derive(Debug, Default)] pub struct PartitionOffsets(OnceLock>); impl PartitionOffsets { - /// Publishes the finished task's offsets. Returns an error if called more than once, which - /// would mean two writers shared one slot and a reader could observe either one's offsets. + /// Publishes the finished task's offsets. Errors if called more than once. pub fn set(&self, offsets: Vec) -> Result<()> { self.0.set(offsets).map_err(|_| { DataFusionError::Execution( @@ -81,9 +76,7 @@ pub enum ShuffleWriterDestination { Local { /// Path of the local shuffle data file. output_data_file: String, - /// Receives the partition offsets once the writer finishes, so the caller can read them - /// in memory rather than through a temporary index file. Shared with the plan that owns - /// this destination, and set exactly once per task. + /// One offset per partition written, plus a trailing total. partition_offsets: Arc, }, /// Pushes complete encoded partition blocks to a task-owned callback. @@ -138,9 +131,8 @@ pub struct ShuffleWriterExec { } impl ShuffleWriterExec { - /// Creates a shuffle writer that writes partition data to a local file and publishes its - /// partition offsets in memory, through a fresh [`PartitionOffsets`] slot that - /// [`Self::partition_offsets`] hands back to the caller. + /// Creates a shuffle writer that writes partition data to a local file and exposes its + /// partition offsets. #[allow(clippy::too_many_arguments)] pub fn try_new( input: Arc, @@ -165,8 +157,8 @@ impl ShuffleWriterExec { ) } - /// The slot carrying this task's partition offsets, for a local destination. `None` for a - /// remote destination, where the pusher reports partition lengths instead. + /// Returns this task's partition offsets, for a local destination. `None` for a remote + /// destination, where the pusher reports partition lengths instead. pub fn partition_offsets(&self) -> Option<&Arc> { match &self.destination { ShuffleWriterDestination::Local { @@ -1657,7 +1649,7 @@ mod test { .unwrap(); assert!(data.is_empty(), "Data file should be empty with zero rows"); - // Published offsets should be all zero + // partition offsets should be all zero let offsets = exec .partition_offsets() .expect("local destination publishes offsets") diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 2a8f791762c..984911470b9 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -306,9 +306,6 @@ impl PartitionWriter for LocalPartitionWriter { // add one extra offset at last to ease partition length computation self.offsets[self.num_output_partitions] = final_offset; - // The offsets go straight to whoever is driving this plan. Writing them to a temporary - // index file first would cost every task a create, write, read and unlink on top of the - // index the caller ultimately commits. let offsets = self .offsets .iter() 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 817de8f114d..4b10019341f 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 @@ -241,13 +241,6 @@ class CometNativeShuffleWriter[K, V]( val output = localOutput.get val tempDataFilePath = Paths.get(output.dataFile) - // The writer emits one offset per partition it actually wrote plus a trailing total, so - // lengths are successive differences and their count comes from the offsets themselves. - // It is not `effectivePartitionCount`: `isSinglePartitioning` serializes a range - // partitioning whose sampled bounds are empty as SinglePartition, so native writes one - // partition while the declared output partitioning still reports more. Sizing this from - // the offsets keeps the behaviour the index file gave, which was also sized by what the - // writer produced. require( partitionOffsets != null && partitionOffsets.length >= 1, "Native shuffle returned no partition offsets") @@ -300,7 +293,7 @@ class CometNativeShuffleWriter[K, V]( .build()) case None => // Keep the legacy path for older native libraries while newer libraries use the - // destination. Partition offsets come back over JNI, so no index path travels in the plan. + // destination. Partition offsets come back over JNI. shuffleWriterBuilder.setOutputDataFile(dataFile) shuffleWriterBuilder.setPartitionWriter( OperatorOuterClass.PartitionWriter From 377f99bae34df1c072d3f3cb3bf010758435b9cc Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 13 Sep 2026 15:57:52 +0800 Subject: [PATCH 6/8] perf: spill every shuffle partition of a task into one file The native shuffle writer spilled each output partition to its own temporary file, created on that partition's first spill and held open until the task finished. A task with P partitions that spilled held up to P spill files open at once, each with a create and an unlink, and every spill scattered its bytes across P files. A task now spills to a single file. Each write appends the partition's blocks and records the range they occupy, and finish_partition copies a partition's ranges into the output in write order. Correctness does not depend on the order partitions are written in, which PartitionWriter leaves unspecified. A single spill round gives contiguous ascending ranges, so the merge reads sequentially. A write that fails partway can leave uncounted bytes in the shared file, which would shift every later range, so the spill refuses further writes and range reads after a failure. The merge also checks each copy's length, so a spill file shorter than its ranges fails the partition instead of writing it short. Closes #3859. Co-Authored-By: Claude Opus 5 --- native/shuffle/src/shuffle_writer.rs | 33 +- .../writers/local/local_partition_writer.rs | 231 +++++++++++--- native/shuffle/src/writers/local/spill.rs | 286 ++++++++++++------ 3 files changed, 403 insertions(+), 147 deletions(-) diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index d6aaf0ae0e9..b1930f7d79a 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -604,21 +604,19 @@ mod test { repartitioner.insert_batch(batch.clone()).await.unwrap(); - { - let spill_writers = repartitioner.partition_writer().get_spill_writers(); - assert_eq!(spill_writers.len(), 2); - - assert!(!spill_writers[0].has_spill_file()); - assert!(!spill_writers[1].has_spill_file()); - } + assert!(!repartitioner + .partition_writer() + .get_spill() + .has_spill_file()); repartitioner.spill(0).unwrap(); - // after spill, there should be spill files + // after spill, both partitions' blocks are in the one spill file { - let spill_writers = repartitioner.partition_writer().get_spill_writers(); - assert!(spill_writers[0].has_spill_file()); - assert!(spill_writers[1].has_spill_file()); + let spill = repartitioner.partition_writer().get_spill(); + assert!(spill.has_spill_file()); + assert!(!spill.ranges(0).unwrap().is_empty()); + assert!(!spill.ranges(1).unwrap().is_empty()); } // insert another batch after spilling @@ -884,13 +882,14 @@ mod test { } repartitioner.shuffle_write().unwrap(); - let actual_spilled_bytes: usize = repartitioner + let actual_spilled_bytes = repartitioner .partition_writer() - .get_spill_writers() - .iter() - .filter_map(|writer| writer.path().unwrap()) - .map(|path| usize::try_from(std::fs::metadata(path).unwrap().len()).unwrap()) - .sum(); + .get_spill() + .path() + .unwrap() + .map_or(0, |path| { + usize::try_from(std::fs::metadata(path).unwrap().len()).unwrap() + }); assert_eq!( spilled_bytes.value(), actual_spilled_bytes, diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 984911470b9..1089b999907 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -16,7 +16,7 @@ // under the License. use crate::metrics::ShufflePartitionerMetrics; -use crate::writers::local::spill::SpillWriter; +use crate::writers::local::spill::PartitionedSpill; use crate::writers::partition_writer::PartitionWriter; use crate::writers::BufBatchWriter; use crate::{PartitionOffsets, ShuffleBlockWriter}; @@ -24,7 +24,7 @@ use arrow::array::RecordBatch; use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; use std::fs::{File, OpenOptions}; -use std::io::{BufWriter, Seek, Write}; +use std::io::{BufWriter, Read, Seek, SeekFrom, Write}; use std::sync::Arc; /// Output target for the shuffle data file. @@ -54,10 +54,11 @@ enum DataOutput { Multi { output_writer: BufWriter, shuffle_block_writer: ShuffleBlockWriter, - /// One spill file per output partition, buffered until `finish_partition` - /// merges them into the shuffle output. - spill_writers: Vec, - /// Runtime used to allocate the temporary spill files. + /// Spilled blocks for every partition, in one file. + spill: PartitionedSpill, + /// Read handle on the spill file and its position, opened on first use. + spill_reader: Option<(File, u64)>, + /// Runtime used to allocate the temporary spill file. runtime: Arc, /// Byte buffer recycled through the short-lived per-partition `BufBatchWriter`s. /// Partitions are written strictly one at a time, so a single buffer keeps its @@ -116,19 +117,17 @@ impl LocalPartitionWriter { } } else { let output_writer = BufWriter::with_capacity(write_buffer_size, output_file); - let spill_writers = (0..num_output_partitions) - .map(|_| { - SpillWriter::try_new( - shuffle_block_writer.clone(), - write_buffer_size, - batch_size, - ) - }) - .collect::>>()?; + let spill = PartitionedSpill::new( + shuffle_block_writer.clone(), + write_buffer_size, + batch_size, + num_output_partitions, + ); DataOutput::Multi { output_writer, shuffle_block_writer, - spill_writers, + spill, + spill_reader: None, runtime, recycled_buffer: Vec::new(), } @@ -145,10 +144,10 @@ impl LocalPartitionWriter { } #[cfg(test)] - pub(crate) fn get_spill_writers(&self) -> &Vec { + pub(crate) fn get_spill(&self) -> &PartitionedSpill { match &self.data_output { - DataOutput::Multi { spill_writers, .. } => spill_writers, - DataOutput::Single { .. } => panic!("single-partition output has no spill writers"), + DataOutput::Multi { spill, .. } => spill, + DataOutput::Single { .. } => panic!("single-partition output does not spill"), } } } @@ -183,15 +182,12 @@ impl PartitionWriter for LocalPartitionWriter { } } DataOutput::Multi { - spill_writers, + spill, runtime, recycled_buffer, .. } => { - // Multi-partition output buffers each partition's batches into its own - // spill file. `finish_partition` later merges the spill files (and any - // remaining in-memory batches) into the shuffle output in partition order. - spill_writers[pid].write(iter, runtime, metrics, recycled_buffer)?; + spill.write(pid, iter, runtime, metrics, recycled_buffer)?; } } @@ -231,23 +227,41 @@ impl PartitionWriter for LocalPartitionWriter { DataOutput::Multi { output_writer, shuffle_block_writer, - spill_writers, + spill, + spill_reader, recycled_buffer, .. } => { self.offsets[pid] = output_writer.stream_position()?; - // if we wrote a spill file for this partition then copy the - // contents into the shuffle file - if let Some(writer) = spill_writers.get(pid) { - if let Some(spill_path) = writer.path()? { - // Use raw File handle (not BufReader) so that std::io::copy - // can use copy_file_range/sendfile for zero-copy on Linux. - let mut spill_file = File::open(spill_path)?; - let mut write_timer = metrics.write_time.timer(); - std::io::copy(&mut spill_file, output_writer)?; - write_timer.stop(); + let ranges = spill.ranges(pid)?; + if !ranges.is_empty() { + if spill_reader.is_none() { + let path = spill.path()?.ok_or_else(|| { + DataFusionError::Internal( + "shuffle spill ranges recorded without a spill file".to_string(), + ) + })?; + *spill_reader = Some((File::open(path)?, 0)); } + let (spill_file, position) = spill_reader.as_mut().unwrap(); + let mut write_timer = metrics.write_time.timer(); + for range in ranges { + if *position != range.start { + spill_file.seek(SeekFrom::Start(range.start))?; + } + // raw File, not BufReader, so the copy can use copy_file_range on Linux + let len = range.end - range.start; + let copied = + std::io::copy(&mut Read::by_ref(spill_file).take(len), output_writer)?; + if copied != len { + return Err(DataFusionError::Execution(format!( + "shuffle spill file truncated: copied {copied} of {len} bytes" + ))); + } + *position = range.end; + } + write_timer.stop(); } // Write in memory batches to output data file. Each partition uses its @@ -329,7 +343,10 @@ mod tests { use crate::writers::local::spill::pathless_backend; use crate::CompressionCodec; use arrow::array::Int64Array; + use arrow::compute::concat_batches; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; fn test_batch() -> RecordBatch { @@ -345,6 +362,15 @@ mod tests { batch: &RecordBatch, dir: &tempfile::TempDir, runtime: Arc, + ) -> LocalPartitionWriter { + partition_writer_with(batch, 2, dir, runtime) + } + + fn partition_writer_with( + batch: &RecordBatch, + num_partitions: usize, + dir: &tempfile::TempDir, + runtime: Arc, ) -> LocalPartitionWriter { let block_writer = ShuffleBlockWriter::try_new(batch.schema_ref().as_ref(), CompressionCodec::None) @@ -353,7 +379,7 @@ mod tests { dir.path().join("data.out").to_str().unwrap().to_string(), Arc::new(PartitionOffsets::default()), block_writer, - 2, + num_partitions, // batch_size below the row count so the write serializes into the scratch. 10, 1 << 20, @@ -412,4 +438,137 @@ mod tests { "unexpected error: {err}" ); } + + fn int_batch(values: std::ops::Range) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from_iter_values(values))]).unwrap() + } + + fn decode_blocks(bytes: &[u8]) -> Vec { + let mut batches = Vec::new(); + let mut pos = 0; + while pos < bytes.len() { + let len = u64::from_le_bytes(bytes[pos..pos + 8].try_into().unwrap()) as usize; + batches.push(crate::read_ipc_compressed(&bytes[pos + 16..pos + 8 + len]).unwrap()); + pos += 8 + len; + } + batches + } + + fn count_files(dir: &std::path::Path) -> usize { + std::fs::read_dir(dir) + .unwrap() + .map(|entry| { + let path = entry.unwrap().path(); + if path.is_dir() { + count_files(&path) + } else { + 1 + } + }) + .sum() + } + + /// Partition data spread over spill rounds written in different partition orders, plus a + /// final in-memory batch, reads back per partition in write order. + #[test] + fn spilled_partitions_read_back_in_write_order() { + let dir = tempfile::tempdir().unwrap(); + let schema = test_batch().schema(); + let mut writer = + partition_writer_with(&test_batch(), 4, &dir, Arc::new(RuntimeEnv::default())); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut expected: Vec> = vec![Vec::new(); 4]; + let mut next = 0i64; + let mut batch = || { + let batch = int_batch(next..next + 10); + next += 10; + batch + }; + + for order in [[3, 2, 1, 0], [0, 1, 2, 3]] { + for pid in order { + let b = batch(); + expected[pid].push(b.clone()); + writer + .write(pid, &mut vec![Ok(b)].into_iter(), &metrics) + .unwrap(); + } + } + for (pid, rows) in expected.iter_mut().enumerate() { + let b = batch(); + rows.push(b.clone()); + writer + .finish_partition(pid, &mut vec![Ok(b)].into_iter(), &metrics) + .unwrap(); + } + writer.finish_all(&metrics).unwrap(); + + let offsets = writer.partition_offsets.get().unwrap().to_vec(); + let data = std::fs::read(dir.path().join("data.out")).unwrap(); + for (pid, rows) in expected.iter().enumerate() { + let bytes = &data[offsets[pid] as usize..offsets[pid + 1] as usize]; + let actual = concat_batches(&schema, &decode_blocks(bytes)).unwrap(); + assert_eq!( + actual, + concat_batches(&schema, rows).unwrap(), + "partition {pid}" + ); + } + } + + /// A task spills to one file however many partitions it has. + #[test] + fn spilling_every_partition_creates_one_file() { + let spill_dir = tempfile::tempdir().unwrap(); + let output_dir = tempfile::tempdir().unwrap(); + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_disk_manager_builder(DiskManagerBuilder::default().with_mode( + DiskManagerMode::Directories(vec![spill_dir.path().to_path_buf()]), + )) + .build() + .unwrap(), + ); + let num_partitions = 64; + let mut writer = partition_writer_with(&test_batch(), num_partitions, &output_dir, runtime); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + + for _ in 0..3 { + for pid in 0..num_partitions { + writer + .write(pid, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); + } + } + assert_eq!(count_files(spill_dir.path()), 1); + } + + /// A spill file shorter than its recorded ranges fails the task instead of writing a short + /// partition. + #[test] + fn finish_partition_fails_when_spill_file_is_truncated() { + let dir = tempfile::tempdir().unwrap(); + let mut writer = partition_writer(&test_batch(), &dir, Arc::new(RuntimeEnv::default())); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + writer + .write(0, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); + + let path = writer.get_spill().path().unwrap().unwrap().to_path_buf(); + std::fs::OpenOptions::new() + .write(true) + .open(&path) + .unwrap() + .set_len(1) + .unwrap(); + + let err = writer + .finish_partition(0, &mut std::iter::empty(), &metrics) + .expect_err("a truncated spill file must fail the partition"); + assert!( + err.to_string().contains("truncated"), + "unexpected error: {err}" + ); + } } diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index 77fe009046d..714af2c249f 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -23,6 +23,7 @@ use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::execution::SpillFile as DfSpillFile; use datafusion::execution::SpillWriter as DfSpillWriter; +use std::ops::Range; use std::sync::Arc; struct ActiveSpillFile { @@ -30,108 +31,120 @@ struct ActiveSpillFile { writer: Box, } -pub(crate) struct SpillWriter { +/// One spill file shared by every output partition of a task, with the ranges each partition's +/// blocks occupy. +pub(crate) struct PartitionedSpill { shuffle_block_writer: ShuffleBlockWriter, write_buffer_size: usize, batch_size: usize, spill_file: Option, + /// Bytes appended to the spill file so far. + len: u64, + /// Per partition, the spill file ranges holding its blocks, in write order. + ranges: Vec>>, + /// Set when a write fails partway, after which `len` may not match the file. + failed: bool, } -impl SpillWriter { - pub(crate) fn try_new( +impl PartitionedSpill { + pub(crate) fn new( shuffle_block_writer: ShuffleBlockWriter, write_buffer_size: usize, batch_size: usize, - ) -> datafusion::common::Result { - Ok(Self { + num_partitions: usize, + ) -> Self { + Self { shuffle_block_writer, write_buffer_size, batch_size, spill_file: None, - }) + len: 0, + ranges: vec![Vec::new(); num_partitions], + failed: false, + } } - /// `recycled_buffer` is a scratch byte buffer shared by the sequential per-partition - /// spill writes; it is left drained on return so one buffer's capacity serves every - /// partition instead of each write regrowing its own. + /// Appends partition `pid`'s batches to the spill file. `recycled_buffer` is left drained, + /// including on error. pub(crate) fn write>>( &mut self, + pid: usize, iter: &mut I, runtime: &RuntimeEnv, metrics: &ShufflePartitionerMetrics, recycled_buffer: &mut Vec, ) -> datafusion::common::Result<()> { - if let Some(batch) = iter.next() { - self.ensure_spill_file_created(runtime)?; - - let result = (|| { - let mut buf_batch_writer = BufBatchWriter::new( - &mut self.shuffle_block_writer, - &mut self.spill_file.as_mut().unwrap().writer, - self.write_buffer_size, - self.batch_size, - ); + self.check_usable()?; + let Some(batch) = iter.next() else { + return Ok(()); + }; + self.ensure_spill_file_created(runtime)?; + + let result = (|| { + let mut buf_batch_writer = BufBatchWriter::new( + &mut self.shuffle_block_writer, + &mut self.spill_file.as_mut().unwrap().writer, + self.write_buffer_size, + self.batch_size, + ); + buf_batch_writer.write( + &batch?, + recycled_buffer, + &metrics.encode_time, + &metrics.write_time, + )?; + for batch in iter.by_ref() { + let batch = batch?; buf_batch_writer.write( - &batch?, - recycled_buffer, - &metrics.encode_time, - &metrics.write_time, - )?; - for batch in iter.by_ref() { - let batch = batch?; - buf_batch_writer.write( - &batch, - recycled_buffer, - &metrics.encode_time, - &metrics.write_time, - )?; - } - buf_batch_writer.flush( + &batch, recycled_buffer, &metrics.encode_time, &metrics.write_time, )?; - // `SpillWriter` is not `Seek`, so bytes are tracked by the writer itself rather - // than measured via stream position. - let bytes_written = buf_batch_writer.bytes_written(); - usize::try_from(bytes_written).map_err(|_| { - DataFusionError::Execution(format!( - "Spill file byte count exceeds platform capacity: {bytes_written}" - )) - }) - })(); - // An errored spill must hand back a drained buffer, or its bytes leak into - // the next partition's block. - let total_bytes_written = result.inspect_err(|_| recycled_buffer.clear())?; - metrics.spilled_bytes.add(total_bytes_written); + } + buf_batch_writer.flush(recycled_buffer, &metrics.encode_time, &metrics.write_time)?; + Ok::<_, DataFusionError>(buf_batch_writer.bytes_written()) + })(); + + let bytes_written = match result { + Ok(bytes_written) => bytes_written, + Err(error) => { + // bytes may already be in the file, so later ranges could not be trusted + self.failed = true; + recycled_buffer.clear(); + return Err(error); + } + }; + + if bytes_written > 0 { + let start = self.len; + self.len += bytes_written; + self.ranges[pid].push(start..self.len); } + metrics + .spilled_bytes + .add(usize::try_from(bytes_written).map_err(|_| { + DataFusionError::Execution(format!( + "Spill file byte count exceeds platform capacity: {bytes_written}" + )) + })?); Ok(()) } - fn ensure_spill_file_created( - &mut self, - runtime: &RuntimeEnv, - ) -> datafusion::common::Result<()> { - if self.spill_file.is_none() { - // Spill file is not yet created, create it - let temp_file = runtime - .disk_manager - .create_tmp_file("shuffle writer spill")?; - let writer = temp_file.open_writer()?; - self.spill_file = Some(ActiveSpillFile { temp_file, writer }); - } - Ok(()) + /// The spill file ranges holding partition `pid`'s blocks, in write order. + pub(crate) fn ranges(&self, pid: usize) -> datafusion::common::Result<&[Range]> { + self.check_usable()?; + Ok(&self.ranges[pid]) } - /// Local filesystem path holding this partition's spilled bytes. + /// Local filesystem path of the spill file. /// - /// * `Ok(None)` — nothing was spilled for this partition. + /// * `Ok(None)` — nothing was spilled. /// * `Ok(Some(path))` — the spilled bytes live at `path`. /// * `Err(..)` — bytes were spilled but the backend exposes no local path. /// - /// The last case must stay distinct from `Ok(None)`: a caller that treated it as - /// "nothing to copy" would drop the spilled bytes while still recording the - /// partition offsets, silently truncating the partition in the shuffle file. + /// The last case must stay distinct from `Ok(None)`, or spilled bytes would be dropped while + /// partition offsets still counted them. pub(crate) fn path(&self) -> datafusion::common::Result> { match self.spill_file.as_ref() { None => Ok(None), @@ -146,6 +159,29 @@ impl SpillWriter { } } + fn check_usable(&self) -> datafusion::common::Result<()> { + if self.failed { + return Err(DataFusionError::Execution( + "Shuffle spill file is unusable after a failed write".to_string(), + )); + } + Ok(()) + } + + fn ensure_spill_file_created( + &mut self, + runtime: &RuntimeEnv, + ) -> datafusion::common::Result<()> { + if self.spill_file.is_none() { + let temp_file = runtime + .disk_manager + .create_tmp_file("shuffle writer spill")?; + let writer = temp_file.open_writer()?; + self.spill_file = Some(ActiveSpillFile { temp_file, writer }); + } + Ok(()) + } + #[cfg(test)] pub(crate) fn has_spill_file(&self) -> bool { self.spill_file.is_some() @@ -251,32 +287,34 @@ mod tests { .unwrap() } - fn spill_writer(batch: &RecordBatch, batch_size: usize) -> SpillWriter { + fn partitioned_spill(batch: &RecordBatch, num_partitions: usize) -> PartitionedSpill { let block_writer = ShuffleBlockWriter::try_new(batch.schema_ref().as_ref(), CompressionCodec::None) .unwrap(); - SpillWriter::try_new(block_writer, 1 << 20, batch_size).unwrap() + // batch_size below the row count so a write serializes into the scratch + PartitionedSpill::new(block_writer, 1 << 20, 10, num_partitions) } - /// A spill whose batch iterator fails after a batch was already encoded must hand - /// back a drained scratch; leftover bytes would land in the next partition's block. - #[test] - fn write_error_drains_recycled_buffer() { - let batch = test_batch(); - // batch_size below the row count so the first write serializes into the scratch. - let mut spill = spill_writer(&batch, 10); - let runtime = RuntimeEnv::default(); - let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); - let mut recycled = Vec::new(); + fn metrics() -> ShufflePartitionerMetrics { + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0) + } + + fn failing_write(spill: &mut PartitionedSpill, recycled: &mut Vec) { let mut iter = vec![ - Ok(batch), + Ok(test_batch()), Err(DataFusionError::Execution("injected failure".to_string())), ] .into_iter(); - assert!(spill - .write(&mut iter, &runtime, &metrics, &mut recycled) + .write(0, &mut iter, &RuntimeEnv::default(), &metrics(), recycled) .is_err()); + } + + #[test] + fn write_error_drains_recycled_buffer() { + let mut spill = partitioned_spill(&test_batch(), 2); + let mut recycled = Vec::new(); + failing_write(&mut spill, &mut recycled); assert!( recycled.is_empty(), "errored spill left {} bytes in the recycled buffer", @@ -284,34 +322,94 @@ mod tests { ); } - /// A partition that never spilled has no path, and that is not an error. + #[test] + fn failed_write_makes_spill_unusable() { + let mut spill = partitioned_spill(&test_batch(), 2); + let mut recycled = Vec::new(); + failing_write(&mut spill, &mut recycled); + + let err = spill + .write( + 1, + &mut vec![Ok(test_batch())].into_iter(), + &RuntimeEnv::default(), + &metrics(), + &mut recycled, + ) + .expect_err("write after a failed write"); + assert!( + err.to_string().contains("unusable"), + "unexpected error: {err}" + ); + assert!(spill.ranges(1).is_err()); + } + + #[test] + fn partitions_share_one_file_in_write_order() { + let mut spill = partitioned_spill(&test_batch(), 2); + let runtime = RuntimeEnv::default(); + let mut recycled = Vec::new(); + for pid in [1, 0, 1] { + spill + .write( + pid, + &mut vec![Ok(test_batch())].into_iter(), + &runtime, + &metrics(), + &mut recycled, + ) + .unwrap(); + } + + let first = spill.ranges(1).unwrap().to_vec(); + let second = spill.ranges(0).unwrap().to_vec(); + assert_eq!(first.len(), 2); + assert_eq!(second.len(), 1); + assert_eq!(first[0].start, 0); + assert_eq!(second[0].start, first[0].end); + assert_eq!(first[1].start, second[0].end); + } + + #[test] + fn empty_write_records_no_range() { + let mut spill = partitioned_spill(&test_batch(), 2); + spill + .write( + 0, + &mut std::iter::empty(), + &RuntimeEnv::default(), + &metrics(), + &mut Vec::new(), + ) + .unwrap(); + assert!(!spill.has_spill_file()); + assert!(spill.ranges(0).unwrap().is_empty()); + } + #[test] fn path_is_none_when_nothing_spilled() { - let batch = test_batch(); - let spill = spill_writer(&batch, 10); + let spill = partitioned_spill(&test_batch(), 2); assert!(!spill.has_spill_file()); assert_eq!(spill.path().unwrap(), None); } - /// Spilling to a backend with no local path must report an error rather than the - /// `None` that means "nothing spilled" — see `path`'s doc comment. #[test] fn path_errors_when_backend_has_no_local_path() { - let batch = test_batch(); - let mut spill = spill_writer(&batch, 10); - let runtime = pathless_backend::runtime(); - let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); - let mut recycled = Vec::new(); - let mut iter = vec![Ok(batch)].into_iter(); - + let mut spill = partitioned_spill(&test_batch(), 2); spill - .write(&mut iter, &runtime, &metrics, &mut recycled) + .write( + 0, + &mut vec![Ok(test_batch())].into_iter(), + &pathless_backend::runtime(), + &metrics(), + &mut Vec::new(), + ) .unwrap(); assert!(spill.has_spill_file()); let err = spill .path() - .expect_err("a spill file with no local path must not look like an empty partition"); + .expect_err("a spill file with no local path must not look like nothing spilled"); assert!( err.to_string().contains("no local path"), "unexpected error: {err}" From fc9d64b59a6bd0cf5672bda0f6a2ca844c88ee01 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 13 Sep 2026 20:41:36 +0800 Subject: [PATCH 7/8] perf: merge short spill ranges with one positional read Copying a spill range with io::copy costs an lseek, two statx calls and a copy_file_range. Ranges that fit in the write buffer are now read with one read_exact_at into a scratch buffer; longer ranges still use io::copy. Co-Authored-By: Claude Opus 5 --- .../writers/local/local_partition_writer.rs | 209 +++++++++++------- 1 file changed, 129 insertions(+), 80 deletions(-) diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 1089b999907..5d2475f898b 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -24,7 +24,9 @@ use arrow::array::RecordBatch; use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; use std::fs::{File, OpenOptions}; -use std::io::{BufWriter, Read, Seek, SeekFrom, Write}; +use std::io::{BufWriter, ErrorKind, Read, Seek, SeekFrom, Write}; +use std::ops::Range; +use std::os::unix::fs::FileExt; use std::sync::Arc; /// Output target for the shuffle data file. @@ -56,8 +58,9 @@ enum DataOutput { shuffle_block_writer: ShuffleBlockWriter, /// Spilled blocks for every partition, in one file. spill: PartitionedSpill, - /// Read handle on the spill file and its position, opened on first use. - spill_reader: Option<(File, u64)>, + /// Read handle on the spill file and a write-buffer-sized scratch for its ranges, + /// opened on first use. + spill_reader: Option<(File, Vec)>, /// Runtime used to allocate the temporary spill file. runtime: Arc, /// Byte buffer recycled through the short-lived per-partition `BufBatchWriter`s. @@ -242,24 +245,12 @@ impl PartitionWriter for LocalPartitionWriter { "shuffle spill ranges recorded without a spill file".to_string(), ) })?; - *spill_reader = Some((File::open(path)?, 0)); + *spill_reader = Some((File::open(path)?, vec![0; write_buffer_size])); } - let (spill_file, position) = spill_reader.as_mut().unwrap(); + let (spill_file, buffer) = spill_reader.as_mut().unwrap(); let mut write_timer = metrics.write_time.timer(); for range in ranges { - if *position != range.start { - spill_file.seek(SeekFrom::Start(range.start))?; - } - // raw File, not BufReader, so the copy can use copy_file_range on Linux - let len = range.end - range.start; - let copied = - std::io::copy(&mut Read::by_ref(spill_file).take(len), output_writer)?; - if copied != len { - return Err(DataFusionError::Execution(format!( - "shuffle spill file truncated: copied {copied} of {len} bytes" - ))); - } - *position = range.end; + copy_spill_range(spill_file, buffer, range, output_writer)?; } write_timer.stop(); } @@ -337,6 +328,45 @@ impl PartitionWriter for LocalPartitionWriter { } } +/// Appends `range` of the spill file to `output`, reading it through `buffer` when it fits. +fn copy_spill_range( + spill_file: &mut File, + buffer: &mut [u8], + range: &Range, + output: &mut BufWriter, +) -> datafusion::common::Result<()> { + let len = range.end - range.start; + let truncated = || { + DataFusionError::Execution(format!( + "shuffle spill file truncated: range {range:?} extends past its end" + )) + }; + match usize::try_from(len) + .ok() + .and_then(|len| buffer.get_mut(..len)) + { + // one pread instead of io::copy's lseek, two statx and copy_file_range + Some(chunk) => { + spill_file + .read_exact_at(chunk, range.start) + .map_err(|e| match e.kind() { + ErrorKind::UnexpectedEof => truncated(), + _ => e.into(), + })?; + output.write_all(chunk)?; + } + None => { + spill_file.seek(SeekFrom::Start(range.start))?; + // raw File, not BufReader, so the copy can use copy_file_range on Linux + let copied = std::io::copy(&mut Read::by_ref(spill_file).take(len), output)?; + if copied != len { + return Err(truncated()); + } + } + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -363,12 +393,13 @@ mod tests { dir: &tempfile::TempDir, runtime: Arc, ) -> LocalPartitionWriter { - partition_writer_with(batch, 2, dir, runtime) + partition_writer_with(batch, 2, 1 << 20, dir, runtime) } fn partition_writer_with( batch: &RecordBatch, num_partitions: usize, + write_buffer_size: usize, dir: &tempfile::TempDir, runtime: Arc, ) -> LocalPartitionWriter { @@ -382,7 +413,7 @@ mod tests { num_partitions, // batch_size below the row count so the write serializes into the scratch. 10, - 1 << 20, + write_buffer_size, runtime, ) .unwrap() @@ -470,50 +501,59 @@ mod tests { } /// Partition data spread over spill rounds written in different partition orders, plus a - /// final in-memory batch, reads back per partition in write order. + /// final in-memory batch, reads back per partition in write order, whether the spilled + /// ranges are read through the scratch buffer or copied. #[test] fn spilled_partitions_read_back_in_write_order() { - let dir = tempfile::tempdir().unwrap(); - let schema = test_batch().schema(); - let mut writer = - partition_writer_with(&test_batch(), 4, &dir, Arc::new(RuntimeEnv::default())); - let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); - let mut expected: Vec> = vec![Vec::new(); 4]; - let mut next = 0i64; - let mut batch = || { - let batch = int_batch(next..next + 10); - next += 10; - batch - }; - - for order in [[3, 2, 1, 0], [0, 1, 2, 3]] { - for pid in order { + // a 64-byte write buffer is smaller than one block, so every range is copied + for write_buffer_size in [1 << 20, 64] { + let dir = tempfile::tempdir().unwrap(); + let schema = test_batch().schema(); + let mut writer = partition_writer_with( + &test_batch(), + 4, + write_buffer_size, + &dir, + Arc::new(RuntimeEnv::default()), + ); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut expected: Vec> = vec![Vec::new(); 4]; + let mut next = 0i64; + let mut batch = || { + let batch = int_batch(next..next + 10); + next += 10; + batch + }; + + for order in [[3, 2, 1, 0], [0, 1, 2, 3]] { + for pid in order { + let b = batch(); + expected[pid].push(b.clone()); + writer + .write(pid, &mut vec![Ok(b)].into_iter(), &metrics) + .unwrap(); + } + } + for (pid, rows) in expected.iter_mut().enumerate() { let b = batch(); - expected[pid].push(b.clone()); + rows.push(b.clone()); writer - .write(pid, &mut vec![Ok(b)].into_iter(), &metrics) + .finish_partition(pid, &mut vec![Ok(b)].into_iter(), &metrics) .unwrap(); } - } - for (pid, rows) in expected.iter_mut().enumerate() { - let b = batch(); - rows.push(b.clone()); - writer - .finish_partition(pid, &mut vec![Ok(b)].into_iter(), &metrics) - .unwrap(); - } - writer.finish_all(&metrics).unwrap(); - - let offsets = writer.partition_offsets.get().unwrap().to_vec(); - let data = std::fs::read(dir.path().join("data.out")).unwrap(); - for (pid, rows) in expected.iter().enumerate() { - let bytes = &data[offsets[pid] as usize..offsets[pid + 1] as usize]; - let actual = concat_batches(&schema, &decode_blocks(bytes)).unwrap(); - assert_eq!( - actual, - concat_batches(&schema, rows).unwrap(), - "partition {pid}" - ); + writer.finish_all(&metrics).unwrap(); + + let offsets = writer.partition_offsets.get().unwrap().to_vec(); + let data = std::fs::read(dir.path().join("data.out")).unwrap(); + for (pid, rows) in expected.iter().enumerate() { + let bytes = &data[offsets[pid] as usize..offsets[pid + 1] as usize]; + let actual = concat_batches(&schema, &decode_blocks(bytes)).unwrap(); + assert_eq!( + actual, + concat_batches(&schema, rows).unwrap(), + "partition {pid}, write buffer {write_buffer_size}" + ); + } } } @@ -531,7 +571,8 @@ mod tests { .unwrap(), ); let num_partitions = 64; - let mut writer = partition_writer_with(&test_batch(), num_partitions, &output_dir, runtime); + let mut writer = + partition_writer_with(&test_batch(), num_partitions, 1 << 20, &output_dir, runtime); let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); for _ in 0..3 { @@ -545,30 +586,38 @@ mod tests { } /// A spill file shorter than its recorded ranges fails the task instead of writing a short - /// partition. + /// partition, whether the range is read through the scratch buffer or copied. #[test] fn finish_partition_fails_when_spill_file_is_truncated() { - let dir = tempfile::tempdir().unwrap(); - let mut writer = partition_writer(&test_batch(), &dir, Arc::new(RuntimeEnv::default())); - let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); - writer - .write(0, &mut vec![Ok(test_batch())].into_iter(), &metrics) - .unwrap(); + for write_buffer_size in [1 << 20, 64] { + let dir = tempfile::tempdir().unwrap(); + let mut writer = partition_writer_with( + &test_batch(), + 2, + write_buffer_size, + &dir, + Arc::new(RuntimeEnv::default()), + ); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + writer + .write(0, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); - let path = writer.get_spill().path().unwrap().unwrap().to_path_buf(); - std::fs::OpenOptions::new() - .write(true) - .open(&path) - .unwrap() - .set_len(1) - .unwrap(); + let path = writer.get_spill().path().unwrap().unwrap().to_path_buf(); + std::fs::OpenOptions::new() + .write(true) + .open(&path) + .unwrap() + .set_len(1) + .unwrap(); - let err = writer - .finish_partition(0, &mut std::iter::empty(), &metrics) - .expect_err("a truncated spill file must fail the partition"); - assert!( - err.to_string().contains("truncated"), - "unexpected error: {err}" - ); + let err = writer + .finish_partition(0, &mut std::iter::empty(), &metrics) + .expect_err("a truncated spill file must fail the partition"); + assert!( + err.to_string().contains("truncated"), + "write buffer {write_buffer_size}: unexpected error: {err}" + ); + } } } From a524afb22296af42ccc7d4e9dd511729975c6dcc Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 13 Sep 2026 20:45:09 +0800 Subject: [PATCH 8/8] perf: buffer spill writes across partitions Each partition's spill write ended with a flush, so a spill round issued one write syscall per partition. The spill file's writer is now buffered across partitions and flushed before the merge reads it. Co-Authored-By: Claude Opus 5 --- .../writers/local/local_partition_writer.rs | 13 +++- native/shuffle/src/writers/local/spill.rs | 64 ++++++++++++++++++- 2 files changed, 72 insertions(+), 5 deletions(-) diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 5d2475f898b..cf8ce6f5574 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -237,6 +237,9 @@ impl PartitionWriter for LocalPartitionWriter { } => { self.offsets[pid] = output_writer.stream_position()?; + let mut flush_timer = metrics.write_time.timer(); + spill.flush()?; + flush_timer.stop(); let ranges = spill.ranges(pid)?; if !ranges.is_empty() { if spill_reader.is_none() { @@ -599,8 +602,14 @@ mod tests { Arc::new(RuntimeEnv::default()), ); let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + for pid in 0..2 { + writer + .write(pid, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); + } + // finishing partition 0 flushes the buffered spill bytes before the file is cut writer - .write(0, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .finish_partition(0, &mut std::iter::empty(), &metrics) .unwrap(); let path = writer.get_spill().path().unwrap().unwrap().to_path_buf(); @@ -612,7 +621,7 @@ mod tests { .unwrap(); let err = writer - .finish_partition(0, &mut std::iter::empty(), &metrics) + .finish_partition(1, &mut std::iter::empty(), &metrics) .expect_err("a truncated spill file must fail the partition"); assert!( err.to_string().contains("truncated"), diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index 714af2c249f..a5d541aa51f 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -23,12 +23,33 @@ use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::execution::SpillFile as DfSpillFile; use datafusion::execution::SpillWriter as DfSpillWriter; +use std::io::{BufWriter, Write}; use std::ops::Range; use std::sync::Arc; struct ActiveSpillFile { temp_file: Arc, - writer: Box, + /// Shared by every partition; bytes reach the file when it fills or on + /// [`PartitionedSpill::flush`]. + writer: BufWriter>, +} + +/// Forwards writes but ignores `flush`, so `BufBatchWriter`'s flush at the end of each partition +/// leaves the bytes buffered. +struct DeferFlush<'a, W: Write>(&'a mut W); + +impl Write for DeferFlush<'_, W> { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.write(buf) + } + + fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + self.0.write_all(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } } /// One spill file shared by every output partition of a task, with the ranges each partition's @@ -83,7 +104,7 @@ impl PartitionedSpill { let result = (|| { let mut buf_batch_writer = BufBatchWriter::new( &mut self.shuffle_block_writer, - &mut self.spill_file.as_mut().unwrap().writer, + DeferFlush(&mut self.spill_file.as_mut().unwrap().writer), self.write_buffer_size, self.batch_size, ); @@ -137,6 +158,19 @@ impl PartitionedSpill { Ok(&self.ranges[pid]) } + /// Writes buffered spill bytes to the spill file. + pub(crate) fn flush(&mut self) -> datafusion::common::Result<()> { + self.check_usable()?; + if let Some(spill_file) = self.spill_file.as_mut() { + if let Err(error) = spill_file.writer.flush() { + // the file holds an unknown prefix of the buffered bytes + self.failed = true; + return Err(error.into()); + } + } + Ok(()) + } + /// Local filesystem path of the spill file. /// /// * `Ok(None)` — nothing was spilled. @@ -176,7 +210,7 @@ impl PartitionedSpill { let temp_file = runtime .disk_manager .create_tmp_file("shuffle writer spill")?; - let writer = temp_file.open_writer()?; + let writer = BufWriter::with_capacity(self.write_buffer_size, temp_file.open_writer()?); self.spill_file = Some(ActiveSpillFile { temp_file, writer }); } Ok(()) @@ -370,6 +404,30 @@ mod tests { assert_eq!(first[1].start, second[0].end); } + #[test] + fn writes_stay_buffered_until_flush() { + let mut spill = partitioned_spill(&test_batch(), 2); + let runtime = RuntimeEnv::default(); + let mut recycled = Vec::new(); + for pid in [0, 1] { + spill + .write( + pid, + &mut vec![Ok(test_batch())].into_iter(), + &runtime, + &metrics(), + &mut recycled, + ) + .unwrap(); + } + let path = spill.path().unwrap().unwrap().to_path_buf(); + let spilled = spill.ranges(1).unwrap()[0].end; + + assert_eq!(std::fs::metadata(&path).unwrap().len(), 0); + spill.flush().unwrap(); + assert_eq!(std::fs::metadata(&path).unwrap().len(), spilled); + } + #[test] fn empty_write_records_no_range() { let mut spill = partitioned_spill(&test_batch(), 2);