From 6658894632adf114fb86313c282b2ecf95c8453f Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Thu, 27 Aug 2026 01:26:15 +0000 Subject: [PATCH 1/2] feat: bind task-owned RSS callbacks to native shuffle plans --- native/core/src/execution/jni_api.rs | 49 ++- native/core/src/execution/planner.rs | 393 ++++++++++++++++-- .../src/shuffle_partition_pusher.rs | 10 + .../org/apache/comet/CometExecIterator.scala | 39 +- .../main/scala/org/apache/comet/Native.scala | 14 + .../comet/exec/CometNativeShuffleSuite.scala | 198 ++++++++- 6 files changed, 654 insertions(+), 49 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index af945adf9a..bb218bf108 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -24,7 +24,7 @@ use crate::{ metrics::utils::update_comet_metric, planner::PhysicalPlanner, serde::to_arrow_datatype, shuffle::spark_unsafe::row::process_sorted_row_partition, sort::RdxSort, }, - jvm_bridge::JVMClasses, + jvm_bridge::{JVMClasses, JavaShufflePartitionPusher, ShufflePartitionPusher}, }; use std::collections::HashSet; @@ -393,6 +393,9 @@ struct ExecutionContext { /// it has to travel with the plan. `None` when no driving Spark task is present (unit tests, /// direct native driver runs). Lifetime is as for `task_context` above. pub class_loader: Option>>>, + /// Task-owned remote shuffle callback, registered before native planning starts. + /// The callback owns a JNI global reference and can safely run on Tokio workers. + pub shuffle_partition_pusher: Option>, /// Removes this context's tracing memory-pool entry on every exit path. memory_pool_registration: Option, } @@ -583,6 +586,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( tracing_event_name, task_context, class_loader, + shuffle_partition_pusher: None, memory_pool_registration, }); @@ -591,6 +595,44 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( }) } +/// Binds one task-owned shuffle callback before native execution is initialized. +/// +/// Keeping callback registration separate preserves the existing `createPlan` JNI ABI for +/// all local shuffle and non-shuffle callers. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_Native_setShufflePartitionPusher( + e: EnvUnowned, + _class: JClass, + exec_context: jlong, + callback: JObject, +) { + try_unwrap_or_throw(&e, |env| { + if exec_context == 0 { + return Err(CometError::NullPointer( + "Remote shuffle callback requires a valid native execution plan".to_string(), + )); + } + + let exec_context = get_execution_context(exec_context); + if exec_context.root_op.is_some() { + return Err(CometError::Internal( + "Remote shuffle callback cannot be registered after native execution starts" + .to_string(), + )); + } + + if exec_context.shuffle_partition_pusher.is_some() { + return Err(CometError::Internal( + "Remote shuffle callback has already been registered for this task".to_string(), + )); + } + + let pusher = JavaShufflePartitionPusher::try_new(env, &callback)?; + exec_context.shuffle_partition_pusher = Some(Arc::new(pusher)); + Ok(()) + }) +} + /// Configure DataFusion session context. fn prepare_datafusion_session_context( batch_size: usize, @@ -831,7 +873,10 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( .with_exec_id(exec_context_id) .with_sql_text_pool(&exec_context.spark_plan) .with_task_context(exec_context.task_context.clone()) - .with_class_loader(exec_context.class_loader.clone()); + .with_class_loader(exec_context.class_loader.clone()) + .with_shuffle_partition_pusher( + exec_context.shuffle_partition_pusher.clone(), + ); let (scans, shuffle_scans, root_op) = planner.create_plan( &exec_context.spark_plan, &mut exec_context.input_sources.clone(), diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 7fb1c89b96..9d2a03f29a 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -45,9 +45,9 @@ use crate::execution::{ planner::expression_registry::ExpressionRegistry, planner::operator_registry::OperatorRegistry, serde::to_arrow_datatype, - shuffle::{SchemaAlignExec, ShuffleWriterExec}, + shuffle::{SchemaAlignExec, ShuffleWriterDestination, ShuffleWriterExec}, }; -use crate::jvm_bridge::{jni_call, JVMClasses}; +use crate::jvm_bridge::{jni_call, JVMClasses, JavaShufflePartitionPusher, ShufflePartitionPusher}; use arrow::compute::CastOptions; use arrow::datatypes::{DataType, Field, FieldRef, Schema, TimeUnit, DECIMAL128_MAX_PRECISION}; use arrow::ffi_stream::FFI_ArrowArrayStream; @@ -265,6 +265,9 @@ pub struct PhysicalPlanner { /// `ExecutionContext`; see that struct for the propagation rationale. `None` when no driving /// Spark task is available. class_loader: Option>>>, + /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task + /// thread before native planning. Only explicit RSS destinations may use it. + shuffle_partition_pusher: Option>, } impl Default for PhysicalPlanner { @@ -283,6 +286,7 @@ impl PhysicalPlanner { sql_text_pool: vec![], task_context: None, class_loader: None, + shuffle_partition_pusher: None, } } @@ -383,6 +387,18 @@ impl PhysicalPlanner { self } + /// Attach the remote-shuffle callback owned by the Spark task that created this plan. + /// + /// The callback is retained by any RSS shuffle writer in the physical plan and must not be + /// supplied to a legacy or explicitly local shuffle writer. + pub fn with_shuffle_partition_pusher( + mut self, + shuffle_partition_pusher: Option>, + ) -> Self { + self.shuffle_partition_pusher = shuffle_partition_pusher; + self + } + /// Return session context of this planner. pub fn session_ctx(&self) -> &Arc { &self.session_ctx @@ -1835,18 +1851,18 @@ impl PhysicalPlanner { ))), }?; - let (output_data_file, output_index_file) = shuffle_output_files(writer)?; + let destination = + shuffle_writer_destination(writer, self.shuffle_partition_pusher.as_ref())?; let write_buffer_size = writer.write_buffer_size as usize; // Zero on the wire means the limit is disabled; normalize it here so the writer // only ever sees a real limit or none at all. let max_buffer_bytes = (writer.max_buffer_bytes > 0).then_some(writer.max_buffer_bytes as usize); - let shuffle_writer = Arc::new(ShuffleWriterExec::try_new( + let shuffle_writer = Arc::new(ShuffleWriterExec::try_new_with_destination( writer_input, partitioning, codec, - output_data_file.to_string(), - output_index_file.to_string(), + destination, writer.tracing_enabled, write_buffer_size, max_buffer_bytes, @@ -3892,16 +3908,28 @@ fn align_shuffle_writer_input( .map_err(|e| ExecutionError::DataFusionError(e.to_string())) } -/// Resolves local shuffle output paths from the destination-aware writer descriptor. +/// Resolves a native shuffle writer's destination and binds its task-owned callback, if any. /// /// Plans serialized before partition-writer descriptors were introduced only contain the /// top-level paths. New local plans populate both representations so older native binaries can /// still consume them; when both are present they must agree to avoid ambiguous destinations. -fn shuffle_output_files( +/// RSS plans must not contain local output paths and require the callback supplied by their own +/// Spark task. Supplying a remote callback for a local destination is also rejected. +fn shuffle_writer_destination( writer: &spark_operator::ShuffleWriter, -) -> Result<(&str, &str), ExecutionError> { + shuffle_partition_pusher: Option<&Arc>, +) -> Result { let Some(partition_writer) = writer.partition_writer.as_ref() else { - return Ok((&writer.output_data_file, &writer.output_index_file)); + if shuffle_partition_pusher.is_some() { + return Err(GeneralError( + "Local shuffle partition writer cannot use a remote shuffle callback".to_string(), + )); + } + + return Ok(ShuffleWriterDestination::Local { + output_data_file: writer.output_data_file.clone(), + output_index_file: writer.output_index_file.clone(), + }); }; match partition_writer.writer.as_ref() { @@ -3938,13 +3966,37 @@ fn shuffle_output_files( )); } - Ok((&local.output_data_file, &local.output_index_file)) + if shuffle_partition_pusher.is_some() { + return Err(GeneralError( + "Local shuffle partition writer cannot use a remote shuffle callback" + .to_string(), + )); + } + + Ok(ShuffleWriterDestination::Local { + output_data_file: local.output_data_file.clone(), + output_index_file: local.output_index_file.clone(), + }) + } + Some(spark_operator::partition_writer::Writer::Rss(_)) => { + if !writer.output_data_file.is_empty() || !writer.output_index_file.is_empty() { + return Err(GeneralError( + "RSS shuffle partition writer cannot have local output files".to_string(), + )); + } + + let pusher = shuffle_partition_pusher.ok_or_else(|| { + GeneralError( + "RSS shuffle partition writer requires a task-owned remote shuffle callback" + .to_string(), + ) + })?; + + Ok(ShuffleWriterDestination::Rss { + pusher: Arc::clone(pusher), + max_frame_size: JavaShufflePartitionPusher::MAX_PAYLOAD_SIZE, + }) } - Some(spark_operator::partition_writer::Writer::Rss(_)) => Err(GeneralError( - "RSS shuffle partition writers are not supported until remote shuffle execution is \ - enabled" - .to_string(), - )), None => Err(GeneralError( "Shuffle partition writer has no destination".to_string(), )), @@ -4744,7 +4796,13 @@ fn needs_fields_coercion(sig: &TypeSignature) -> bool { #[cfg(test)] mod tests { use futures::{poll, StreamExt}; - use std::{sync::Arc, task::Poll}; + use std::{ + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + task::Poll, + }; use arrow::array::{ Array, DictionaryArray, Int32Array, Int8Array, ListArray, RecordBatch, StringArray, @@ -4765,7 +4823,10 @@ mod tests { use tempfile::TempDir; use tokio::sync::mpsc; - use crate::execution::{operators::InputBatch, planner::PhysicalPlanner}; + use crate::execution::{ + operators::InputBatch, planner::PhysicalPlanner, shuffle::ShuffleWriterDestination, + }; + use crate::jvm_bridge::{JavaShufflePartitionPusher, ShufflePartitionPusher}; use crate::execution::operators::ExecutionError; use crate::execution::planner::literal_to_array_ref; @@ -4783,6 +4844,22 @@ mod tests { }; use datafusion_comet_spark_expr::EvalMode; + #[derive(Default)] + struct RecordingShufflePartitionPusher { + pushes: AtomicUsize, + } + + impl ShufflePartitionPusher for RecordingShufflePartitionPusher { + fn push_partition_data( + &self, + _partition_id: i32, + _data: &[u8], + ) -> datafusion::common::Result<()> { + self.pushes.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + } + fn local_shuffle_partition_writer( output_data_file: &str, output_index_file: &str, @@ -4797,6 +4874,31 @@ mod tests { } } + fn rss_shuffle_partition_writer() -> spark_operator::PartitionWriter { + spark_operator::PartitionWriter { + writer: Some(spark_operator::partition_writer::Writer::Rss( + spark_operator::RssPartitionWriter {}, + )), + } + } + + 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, + } => { + assert_eq!(output_data_file, expected_data_file); + assert_eq!(output_index_file, expected_index_file); + } + destination => panic!("expected a local shuffle destination, got {destination:?}"), + } + } + #[test] fn shuffle_partition_writer_legacy_paths_remain_supported() { let writer = spark_operator::ShuffleWriter { @@ -4805,10 +4907,7 @@ mod tests { ..Default::default() }; - assert_eq!( - super::shuffle_output_files(&writer).unwrap(), - ("legacy.data", "legacy.index") - ); + assert_local_shuffle_destination(&writer, "legacy.data", "legacy.index"); } #[test] @@ -4821,10 +4920,7 @@ mod tests { ..Default::default() }; - assert_eq!( - super::shuffle_output_files(&writer).unwrap(), - ("shuffle.data", "shuffle.index") - ); + assert_local_shuffle_destination(&writer, "shuffle.data", "shuffle.index"); } #[test] @@ -4839,10 +4935,7 @@ mod tests { ..Default::default() }; - assert_eq!( - super::shuffle_output_files(&writer).unwrap(), - ("shuffle.data", "shuffle.index") - ); + assert_local_shuffle_destination(&writer, "shuffle.data", "shuffle.index"); } #[test] @@ -4856,7 +4949,7 @@ mod tests { ..Default::default() }; - let error = super::shuffle_output_files(&writer).unwrap_err(); + let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); assert!( error.to_string().contains("output data file conflicts"), "unexpected error: {error}" @@ -4874,7 +4967,7 @@ mod tests { ..Default::default() }; - let error = super::shuffle_output_files(&writer).unwrap_err(); + let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); assert!( error.to_string().contains("output index file conflicts"), "unexpected error: {error}" @@ -4889,7 +4982,7 @@ mod tests { ..Default::default() }; - let error = super::shuffle_output_files(&writer).unwrap_err(); + let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); assert!( error.to_string().contains("missing its output data file"), "unexpected error: {error}" @@ -4904,7 +4997,7 @@ mod tests { ..Default::default() }; - let error = super::shuffle_output_files(&writer).unwrap_err(); + let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); assert!( error.to_string().contains("missing its output index file"), "unexpected error: {error}" @@ -4918,7 +5011,7 @@ mod tests { ..Default::default() }; - let error = super::shuffle_output_files(&writer).unwrap_err(); + let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); assert!( error.to_string().contains("has no destination"), "unexpected error: {error}" @@ -4926,23 +5019,237 @@ mod tests { } #[test] - fn shuffle_partition_writer_rejects_rss_until_execution_is_supported() { + fn shuffle_partition_writer_requires_a_task_callback_for_rss() { let writer = spark_operator::ShuffleWriter { - partition_writer: Some(spark_operator::PartitionWriter { - writer: Some(spark_operator::partition_writer::Writer::Rss( - spark_operator::RssPartitionWriter {}, - )), - }), + partition_writer: Some(rss_shuffle_partition_writer()), + ..Default::default() + }; + + let error = super::shuffle_writer_destination(&writer, None).unwrap_err(); + assert!( + error + .to_string() + .contains("requires a task-owned remote shuffle callback"), + "unexpected error: {error}" + ); + } + + #[test] + fn shuffle_partition_writer_binds_its_rss_task_callback() { + let writer = spark_operator::ShuffleWriter { + partition_writer: Some(rss_shuffle_partition_writer()), ..Default::default() }; + let callback: Arc = + Arc::new(RecordingShufflePartitionPusher::default()); + + match super::shuffle_writer_destination(&writer, Some(&callback)).unwrap() { + ShuffleWriterDestination::Rss { + pusher, + max_frame_size, + } => { + assert!(Arc::ptr_eq(&pusher, &callback)); + assert_eq!(max_frame_size, JavaShufflePartitionPusher::MAX_PAYLOAD_SIZE); + } + destination => panic!("expected an RSS shuffle destination, got {destination:?}"), + } + } - let error = super::shuffle_output_files(&writer).unwrap_err(); + #[test] + fn shuffle_partition_writer_rejects_rss_with_legacy_data_path() { + let writer = spark_operator::ShuffleWriter { + output_data_file: "legacy.data".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("RSS shuffle partition writers"), + error.to_string().contains("cannot have local output files"), "unexpected error: {error}" ); } + #[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 = + Arc::new(RecordingShufflePartitionPusher::default()); + + let error = super::shuffle_writer_destination(&writer, Some(&callback)).unwrap_err(); + assert!( + error + .to_string() + .contains("cannot use a remote shuffle callback"), + "unexpected error: {error}" + ); + } + + #[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", + )), + ..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 use a remote shuffle callback"), + "unexpected error: {error}" + ); + } + + #[test] + fn shuffle_partition_writer_callbacks_remain_isolated_between_tasks() { + let writer = spark_operator::ShuffleWriter { + partition_writer: Some(rss_shuffle_partition_writer()), + ..Default::default() + }; + let first_recorder = Arc::new(RecordingShufflePartitionPusher::default()); + let second_recorder = Arc::new(RecordingShufflePartitionPusher::default()); + let first_callback: Arc = + Arc::::clone(&first_recorder); + let second_callback: Arc = + Arc::::clone(&second_recorder); + let first_planner = PhysicalPlanner::default() + .with_shuffle_partition_pusher(Some(Arc::clone(&first_callback))); + let second_planner = PhysicalPlanner::default() + .with_shuffle_partition_pusher(Some(Arc::clone(&second_callback))); + + let ShuffleWriterDestination::Rss { + pusher: first_pusher, + .. + } = super::shuffle_writer_destination( + &writer, + first_planner.shuffle_partition_pusher.as_ref(), + ) + .unwrap() + else { + panic!("expected an RSS shuffle destination for the first task"); + }; + let ShuffleWriterDestination::Rss { + pusher: second_pusher, + .. + } = super::shuffle_writer_destination( + &writer, + second_planner.shuffle_partition_pusher.as_ref(), + ) + .unwrap() + else { + panic!("expected an RSS shuffle destination for the second task"); + }; + + assert!(Arc::ptr_eq(&first_pusher, &first_callback)); + assert!(Arc::ptr_eq(&second_pusher, &second_callback)); + assert!(!Arc::ptr_eq(&first_pusher, &second_pusher)); + + first_pusher.push_partition_data(0, b"first").unwrap(); + assert_eq!(first_recorder.pushes.load(Ordering::Relaxed), 1); + assert_eq!(second_recorder.pushes.load(Ordering::Relaxed), 0); + + second_pusher.push_partition_data(0, b"second").unwrap(); + assert_eq!(first_recorder.pushes.load(Ordering::Relaxed), 1); + assert_eq!(second_recorder.pushes.load(Ordering::Relaxed), 1); + } + + #[test] + fn shuffle_partition_writer_plans_and_executes_rss_with_its_task_callback() { + let scan = Operator { + plan_id: 0, + sql_text_pool: vec![], + children: vec![], + op_struct: Some(OpStruct::Scan(spark_operator::Scan { + fields: vec![], + source: "rss-task-input".to_string(), + })), + }; + let shuffle = Operator { + plan_id: 1, + sql_text_pool: vec![], + children: vec![scan], + op_struct: Some(OpStruct::ShuffleWriter(spark_operator::ShuffleWriter { + partitioning: Some(datafusion_comet_proto::spark_partitioning::Partitioning { + partitioning_struct: Some( + datafusion_comet_proto::spark_partitioning::partitioning::PartitioningStruct::SinglePartition( + datafusion_comet_proto::spark_partitioning::SinglePartition {}, + ), + ), + }), + partition_writer: Some(rss_shuffle_partition_writer()), + write_buffer_size: 1024, + ..Default::default() + })), + }; + let recorder = Arc::new(RecordingShufflePartitionPusher::default()); + let callback: Arc = + Arc::::clone(&recorder); + let planner = PhysicalPlanner::default().with_shuffle_partition_pusher(Some(callback)); + let (mut scans, shuffle_scans, plan) = + planner.create_plan(&shuffle, &mut vec![], 1).unwrap(); + + assert!(shuffle_scans.is_empty()); + scans[0].set_input_batch(InputBatch::Batch(vec![], 17)); + let mut stream = plan + .native_plan + .execute(0, planner.session_ctx().task_ctx()) + .unwrap(); + + tokio::runtime::Runtime::new() + .unwrap() + .block_on(async move { + let mut eof_sent = false; + + loop { + match poll!(stream.next()) { + Poll::Ready(Some(result)) => { + let result = result.unwrap(); + panic!("shuffle writer must not produce output batches: {result:?}"); + } + Poll::Ready(None) => break, + Poll::Pending if !eof_sent => { + scans[0].set_input_batch(InputBatch::EOF); + eof_sent = true; + } + Poll::Pending => { + panic!("shuffle writer remained pending after end of input") + } + } + } + }); + + assert_eq!(recorder.pushes.load(Ordering::Relaxed), 1); + } + #[test] fn test_unpack_dictionary_primitive() { let op_scan = Operator { diff --git a/native/jni-bridge/src/shuffle_partition_pusher.rs b/native/jni-bridge/src/shuffle_partition_pusher.rs index b8681ae252..e5269b002a 100644 --- a/native/jni-bridge/src/shuffle_partition_pusher.rs +++ b/native/jni-bridge/src/shuffle_partition_pusher.rs @@ -44,6 +44,12 @@ pub struct JavaShufflePartitionPusher { } impl JavaShufflePartitionPusher { + /// Largest payload that can safely be copied into a JVM byte array. + /// + /// Keep native shuffle frame limits aligned with OpenJDK's conservative + /// soft maximum so oversized frames fail before crossing the JNI boundary. + pub const MAX_PAYLOAD_SIZE: usize = MAX_JVM_ARRAY_LENGTH as usize; + /// Captures the callback while running on an attached JVM thread. pub fn try_new(env: &mut Env<'_>, callback: &JObject<'_>) -> Result { if callback.is_null() { @@ -126,6 +132,10 @@ mod tests { #[test] fn accepts_payload_lengths_up_to_jvm_soft_array_limit() { + assert_eq!( + JavaShufflePartitionPusher::MAX_PAYLOAD_SIZE, + MAX_JVM_ARRAY_LENGTH as usize + ); assert_eq!( JavaShufflePartitionPusher::checked_payload_length(0, 0).unwrap(), 0 diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 57578235f7..0c00352cf0 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -36,6 +36,7 @@ import org.apache.comet.Tracing.withTrace import org.apache.comet.exceptions.CometQueryExecutionException import org.apache.comet.parquet.CometFileKeyUnwrapper import org.apache.comet.serde.Config.ConfigMap +import org.apache.comet.shuffle.ShufflePartitionPusher import org.apache.comet.vector.NativeUtil /** @@ -60,6 +61,8 @@ import org.apache.comet.vector.NativeUtil * The index of the partition. * @param encryptedFilePaths * Paths to encrypted Parquet files that need key unwrapping. + * @param shufflePartitionPusher + * Optional task-owned callback that receives remote shuffle output. */ class CometExecIterator( val id: Long, @@ -72,7 +75,8 @@ class CometExecIterator( broadcastedHadoopConfForEncryption: Option[Broadcast[SerializableConfiguration]] = None, encryptedFilePaths: Seq[String] = Seq.empty, shuffleBlockIterators: Map[Int, CometShuffleBlockIterator] = Map.empty, - taskFilePaths: Seq[String] = Seq.empty) + taskFilePaths: Seq[String] = Seq.empty, + shufflePartitionPusher: Option[ShufflePartitionPusher] = None) extends Iterator[ColumnarBatch] with Logging { @@ -106,7 +110,7 @@ class CometExecIterator( val memoryConfig = CometExecIterator.getMemoryConfig(conf) - nativeLib.createPlan( + val createdPlan = nativeLib.createPlan( id, inputObjects, protobufQueryPlan, @@ -130,6 +134,37 @@ class CometExecIterator( // worker has neither. See CometUdfBridge.evaluate. TaskContext.get(), Thread.currentThread().getContextClassLoader) + + // Bind task-owned callbacks separately to preserve the existing createPlan JNI signature. + try { + shufflePartitionPusher.foreach { pusher => + nativeLib.setShufflePartitionPusher(createdPlan, pusher) + } + createdPlan + } catch { + case failure: Throwable => + // The task-completion listener is not installed until iterator construction succeeds. + try { + nativeUtil.close() + } catch { + case closeFailure: Throwable => failure.addSuppressed(closeFailure) + } + + shuffleBlockIterators.values.foreach { iterator => + try { + iterator.close() + } catch { + case closeFailure: Throwable => failure.addSuppressed(closeFailure) + } + } + + try { + nativeLib.releasePlan(createdPlan) + } catch { + case releaseFailure: Throwable => failure.addSuppressed(releaseFailure) + } + throw failure + } } private var nextBatch: Option[ColumnarBatch] = None diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 51b3e5e41b..630b3b2b7b 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -25,6 +25,7 @@ import org.apache.spark.{CometTaskMemoryManager, TaskContext} import org.apache.spark.sql.comet.CometMetricNode import org.apache.comet.parquet.CometFileKeyUnwrapper +import org.apache.comet.shuffle.ShufflePartitionPusher class Native extends NativeBase { @@ -79,6 +80,19 @@ class Native extends NativeBase { classLoader: ClassLoader): Long // scalastyle:on + /** + * Bind a task-owned remote shuffle callback to an existing native query plan. + * + * The callback must be registered before the plan is executed. Its native global reference is + * released when the plan is released. + * + * @param plan + * the address of the native query plan. + * @param pusher + * the callback that receives complete encoded shuffle blocks. + */ + @native def setShufflePartitionPusher(plan: Long, pusher: ShufflePartitionPusher): Unit + /** * Execute a native query plan based on given input Arrow arrays. * 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 6daa0c10e3..ce7adbc4b9 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -19,6 +19,10 @@ package org.apache.comet.exec +import java.io.{ByteArrayInputStream, IOException} +import java.nio.{ByteBuffer, ByteOrder} + +import scala.collection.mutable import scala.concurrent.duration.DurationInt import scala.util.Random @@ -28,12 +32,17 @@ import org.scalatest.Tag import org.apache.hadoop.fs.Path import org.apache.spark.SparkEnv import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row} +import org.apache.spark.sql.comet.{CometExec, CometMetricNode} +import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions.{col, count, sum} +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.CometConf -import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, Native} +import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass} +import org.apache.comet.shuffle.ShufflePartitionPusher class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper { override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit @@ -50,6 +59,191 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper import testImplicits._ + private def rssShufflePlanBytes: Array[Byte] = { + val scan = OperatorOuterClass.Operator + .newBuilder() + .setScan(OperatorOuterClass.Scan.newBuilder().setSource("RemoteShuffleInput")) + .build() + val partitioning = PartitioningOuterClass.Partitioning + .newBuilder() + .setSinglePartition(PartitioningOuterClass.SinglePartition.getDefaultInstance) + .build() + val writer = OperatorOuterClass.ShuffleWriter + .newBuilder() + .setPartitioning(partitioning) + .setPartitionWriter( + OperatorOuterClass.PartitionWriter + .newBuilder() + .setRss(OperatorOuterClass.RssPartitionWriter.getDefaultInstance) + .build()) + .setWriteBufferSize(1024) + .build() + + OperatorOuterClass.Operator + .newBuilder() + .setShuffleWriter(writer) + .addChildren(scan) + .build() + .toByteArray + } + + test("native shuffle callback registration preserves the existing createPlan JNI signature") { + val createPlan = classOf[Native].getDeclaredMethods.find(_.getName == "createPlan").get + assert(createPlan.getParameterTypes.last == classOf[ClassLoader]) + + val registration = classOf[Native] + .getDeclaredMethod( + "setShufflePartitionPusher", + java.lang.Long.TYPE, + classOf[ShufflePartitionPusher]) + assert(registration.getReturnType == java.lang.Void.TYPE) + } + + test("native RSS shuffle invokes its task-owned JVM callback with a complete frame") { + val planBytes = rssShufflePlanBytes + + val frames = spark.sparkContext + .parallelize(Seq(17), 1) + .mapPartitions { rowCounts => + val batch = new ColumnarBatch(Array.empty[ColumnVector], rowCounts.next()) + val inputs = CometArrowStream.inputObjects( + Iterator.single(batch), + StructType(Nil), + "native-rss-callback-test") + val captured = mutable.ArrayBuffer.empty[(Int, Int, Int, Long)] + val pusher = new ShufflePartitionPusher { + override def pushPartitionData(partitionId: Int, data: Array[Byte], length: Int) + : Unit = { + val encodedLength = ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN).getLong + captured.synchronized { + captured += ((partitionId, data.length, length, encodedLength)) + } + } + } + val iterator = new CometExecIterator( + CometExec.newIterId, + inputs, + 0, + planBytes, + CometMetricNode(Map.empty), + 1, + 0, + shufflePartitionPusher = Some(pusher)) + + try { + while (iterator.hasNext) { + iterator.next().close() + } + captured.iterator + } finally { + iterator.close() + } + } + .collect() + + assert(frames.length == 1) + val (partitionId, arrayLength, callbackLength, encodedLength) = frames.head + assert(partitionId == 0) + assert(arrayLength > java.lang.Long.BYTES) + assert(callbackLength == arrayLength) + assert(encodedLength + java.lang.Long.BYTES == arrayLength) + } + + test("native RSS shuffle preserves the original JVM callback exception") { + val planBytes = rssShufflePlanBytes + + val failures = spark.sparkContext + .parallelize(Seq(17), 1) + .mapPartitions { rowCounts => + val batch = new ColumnarBatch(Array.empty[ColumnVector], rowCounts.next()) + val inputs = CometArrowStream.inputObjects( + Iterator.single(batch), + StructType(Nil), + "native-rss-callback-failure-test") + val expected = new IOException("remote shuffle callback sentinel") + val pusher = new ShufflePartitionPusher { + override def pushPartitionData(partitionId: Int, data: Array[Byte], length: Int): Unit = + throw expected + } + val iterator = new CometExecIterator( + CometExec.newIterId, + inputs, + 0, + planBytes, + CometMetricNode(Map.empty), + 1, + 0, + shufflePartitionPusher = Some(pusher)) + + try { + try { + iterator.hasNext + Iterator.single((false, "callback unexpectedly succeeded")) + } catch { + case actual: IOException => Iterator.single((actual eq expected, actual.getMessage)) + } + } finally { + iterator.close() + } + } + .collect() + + assert(failures.sameElements(Array((true, "remote shuffle callback sentinel")))) + } + + test("failed RSS callback registration closes every shuffle input and preserves its error") { + val planBytes = rssShufflePlanBytes + + val results = spark.sparkContext + .parallelize(Seq(17), 1) + .mapPartitions { rowCounts => + val batch = new ColumnarBatch(Array.empty[ColumnVector], rowCounts.next()) + val inputs = CometArrowStream.inputObjects( + Iterator.single(batch), + StructType(Nil), + "native-rss-registration-failure-test") + val closedInputs = mutable.ArrayBuffer.empty[String] + val failingInput = new ByteArrayInputStream(Array.empty[Byte]) { + override def close(): Unit = { + closedInputs += "failing" + throw new IOException("shuffle input cleanup sentinel") + } + } + val remainingInput = new ByteArrayInputStream(Array.empty[Byte]) { + override def close(): Unit = closedInputs += "remaining" + } + val shuffleInputs = Map( + 0 -> new CometShuffleBlockIterator(failingInput), + 1 -> new CometShuffleBlockIterator(remainingInput)) + + try { + val iterator = new CometExecIterator( + CometExec.newIterId, + inputs, + 0, + planBytes, + CometMetricNode(Map.empty), + 1, + 0, + shuffleBlockIterators = shuffleInputs, + shufflePartitionPusher = Some(null)) + iterator.close() + Iterator.single((false, false, false)) + } catch { + case failure: Throwable => + Iterator.single( + ( + failure.getMessage.contains("Remote shuffle callback must not be null"), + closedInputs.toSet == Set("failing", "remaining"), + failure.getSuppressed.exists( + _.getMessage.contains("shuffle input cleanup sentinel")))) + } + } + .collect() + + assert(results.sameElements(Array((true, true, true)))) + } + test("native shuffle plan preserves local partition writer and legacy output paths") { val dataFile = "/tmp/comet-shuffle.data" val indexFile = "/tmp/comet-shuffle.index" From 69e3cfb8239806648b79b24148607be16504e378 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Thu, 27 Aug 2026 03:28:27 +0000 Subject: [PATCH 2/2] fix: release Arrow inputs after RSS callback registration failure --- .../org/apache/comet/CometExecIterator.scala | 12 ++++++ .../comet/exec/CometNativeShuffleSuite.scala | 41 +++++++++++++++---- 2 files changed, 44 insertions(+), 9 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 0c00352cf0..17d78dfee7 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -21,6 +21,7 @@ package org.apache.comet import java.lang.management.ManagementFactory +import org.apache.arrow.c.ArrowArrayStream import org.apache.hadoop.conf.Configuration import org.apache.spark._ import org.apache.spark.broadcast.Broadcast @@ -150,6 +151,17 @@ class CometExecIterator( case closeFailure: Throwable => failure.addSuppressed(closeFailure) } + // Native only takes ownership of Arrow streams during the first executePlan call. + inputObjects.foreach { + case stream: ArrowArrayStream => + try { + stream.release() + } catch { + case releaseFailure: Throwable => failure.addSuppressed(releaseFailure) + } + case _ => + } + shuffleBlockIterators.values.foreach { iterator => try { iterator.close() 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 ce7adbc4b9..ee76189bd8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -29,6 +29,9 @@ import scala.util.Random import org.scalactic.source.Position import org.scalatest.Tag +import org.apache.arrow.memory.ArrowBuf +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.{Field, Schema} import org.apache.hadoop.fs.Path import org.apache.spark.SparkEnv import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row} @@ -191,17 +194,36 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper assert(failures.sameElements(Array((true, "remote shuffle callback sentinel")))) } - test("failed RSS callback registration closes every shuffle input and preserves its error") { + test( + "failed RSS callback registration closes every Arrow and shuffle input and preserves its error") { val planBytes = rssShufflePlanBytes val results = spark.sparkContext .parallelize(Seq(17), 1) - .mapPartitions { rowCounts => - val batch = new ColumnarBatch(Array.empty[ColumnVector], rowCounts.next()) - val inputs = CometArrowStream.inputObjects( - Iterator.single(batch), - StructType(Nil), - "native-rss-registration-failure-test") + .mapPartitions { _ => + var readerClosed = false + var ownedBuffer: ArrowBuf = null + val arrowInput = CometArrowStream + .stream( + "native-rss-registration-failure-test", + allocator => { + ownedBuffer = allocator.buffer(128) + new ArrowReader(allocator) { + override protected def readSchema(): Schema = + new Schema(java.util.Collections.emptyList[Field]()) + + override def loadNextBatch(): Boolean = false + + override def bytesRead(): Long = 0L + + override protected def closeReadSource(): Unit = { + ownedBuffer.close() + readerClosed = true + } + } + }) + .next() + val inputs = Array[Object](arrowInput) val closedInputs = mutable.ArrayBuffer.empty[String] val failingInput = new ByteArrayInputStream(Array.empty[Byte]) { override def close(): Unit = { @@ -228,12 +250,13 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper shuffleBlockIterators = shuffleInputs, shufflePartitionPusher = Some(null)) iterator.close() - Iterator.single((false, false, false)) + Iterator.single((false, false, false, false)) } catch { case failure: Throwable => Iterator.single( ( failure.getMessage.contains("Remote shuffle callback must not be null"), + readerClosed && ownedBuffer.refCnt() == 0, closedInputs.toSet == Set("failing", "remaining"), failure.getSuppressed.exists( _.getMessage.contains("shuffle input cleanup sentinel")))) @@ -241,7 +264,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } .collect() - assert(results.sameElements(Array((true, true, true)))) + assert(results.sameElements(Array((true, true, true, true)))) } test("native shuffle plan preserves local partition writer and legacy output paths") {