From 5336a9d8e0fe698680fc00869005159d4aedd20c Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 26 Aug 2026 04:07:06 +0000 Subject: [PATCH] feat: add RSS partition writer and task-owned JNI callback --- native/jni-bridge/src/lib.rs | 2 + .../src/shuffle_partition_pusher.rs | 147 ++++++ native/shuffle/src/lib.rs | 2 +- native/shuffle/src/writers/mod.rs | 2 + native/shuffle/src/writers/rss/mod.rs | 472 ++++++++++++++++++ .../src/writers/rss/rss_partition_writer.rs | 214 ++++++++ .../comet/shuffle/ShufflePartitionPusher.java | 35 ++ 7 files changed, 873 insertions(+), 1 deletion(-) create mode 100644 native/jni-bridge/src/shuffle_partition_pusher.rs create mode 100644 native/shuffle/src/writers/rss/mod.rs create mode 100644 native/shuffle/src/writers/rss/rss_partition_writer.rs create mode 100644 spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java diff --git a/native/jni-bridge/src/lib.rs b/native/jni-bridge/src/lib.rs index 8db4c07851a..51fdd736b03 100644 --- a/native/jni-bridge/src/lib.rs +++ b/native/jni-bridge/src/lib.rs @@ -195,6 +195,7 @@ mod comet_s3_credential_dispatcher; mod comet_task_memory_manager; mod comet_udf_bridge; mod shuffle_block_iterator; +mod shuffle_partition_pusher; use arrow_array_stream::ArrowArrayStream; pub use comet_metric_node::*; @@ -202,6 +203,7 @@ pub use comet_s3_credential_dispatcher::CometS3CredentialDispatcher; pub use comet_task_memory_manager::*; use comet_udf_bridge::CometUdfBridge; use shuffle_block_iterator::CometShuffleBlockIterator; +pub use shuffle_partition_pusher::{JavaShufflePartitionPusher, ShufflePartitionPusher}; /// The JVM classes that are used in the JNI calls. #[allow(dead_code)] // we need to keep references to Java items to prevent GC diff --git a/native/jni-bridge/src/shuffle_partition_pusher.rs b/native/jni-bridge/src/shuffle_partition_pusher.rs new file mode 100644 index 00000000000..e2503c77a84 --- /dev/null +++ b/native/jni-bridge/src/shuffle_partition_pusher.rs @@ -0,0 +1,147 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::{check_exception, errors::CometError, JVMClasses}; +use datafusion::common::{DataFusionError, Result}; +use jni::objects::{Global, JMethodID, JObject, JValue}; +use jni::signature::{Primitive, ReturnType}; +use jni::Env; + +/// Receives a complete encoded shuffle block for one output partition. +/// +/// Implementations must remain safe when invoked from native execution +/// threads that do not inherit Spark's task-local JVM state. +pub trait ShufflePartitionPusher: Send + Sync { + /// Sends one complete, length-prefixed Arrow IPC shuffle block. + fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()>; +} + +/// Invokes a task-owned JVM shuffle callback from any native execution thread. +/// +/// The global callback reference keeps both the Java object and its class alive +/// for the lifetime of the cached method ID. No thread-local JNI environment +/// or Spark task context is retained between invocations. +pub struct JavaShufflePartitionPusher { + callback: Global>, + push_method: JMethodID, +} + +impl JavaShufflePartitionPusher { + /// Captures the callback while running on an attached JVM thread. + pub fn try_new(env: &mut Env<'_>, callback: &JObject<'_>) -> Result { + if callback.is_null() { + return Err(DataFusionError::Execution( + "Remote shuffle callback must not be null".to_string(), + )); + } + + let callback_class = env.get_object_class(callback).map_err(CometError::from)?; + let push_method = env + .get_method_id( + &callback_class, + jni::jni_str!("pushPartitionData"), + jni::jni_sig!("(I[BI)V"), + ) + .map_err(CometError::from)?; + let callback = env.new_global_ref(callback).map_err(CometError::from)?; + + Ok(Self { + callback, + push_method, + }) + } + + fn checked_payload_length(partition_id: i32, payload_length: usize) -> Result { + if partition_id < 0 { + return Err(DataFusionError::Execution(format!( + "Remote shuffle partition must be nonnegative, got {partition_id}" + ))); + } + + i32::try_from(payload_length).map_err(|_| { + DataFusionError::Execution(format!( + "Remote shuffle payload size {payload_length} exceeds the JVM array limit" + )) + }) + } +} + +impl ShufflePartitionPusher for JavaShufflePartitionPusher { + fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()> { + let payload_length = Self::checked_payload_length(partition_id, data.len())?; + + JVMClasses::with_env(|env| { + let payload = env.byte_array_from_slice(data).map_err(CometError::from)?; + + // SAFETY: `push_method` was resolved against the callback object's + // class with this exact argument list and void return type. The + // global object reference keeps its defining class alive. + let result = unsafe { + env.call_method_unchecked( + self.callback.as_obj(), + self.push_method, + ReturnType::Primitive(Primitive::Void), + &[ + JValue::Int(partition_id).as_jni(), + JValue::Object(&payload).as_jni(), + JValue::Int(payload_length).as_jni(), + ], + ) + }; + + // Inspect the pending exception before consuming the JNI result so + // its original throwable survives the DataFusion error boundary. + if let Some(exception) = check_exception(env)? { + return Err(exception.into()); + } + + result.map_err(CometError::from)?; + Ok(()) + }) + } +} + +#[cfg(test)] +mod tests { + use super::JavaShufflePartitionPusher; + + #[test] + fn accepts_payload_lengths_representable_by_jvm_arrays() { + assert_eq!( + JavaShufflePartitionPusher::checked_payload_length(0, 0).unwrap(), + 0 + ); + assert_eq!( + JavaShufflePartitionPusher::checked_payload_length(7, i32::MAX as usize).unwrap(), + i32::MAX + ); + } + + #[test] + fn rejects_negative_partition_ids() { + let error = JavaShufflePartitionPusher::checked_payload_length(-1, 1).unwrap_err(); + assert!(error.to_string().contains("partition must be nonnegative")); + } + + #[test] + fn rejects_payload_lengths_larger_than_a_jvm_array() { + let oversized_length = i32::MAX as usize + 1; + let error = + JavaShufflePartitionPusher::checked_payload_length(0, oversized_length).unwrap_err(); + assert!(error.to_string().contains("exceeds the JVM array limit")); + } +} diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 2263ae0dac0..0fd10be1c11 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -29,4 +29,4 @@ pub use comet_partitioning::CometPartitioning; pub use ipc::read_ipc_compressed; pub use schema_align::SchemaAlignExec; pub use shuffle_writer::ShuffleWriterExec; -pub use writers::{CompressionCodec, ShuffleBlockWriter}; +pub use writers::{CompressionCodec, RssPartitionWriter, ShuffleBlockWriter}; diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index 6d330fd12aa..df4bf69122e 100644 --- a/native/shuffle/src/writers/mod.rs +++ b/native/shuffle/src/writers/mod.rs @@ -19,10 +19,12 @@ mod buf_batch_writer; mod checksum; mod local; mod partition_writer; +mod rss; mod shuffle_block_writer; pub(crate) use buf_batch_writer::BufBatchWriter; pub(crate) use checksum::Checksum; pub(crate) use local::local_partition_writer::LocalPartitionWriter; pub(crate) use partition_writer::PartitionWriter; +pub use rss::rss_partition_writer::RssPartitionWriter; pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs new file mode 100644 index 00000000000..c23d742e7dd --- /dev/null +++ b/native/shuffle/src/writers/rss/mod.rs @@ -0,0 +1,472 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub(crate) mod rss_partition_writer; + +#[cfg(test)] +mod tests { + use super::rss_partition_writer::RssPartitionWriter; + use crate::metrics::ShufflePartitionerMetrics; + use crate::writers::PartitionWriter; + use crate::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter}; + use arrow::array::{Array, DictionaryArray, Int32Array}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use arrow::ipc::writer::CompressionContext; + use arrow::record_batch::RecordBatch; + use datafusion::common::{DataFusionError, Result}; + use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time}; + use datafusion_comet_jni_bridge::ShufflePartitionPusher; + use std::io::Cursor; + use std::sync::{Arc, Mutex}; + + type RecordedFrame = (i32, Vec); + + #[derive(Default)] + struct RecordingPusher { + frames: Mutex>, + } + + impl RecordingPusher { + fn frames(&self) -> Vec { + self.frames.lock().unwrap().clone() + } + } + + impl ShufflePartitionPusher for RecordingPusher { + fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()> { + self.frames + .lock() + .unwrap() + .push((partition_id, data.to_vec())); + Ok(()) + } + } + + struct FailingPusher; + + impl ShufflePartitionPusher for FailingPusher { + fn push_partition_data(&self, _partition_id: i32, _data: &[u8]) -> Result<()> { + Err(DataFusionError::Execution( + "remote shuffle callback rejected the frame".to_string(), + )) + } + } + + fn sample_batch(start: i32, rows: i32) -> RecordBatch { + let values = Int32Array::from_iter_values(start..start + rows); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap() + } + + fn metrics() -> ShufflePartitionerMetrics { + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0) + } + + fn writer( + batch: &RecordBatch, + codec: CompressionCodec, + pusher: Arc, + partitions: usize, + max_frame_size: usize, + ) -> RssPartitionWriter { + let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec).unwrap(); + RssPartitionWriter::try_new(block_writer, pusher, partitions, max_frame_size).unwrap() + } + + fn write_batches( + writer: &mut RssPartitionWriter, + partition_id: usize, + batches: Vec, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> { + let mut batches = batches.into_iter().map(Ok); + writer.write(partition_id, &mut batches, metrics) + } + + fn finish_partition( + writer: &mut RssPartitionWriter, + partition_id: usize, + batches: Vec, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> { + let mut batches = batches.into_iter().map(Ok); + writer.finish_partition(partition_id, &mut batches, metrics) + } + + fn decode_frame(frame: &[u8]) -> RecordBatch { + assert!(frame.len() >= 20, "shuffle frame must include its header"); + let payload_length = u64::from_le_bytes(frame[..8].try_into().unwrap()); + assert_eq!( + usize::try_from(payload_length).unwrap() + 8, + frame.len(), + "each callback must receive exactly one complete shuffle frame" + ); + read_ipc_compressed(&frame[16..]).unwrap() + } + + fn encoded_frame_size(batch: &RecordBatch) -> usize { + let block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut frame = Cursor::new(Vec::new()); + let mut compression_context = CompressionContext::default(); + block_writer + .write_batch( + batch, + &mut frame, + &mut compression_context, + &Time::default(), + ) + .unwrap() + } + + #[test] + #[cfg_attr(miri, ignore)] + fn round_trips_all_supported_compression_codecs() { + let batch = sample_batch(10, 128); + + for codec in [ + CompressionCodec::None, + CompressionCodec::Lz4Frame, + CompressionCodec::Snappy, + CompressionCodec::Zstd(1), + ] { + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer(&batch, codec, pusher.clone(), 1, 1024 * 1024); + let metrics = metrics(); + + write_batches(&mut writer, 0, vec![batch.clone()], &metrics).unwrap(); + finish_partition(&mut writer, 0, vec![], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = pusher.frames(); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].0, 0); + assert_eq!(decode_frame(&frames[0].1), batch); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn preserves_dictionary_encoded_batches() { + let values = ["alpha", "beta", "alpha", "gamma"]; + let dictionary: DictionaryArray = values.iter().copied().collect(); + let schema = Arc::new(Schema::new(vec![Field::new( + "dictionary", + dictionary.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(dictionary)]).unwrap(); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::Zstd(1), + pusher.clone(), + 1, + 1024 * 1024, + ); + let metrics = metrics(); + + finish_partition(&mut writer, 0, vec![batch.clone()], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = pusher.frames(); + assert_eq!(frames.len(), 1); + assert_eq!(decode_frame(&frames[0].1), batch); + } + + #[test] + fn sends_each_batch_as_one_complete_frame() { + let first = sample_batch(0, 16); + let second = sample_batch(16, 16); + let third = sample_batch(32, 16); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &first, + CompressionCodec::None, + pusher.clone(), + 1, + 1024 * 1024, + ); + let metrics = metrics(); + + write_batches( + &mut writer, + 0, + vec![first.clone(), second.clone()], + &metrics, + ) + .unwrap(); + finish_partition(&mut writer, 0, vec![third.clone()], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = pusher.frames(); + assert_eq!(frames.len(), 3); + for ((partition_id, frame), expected) in frames.iter().zip([first, second, third].iter()) { + assert_eq!(*partition_id, 0); + assert_eq!(decode_frame(frame), *expected); + } + } + + #[test] + fn routes_out_of_order_writes_to_the_correct_partition() { + let first = sample_batch(0, 8); + let second = sample_batch(8, 8); + let third = sample_batch(16, 8); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &first, + CompressionCodec::None, + pusher.clone(), + 3, + 1024 * 1024, + ); + let metrics = metrics(); + + write_batches(&mut writer, 2, vec![third.clone()], &metrics).unwrap(); + write_batches(&mut writer, 0, vec![first.clone()], &metrics).unwrap(); + write_batches(&mut writer, 1, vec![second.clone()], &metrics).unwrap(); + for partition_id in 0..3 { + finish_partition(&mut writer, partition_id, vec![], &metrics).unwrap(); + } + writer.finish_all(&metrics).unwrap(); + + let frames = pusher.frames(); + assert_eq!(frames.len(), 3); + for ((partition_id, frame), (expected_id, expected_batch)) in frames + .iter() + .zip([(2, third), (0, first), (1, second)].iter()) + { + assert_eq!(partition_id, expected_id); + assert_eq!(decode_frame(frame), *expected_batch); + } + } + + #[test] + fn empty_batches_do_not_invoke_the_callback() { + let batch = sample_batch(0, 4); + let empty = RecordBatch::new_empty(batch.schema()); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 2, + 1024 * 1024, + ); + let metrics = metrics(); + + write_batches(&mut writer, 0, vec![empty.clone()], &metrics).unwrap(); + finish_partition(&mut writer, 0, vec![empty.clone()], &metrics).unwrap(); + finish_partition(&mut writer, 1, vec![empty], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + + assert!(pusher.frames().is_empty()); + } + + #[test] + fn callback_failures_are_preserved() { + let batch = sample_batch(0, 4); + let mut writer = writer( + &batch, + CompressionCodec::None, + Arc::new(FailingPusher), + 1, + 1024 * 1024, + ); + + let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err(); + assert!( + error + .to_string() + .contains("remote shuffle callback rejected the frame"), + "callback error should not be replaced: {error}" + ); + } + + #[test] + fn iterator_failures_are_preserved() { + let batch = sample_batch(0, 4); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 1, + 1024 * 1024, + ); + let mut failed_batches = std::iter::once(Err(DataFusionError::Execution( + "upstream shuffle input failed".to_string(), + ))); + + let error = writer + .write(0, &mut failed_batches, &metrics()) + .unwrap_err(); + assert!(error.to_string().contains("upstream shuffle input failed")); + assert!(pusher.frames().is_empty()); + } + + #[test] + fn accepts_frames_that_exactly_match_the_maximum_size() { + let batch = sample_batch(0, 16); + let frame_size = encoded_frame_size(&batch); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 1, + frame_size, + ); + let metrics = metrics(); + + finish_partition(&mut writer, 0, vec![batch.clone()], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = pusher.frames(); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].1.len(), frame_size); + assert_eq!(decode_frame(&frames[0].1), batch); + } + + #[test] + fn rejects_oversized_frames_without_splitting_or_pushing() { + let batch = sample_batch(0, 16); + let frame_size = encoded_frame_size(&batch); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 1, + frame_size - 1, + ); + + assert!(write_batches(&mut writer, 0, vec![batch], &metrics()).is_err()); + assert!( + pusher.frames().is_empty(), + "an oversized Arrow IPC frame must never be fragmented" + ); + } + + #[test] + fn constructor_rejects_invalid_limits() { + let batch = sample_batch(0, 1); + let pusher: Arc = Arc::new(RecordingPusher::default()); + let block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); + + assert!( + RssPartitionWriter::try_new(block_writer.clone(), pusher.clone(), 0, 1024).is_err() + ); + assert!(RssPartitionWriter::try_new(block_writer.clone(), pusher.clone(), 1, 0).is_err()); + + let excessive_partitions = usize::try_from(i32::MAX).unwrap() + 2; + assert!( + RssPartitionWriter::try_new(block_writer, pusher, excessive_partitions, 1024).is_err() + ); + } + + #[test] + fn rejects_partition_ids_outside_the_configured_range() { + let batch = sample_batch(0, 4); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 2, + 1024 * 1024, + ); + let metrics = metrics(); + + assert!(write_batches(&mut writer, 2, vec![batch.clone()], &metrics).is_err()); + assert!(finish_partition(&mut writer, 2, vec![batch], &metrics).is_err()); + assert!(pusher.frames().is_empty()); + } + + #[test] + fn partition_finalization_must_be_in_ascending_order() { + let batch = sample_batch(0, 4); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer(&batch, CompressionCodec::None, pusher, 2, 1024 * 1024); + let metrics = metrics(); + + assert!(finish_partition(&mut writer, 1, vec![], &metrics).is_err()); + finish_partition(&mut writer, 0, vec![], &metrics).unwrap(); + assert!(finish_partition(&mut writer, 0, vec![], &metrics).is_err()); + finish_partition(&mut writer, 1, vec![], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + } + + #[test] + fn rejects_writes_to_a_finalized_partition() { + let batch = sample_batch(0, 4); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 2, + 1024 * 1024, + ); + let metrics = metrics(); + + finish_partition(&mut writer, 0, vec![], &metrics).unwrap(); + assert!(write_batches(&mut writer, 0, vec![batch], &metrics).is_err()); + assert!(pusher.frames().is_empty()); + } + + #[test] + fn finish_all_requires_every_partition_to_be_finalized() { + let batch = sample_batch(0, 4); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer(&batch, CompressionCodec::None, pusher, 2, 1024 * 1024); + let metrics = metrics(); + + assert!(writer.finish_all(&metrics).is_err()); + finish_partition(&mut writer, 0, vec![], &metrics).unwrap(); + assert!(writer.finish_all(&metrics).is_err()); + finish_partition(&mut writer, 1, vec![], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + } + + #[test] + fn rejects_operations_after_finish_all() { + let batch = sample_batch(0, 4); + let pusher = Arc::new(RecordingPusher::default()); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 1, + 1024 * 1024, + ); + let metrics = metrics(); + + finish_partition(&mut writer, 0, vec![], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + + assert!(write_batches(&mut writer, 0, vec![batch], &metrics).is_err()); + assert!(finish_partition(&mut writer, 0, vec![], &metrics).is_err()); + assert!(writer.finish_all(&metrics).is_err()); + assert!(pusher.frames().is_empty()); + } +} diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs new file mode 100644 index 00000000000..be57870220f --- /dev/null +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -0,0 +1,214 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::metrics::ShufflePartitionerMetrics; +use crate::writers::partition_writer::PartitionWriter; +use crate::ShuffleBlockWriter; +use arrow::array::RecordBatch; +use arrow::ipc::writer::CompressionContext; +use datafusion::common::{DataFusionError, Result}; +use datafusion_comet_jni_bridge::ShufflePartitionPusher; +use std::io::Cursor; +use std::sync::Arc; + +/// Sends complete Comet shuffle blocks to a task-owned remote shuffle pusher. +/// +/// Each callback receives exactly one self-contained, length-prefixed Arrow +/// IPC block. Keeping block boundaries intact lets remote shuffle services +/// concatenate or retrieve the resulting payloads without repairing partially +/// encoded blocks. The encoded frame is checked against `max_frame_size` +/// before it is exposed to the pusher. +/// +/// Partition data may arrive in any order until its partition is finalized. +/// Finalization follows the same ascending-partition contract as the existing +/// local shuffle writer. +pub struct RssPartitionWriter { + block_writer: ShuffleBlockWriter, + pusher: Arc, + num_partitions: usize, + max_frame_size: usize, + compression_context: CompressionContext, + frame: Vec, + next_partition_to_finish: usize, + finished: bool, +} + +impl RssPartitionWriter { + /// Creates a remote writer without retaining any thread-local JNI state. + /// + /// `num_partitions` must be nonzero and each partition identifier must fit + /// in a JVM `int`. `max_frame_size` is the maximum complete encoded shuffle + /// block passed to a callback and must also be nonzero. + pub fn try_new( + block_writer: ShuffleBlockWriter, + pusher: Arc, + num_partitions: usize, + max_frame_size: usize, + ) -> Result { + if num_partitions == 0 { + return Err(DataFusionError::Execution( + "Remote shuffle output requires at least one partition".to_string(), + )); + } + + if num_partitions - 1 > i32::MAX as usize { + return Err(DataFusionError::Execution(format!( + "Remote shuffle partition count {num_partitions} exceeds the JVM partition limit" + ))); + } + + if max_frame_size == 0 { + return Err(DataFusionError::Execution( + "Remote shuffle maximum frame size must be greater than zero".to_string(), + )); + } + + Ok(Self { + block_writer, + pusher, + num_partitions, + max_frame_size, + compression_context: CompressionContext::default(), + frame: Vec::new(), + next_partition_to_finish: 0, + finished: false, + }) + } + + fn validate_writable_partition(&self, partition_id: usize) -> Result { + if self.finished { + return Err(DataFusionError::Execution( + "Remote shuffle writer has already finished".to_string(), + )); + } + + if partition_id >= self.num_partitions { + return Err(DataFusionError::Execution(format!( + "Remote shuffle partition {partition_id} is outside the configured range 0..{}", + self.num_partitions + ))); + } + + if partition_id < self.next_partition_to_finish { + return Err(DataFusionError::Execution(format!( + "Remote shuffle partition {partition_id} has already been finalized" + ))); + } + + i32::try_from(partition_id).map_err(|_| { + DataFusionError::Execution(format!( + "Remote shuffle partition {partition_id} exceeds the JVM partition limit" + )) + }) + } + + fn push_batches( + &mut self, + partition_id: i32, + batches: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + for batch in batches.by_ref() { + let batch = batch?; + self.frame.clear(); + + let encoded_size = self.block_writer.write_batch( + &batch, + &mut Cursor::new(&mut self.frame), + &mut self.compression_context, + &metrics.encode_time, + )?; + + if encoded_size == 0 { + continue; + } + + if encoded_size > self.max_frame_size { + return Err(DataFusionError::Execution(format!( + "Remote shuffle frame size {encoded_size} exceeds the configured maximum {}", + self.max_frame_size + ))); + } + + let mut write_timer = metrics.write_time.timer(); + let result = self.pusher.push_partition_data(partition_id, &self.frame); + write_timer.stop(); + result?; + } + + Ok(()) + } +} + +impl PartitionWriter for RssPartitionWriter { + fn write( + &mut self, + partition_id: usize, + batches: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + let jvm_partition_id = self.validate_writable_partition(partition_id)?; + self.push_batches(jvm_partition_id, batches, metrics) + } + + fn finish_partition( + &mut self, + partition_id: usize, + batches: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + let jvm_partition_id = self.validate_writable_partition(partition_id)?; + if partition_id != self.next_partition_to_finish { + return Err(DataFusionError::Execution(format!( + "Remote shuffle partitions must be finalized in order: expected {}, got {partition_id}", + self.next_partition_to_finish + ))); + } + + self.push_batches(jvm_partition_id, batches, metrics)?; + self.next_partition_to_finish += 1; + Ok(()) + } + + fn finish_all(&mut self, _metrics: &ShufflePartitionerMetrics) -> Result<()> { + if self.finished { + return Err(DataFusionError::Execution( + "Remote shuffle writer has already finished".to_string(), + )); + } + + if self.next_partition_to_finish != self.num_partitions { + return Err(DataFusionError::Execution(format!( + "Remote shuffle writer cannot finish before all partitions are finalized: \ + finalized {} of {}", + self.next_partition_to_finish, self.num_partitions + ))); + } + + self.finished = true; + Ok(()) + } +} diff --git a/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java new file mode 100644 index 00000000000..05771e921ae --- /dev/null +++ b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.shuffle; + +import java.io.IOException; + +/** + * Receives complete encoded shuffle blocks from a native partition writer. + * + *

Instances belong to one Spark task. Implementations must be safe to invoke from native worker + * threads, which do not inherit the Spark task thread's thread-local context. + */ +@FunctionalInterface +public interface ShufflePartitionPusher { + + /** Pushes one complete, length-prefixed Arrow IPC block for the given output partition. */ + void pushPartitionData(int partitionId, byte[] data, int length) throws IOException; +}