From f3706cd921a260e61ff05f13308c16049c5bf36b Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:16:26 +0000 Subject: [PATCH 01/12] feat: add RSS partition writer and JNI callback --- native/Cargo.lock | 1 + native/jni-bridge/Cargo.toml | 1 + native/jni-bridge/src/lib.rs | 1 + .../src/shuffle_partition_pusher.rs | 159 ++++++ .../java/RecordingShufflePartitionPusher.java | 43 ++ .../tests/shuffle_partition_pusher.rs | 204 ++++++++ native/shuffle/src/lib.rs | 2 +- native/shuffle/src/writers/mod.rs | 2 + .../shuffle/src/writers/partition_writer.rs | 17 +- native/shuffle/src/writers/rss/mod.rs | 20 + .../src/writers/rss/rss_partition_writer.rs | 462 ++++++++++++++++++ .../comet/shuffle/ShufflePartitionPusher.java | 43 ++ 12 files changed, 946 insertions(+), 9 deletions(-) create mode 100644 native/jni-bridge/src/shuffle_partition_pusher.rs create mode 100644 native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java create mode 100644 native/jni-bridge/tests/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/Cargo.lock b/native/Cargo.lock index eadd9d995c5..2aa88ee591c 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2024,6 +2024,7 @@ dependencies = [ "paste", "prost", "regex", + "tempfile", "thiserror 2.0.20", ] diff --git a/native/jni-bridge/Cargo.toml b/native/jni-bridge/Cargo.toml index 87e87bcc2be..18ecadf4071 100644 --- a/native/jni-bridge/Cargo.toml +++ b/native/jni-bridge/Cargo.toml @@ -44,3 +44,4 @@ datafusion-comet-common = { workspace = true } [dev-dependencies] jni = { version = "0.22.4", features = ["invocation"] } assertables = "10" +tempfile = "3.26.0" diff --git a/native/jni-bridge/src/lib.rs b/native/jni-bridge/src/lib.rs index 8db4c07851a..96487ac7bf4 100644 --- a/native/jni-bridge/src/lib.rs +++ b/native/jni-bridge/src/lib.rs @@ -33,6 +33,7 @@ use once_cell::sync::OnceCell; use errors::{CometError, CometResult}; pub mod errors; +pub mod shuffle_partition_pusher; enum LocalFrameError { Closure(E), 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..d5bf744acfe --- /dev/null +++ b/native/jni-bridge/src/shuffle_partition_pusher.rs @@ -0,0 +1,159 @@ +// 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 std::sync::Arc; + +use jni::{ + errors::Error as JniError, + objects::{Global, JClass, JMethodID, JObject, JValue}, + signature::{Primitive, ReturnType}, + Env, JavaVM, +}; + +use crate::errors::{CometError, CometResult}; + +/// Task-owned JNI callback. It is resolved only when explicitly constructed, so older JVM +/// artifacts without the RSS interface can still execute existing local Comet plans. +#[derive(Clone, Debug)] +pub struct JavaShufflePartitionPusher { + inner: Arc, + num_partitions: usize, + max_push_bytes: usize, +} + +#[derive(Debug)] +struct Callback { + vm: JavaVM, + object: Global>, + // The class must stay alive for as long as the cached method ID. + _class: Global>, + push_partition_data: JMethodID, +} + +impl JavaShufflePartitionPusher { + pub const JVM_CLASS: &'static str = "org/apache/comet/shuffle/ShufflePartitionPusher"; + + pub fn try_new( + env: &mut Env, + object: &JObject, + num_partitions: usize, + max_push_bytes: usize, + ) -> CometResult { + if object.is_null() { + return Err(CometError::NullPointer( + "ShufflePartitionPusher is null".into(), + )); + } + if num_partitions == 0 || num_partitions > i32::MAX as usize { + return Err(CometError::Config("Invalid RSS partition count".into())); + } + if max_push_bytes == 0 || max_push_bytes > i32::MAX as usize { + return Err(CometError::Config("Invalid RSS push byte limit".into())); + } + + let class = env.find_class(jni::strings::JNIString::new(Self::JVM_CLASS))?; + if !env.is_instance_of(object, &class)? { + return Err(CometError::Config( + "Object does not implement ShufflePartitionPusher".into(), + )); + } + let push_partition_data = env.get_method_id( + &class, + jni::jni_str!("pushPartitionData"), + jni::jni_sig!("(I[BI)I"), + )?; + Ok(Self { + inner: Arc::new(Callback { + vm: env.get_java_vm()?, + object: env.new_global_ref(object)?, + _class: env.new_global_ref(&class)?, + push_partition_data, + }), + num_partitions, + max_push_bytes, + }) + } + + /// Copies a caller-validated complete frame into Java-owned memory. + /// + /// The configured bound is checked before allocating the Java array. This only bounds this + /// synchronous copy; the backend remains responsible for asynchronous admission and commit. + pub fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> CometResult<()> { + if partition_id >= self.num_partitions { + return Err(CometError::Config(format!( + "RSS partition {partition_id} is outside 0..{}", + self.num_partitions + ))); + } + if frame.is_empty() || frame.len() > self.max_push_bytes { + return Err(CometError::Config(format!( + "RSS push size {} is outside 1..={}", + frame.len(), + self.max_push_bytes + ))); + } + let partition_id = i32::try_from(partition_id) + .map_err(|_| CometError::Config("RSS partition ID exceeds jint".into()))?; + let length = i32::try_from(frame.len()) + .map_err(|_| CometError::Config("RSS push size exceeds jint".into()))?; + + // attach_current_thread creates a local frame and captures/clears Java exceptions. Do not + // stringify the captured throwable: the outer Comet JNI boundary can rethrow it unchanged. + let accepted = self + .inner + .vm + .attach_current_thread(|env| -> jni::errors::Result { + let bytes = env.byte_array_from_slice(frame)?; + let args = [ + JValue::Int(partition_id).as_jni(), + JValue::Object(&bytes).as_jni(), + JValue::Int(length).as_jni(), + ]; + // SAFETY: try_new checked the interface and cached its exact (I[BI)I method. + // The object and declaring class are held by global references. + unsafe { + env.call_method_unchecked( + &self.inner.object, + self.inner.push_partition_data, + ReturnType::Primitive(Primitive::Int), + &args, + )? + .i() + } + }) + .map_err(|source| match source { + JniError::CaughtJavaException { + exception, + name, + msg, + .. + } => CometError::JavaException { + class: name, + msg, + throwable: exception, + }, + source => CometError::JNI { source }, + })?; + + if accepted != length { + return Err(CometError::Internal(format!( + "RSS callback accepted {accepted} bytes; expected {length}" + ))); + } + Ok(()) + } +} diff --git a/native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java b/native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java new file mode 100644 index 00000000000..85427a571d9 --- /dev/null +++ b/native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java @@ -0,0 +1,43 @@ +/* + * 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; +import java.util.Arrays; + +public final class RecordingShufflePartitionPusher implements ShufflePartitionPusher { + public int calls; + public int partitionId; + public int adjustment; + public int failureMode; + public byte[] lastBytes; + public final IOException failure = new IOException("recorded push failure"); + + @Override + public int pushPartitionData(int partitionId, byte[] bytes, int length) throws IOException { + calls++; + if (failureMode != 0) { + throw failure; + } + this.partitionId = partitionId; + lastBytes = Arrays.copyOf(bytes, length); + return length + adjustment; + } +} diff --git a/native/jni-bridge/tests/shuffle_partition_pusher.rs b/native/jni-bridge/tests/shuffle_partition_pusher.rs new file mode 100644 index 00000000000..9a7b3a38f76 --- /dev/null +++ b/native/jni-bridge/tests/shuffle_partition_pusher.rs @@ -0,0 +1,204 @@ +// 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. + +// JNI errors retain Java exception information; match the bridge crate's lint policy. +#![allow(clippy::result_large_err)] + +use std::path::PathBuf; +use std::process::Command; +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::ipc::writer::StreamWriter; +use datafusion_comet_jni_bridge::errors::{CometError, CometResult}; +use datafusion_comet_jni_bridge::shuffle_partition_pusher::JavaShufflePartitionPusher; +use jni::objects::{JByteArray, JObject, JValue}; +use jni::{InitArgsBuilder, JNIVersion, JavaVM}; + +fn complete_frame() -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![1, 2]))], + ) + .unwrap(); + let mut ipc = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, schema.as_ref()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + let mut frame = Vec::new(); + frame.extend_from_slice(&(12 + ipc.len() as u64).to_le_bytes()); + frame.extend_from_slice(&1_u64.to_le_bytes()); + frame.extend_from_slice(b"NONE"); + frame.extend_from_slice(&ipc); + frame +} + +#[test] +#[cfg_attr(miri, ignore)] // Miri cannot launch a JVM. +fn rss_jni_callback_contract() { + let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let classes = tempfile::tempdir().unwrap(); + let javac = std::env::var_os("JAVA_HOME") + .map(|home| PathBuf::from(home).join("bin/javac")) + .unwrap_or_else(|| PathBuf::from("javac")); + let output = + Command::new(javac) + .arg("-d") + .arg(classes.path()) + .arg(manifest.join( + "../../spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java", + )) + .arg(manifest.join("tests/java/RecordingShufflePartitionPusher.java")) + .output() + .expect("run javac for the RSS callback fixture"); + assert!( + output.status.success(), + "javac: {}", + String::from_utf8_lossy(&output.stderr) + ); + + let args = InitArgsBuilder::new() + .version(JNIVersion::V1_8) + .option("-Xcheck:jni") + .option(format!("-Djava.class.path={}", classes.path().display())) + .build() + .unwrap(); + let vm = JavaVM::new(args).unwrap(); + let frame = complete_frame(); + let (pusher, fixture) = vm + .attach_current_thread(|env| -> CometResult<_> { + let object = env.new_object( + jni::jni_str!("org/apache/comet/shuffle/RecordingShufflePartitionPusher"), + jni::jni_sig!("()V"), + &[], + )?; + assert!(matches!( + JavaShufflePartitionPusher::try_new(env, &JObject::null(), 2, frame.len()), + Err(CometError::NullPointer(_)) + )); + for (partitions, limit) in [ + (0, frame.len()), + (i32::MAX as usize + 1, frame.len()), + (2, 0), + (2, i32::MAX as usize + 1), + ] { + assert!(matches!( + JavaShufflePartitionPusher::try_new(env, &object, partitions, limit), + Err(CometError::Config(_)) + )); + } + let wrong_type = + env.new_object(jni::jni_str!("java/lang/Object"), jni::jni_sig!("()V"), &[])?; + assert!(matches!( + JavaShufflePartitionPusher::try_new(env, &wrong_type, 2, frame.len()), + Err(CometError::Config(_)) + )); + let pusher = JavaShufflePartitionPusher::try_new(env, &object, 2, frame.len())?; + Ok((pusher, env.new_global_ref(&object)?)) + }) + .unwrap(); + + // Global references and the cached method remain valid on another attached native thread. + let worker_pusher = pusher.clone(); + let worker_frame = frame.clone(); + std::thread::spawn(move || worker_pusher.push_partition_data(1, &worker_frame)) + .join() + .unwrap() + .unwrap(); + + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + assert_eq!( + env.get_field(&fixture, jni::jni_str!("calls"), jni::jni_sig!("I"))? + .i()?, + 1 + ); + assert_eq!( + env.get_field(&fixture, jni::jni_str!("partitionId"), jni::jni_sig!("I"))? + .i()?, + 1 + ); + let bytes = env + .get_field(&fixture, jni::jni_str!("lastBytes"), jni::jni_sig!("[B"))? + .l()?; + // SAFETY: the field descriptor is byte[] and into_raw transfers this local reference. + let bytes = unsafe { JByteArray::from_raw(env, bytes.into_raw()) }; + assert_eq!(env.convert_byte_array(&bytes)?, frame); + Ok(()) + }) + .unwrap(); + + assert!(pusher.push_partition_data(2, &frame).is_err()); + assert!(pusher.push_partition_data(0, &[]).is_err()); + assert!(pusher + .push_partition_data(0, &vec![0; frame.len() + 1]) + .is_err()); + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + assert_eq!( + env.get_field(&fixture, jni::jni_str!("calls"), jni::jni_sig!("I"))? + .i()?, + 1 + ); + Ok(()) + }) + .unwrap(); + + for adjustment in [-(frame.len() as i32), -1, 1] { + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + env.set_field( + &fixture, + jni::jni_str!("adjustment"), + jni::jni_sig!("I"), + JValue::Int(adjustment), + ) + }) + .unwrap(); + assert!(matches!( + pusher.push_partition_data(0, &frame), + Err(CometError::Internal(_)) + )); + } + + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + env.set_field( + &fixture, + jni::jni_str!("failureMode"), + jni::jni_sig!("I"), + JValue::Int(1), + ) + }) + .unwrap(); + let error = pusher.push_partition_data(0, &frame).unwrap_err(); + let CometError::JavaException { throwable, .. } = error else { + panic!("original Java throwable was not preserved: {error:?}"); + }; + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + let expected = env + .get_field( + &fixture, + jni::jni_str!("failure"), + jni::jni_sig!("Ljava/io/IOException;"), + )? + .l()?; + assert!(env.is_same_object(&expected, &throwable)?); + Ok(()) + }) + .unwrap(); +} diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 2263ae0dac0..a7b43fabbaa 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, PartitionPusher, RssPartitionWriter, ShuffleBlockWriter}; diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index 6d330fd12aa..d92453bdc49 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::{PartitionPusher, RssPartitionWriter}; pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/writers/partition_writer.rs b/native/shuffle/src/writers/partition_writer.rs index 25b0e598df8..4bcb60f457b 100644 --- a/native/shuffle/src/writers/partition_writer.rs +++ b/native/shuffle/src/writers/partition_writer.rs @@ -22,9 +22,8 @@ use arrow::record_batch::RecordBatch; /// /// Decouples partitioning from storage: partitioners only produce partitioned /// `RecordBatch` streams, while implementations of this trait own how those -/// batches are stored and finalized. [`LocalPartitionWriter`] implements the -/// local file behavior; other backends (e.g. a remote shuffle writer) can be -/// added without changing the partitioners. +/// batches are stored and finalized. [`LocalPartitionWriter`] implements local file output; +/// [`RssPartitionWriter`] sends encoded frames through a remote shuffle callback. /// /// A partitioner drives a writer as: any number of /// [`write`](PartitionWriter::write) calls to stage batches, then one @@ -32,13 +31,15 @@ use arrow::record_batch::RecordBatch; /// ascending id order, then a single [`finish_all`](PartitionWriter::finish_all). /// /// [`LocalPartitionWriter`]: crate::writers::local::local_partition_writer::LocalPartitionWriter +/// [`RssPartitionWriter`]: crate::RssPartitionWriter pub(crate) trait PartitionWriter: Send + Sync { /// Stages the batches from `iter` for partition `pid` without finalizing it. /// /// Used to stream single-partition output and to stage multi-partition /// spilled batches. A partition may be written multiple times and in any - /// order; staged data is only guaranteed visible after - /// [`finish_partition`](PartitionWriter::finish_partition). + /// order; an implementation may retain staged data until + /// [`finish_partition`](PartitionWriter::finish_partition). Writer-side finalization does + /// not replace a remote backend's task-level commit protocol. fn write( &mut self, pid: usize, @@ -63,9 +64,9 @@ pub(crate) trait PartitionWriter: Send + Sync { where I: Iterator>; - /// Completes the shuffle write, flushing output and emitting the partition - /// index. Called exactly once, after the last - /// [`finish_partition`](PartitionWriter::finish_partition). + /// Completes writer-side output. Local writers flush and emit the partition index; a remote + /// writer's owner must separately complete its task-level commit protocol. Called exactly + /// once, after the last [`finish_partition`](PartitionWriter::finish_partition). fn finish_all(&mut self, metrics: &ShufflePartitionerMetrics) -> datafusion::common::Result<()>; } diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs new file mode 100644 index 00000000000..f20630258bd --- /dev/null +++ b/native/shuffle/src/writers/rss/mod.rs @@ -0,0 +1,20 @@ +// 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. + +mod rss_partition_writer; + +pub use rss_partition_writer::{PartitionPusher, RssPartitionWriter}; 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..b4feaef5199 --- /dev/null +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -0,0 +1,462 @@ +// 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 std::io::{self, Cursor, Seek, SeekFrom, Write}; + +use arrow::record_batch::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use datafusion::physical_plan::metrics::Time; +use datafusion_comet_jni_bridge::shuffle_partition_pusher::JavaShufflePartitionPusher; + +use crate::metrics::ShufflePartitionerMetrics; +use crate::writers::PartitionWriter; +use crate::ShuffleBlockWriter; + +/// Backend-neutral, all-or-error transport for one complete Comet frame. +/// +/// Acceptance is not a remote map commit. Implementations own asynchronous admission, retry, +/// cancellation, and commit; they must not split a frame into independently interleavable pushes. +pub trait PartitionPusher: Send + Sync { + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()>; +} + +impl PartitionPusher for JavaShufflePartitionPusher { + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { + JavaShufflePartitionPusher::push_partition_data(self, partition_id, frame) + .map_err(Into::into) + } +} + +/// Encodes already-partitioned batches using the existing Comet format and sends one frame per +/// callback. This foundation is not selected by the native planner yet. +/// +/// There are no retained per-reducer buffers. The byte limit caps encoded output, not Arrow's +/// encoding scratch space or the backend's asynchronous memory. Production admission, row +/// splitting, integrity, and map-commit handling must be supplied before enabling remote plans. +pub struct RssPartitionWriter { + pusher: P, + block_writer: ShuffleBlockWriter, + num_partitions: usize, + max_frame_bytes: usize, + next_partition: usize, + finished: bool, + failed: bool, +} + +impl RssPartitionWriter

{ + pub fn try_new( + pusher: P, + block_writer: ShuffleBlockWriter, + num_partitions: usize, + max_frame_bytes: usize, + ) -> Result { + if num_partitions == 0 || num_partitions > i32::MAX as usize { + return Err(DataFusionError::Configuration( + "Invalid RSS partition count".into(), + )); + } + if !(20..=i32::MAX as usize).contains(&max_frame_bytes) { + return Err(DataFusionError::Configuration( + "Invalid RSS frame byte limit".into(), + )); + } + Ok(Self { + pusher, + block_writer, + num_partitions, + max_frame_bytes, + next_partition: 0, + finished: false, + failed: false, + }) + } + + /// Encodes and synchronously submits one batch. A failed writer cannot be reused. + pub fn write_batch( + &mut self, + partition_id: usize, + batch: &RecordBatch, + encode_time: &Time, + write_time: &Time, + ) -> Result<()> { + let result = (|| { + self.check_partition(partition_id)?; + let mut output = BoundedBuffer::new(self.max_frame_bytes); + self.block_writer + .write_batch(batch, &mut output, encode_time)?; + let frame = output.inner.get_ref(); + if !frame.is_empty() { + let _timer = write_time.timer(); + self.pusher.push_partition_data(partition_id, frame)?; + } + Ok(()) + })(); + if result.is_err() { + self.failed = true; + } + result + } + + fn check_partition(&self, partition_id: usize) -> Result<()> { + if self.failed || self.finished { + return Err(DataFusionError::Execution("RSS writer is closed".into())); + } + if partition_id >= self.num_partitions || partition_id < self.next_partition { + return Err(DataFusionError::Execution(format!( + "RSS partition {partition_id} is invalid or already finalized" + ))); + } + Ok(()) + } + + fn write_batches( + &mut self, + partition_id: usize, + iter: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + let result = (|| { + self.check_partition(partition_id)?; + for batch in iter { + self.write_batch( + partition_id, + &batch?, + &metrics.encode_time, + &metrics.write_time, + )?; + } + Ok(()) + })(); + if result.is_err() { + // Earlier batches may already have been accepted. An input error must not allow + // this attempt to be resumed or finalized as if its output were complete. + self.failed = true; + } + result + } +} + +impl PartitionWriter for RssPartitionWriter

{ + fn write( + &mut self, + pid: usize, + iter: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + self.write_batches(pid, iter, metrics) + } + + fn finish_partition( + &mut self, + pid: usize, + iter: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + if pid != self.next_partition { + return Err(DataFusionError::Execution(format!( + "Expected RSS partition {}, got {pid}", + self.next_partition + ))); + } + self.write_batches(pid, iter, metrics)?; + self.next_partition += 1; + Ok(()) + } + + fn finish_all(&mut self, _metrics: &ShufflePartitionerMetrics) -> Result<()> { + if self.failed || self.finished || self.next_partition != self.num_partitions { + return Err(DataFusionError::Execution( + "RSS partitions are not ready for finalization".into(), + )); + } + // Remote map commit belongs to the task-owned JVM adapter, not this encoder. + self.finished = true; + Ok(()) + } +} + +/// A seekable encoder destination that rejects an oversized frame before growing its buffer. +struct BoundedBuffer { + inner: Cursor>, + limit: usize, +} + +impl BoundedBuffer { + fn new(limit: usize) -> Self { + Self { + inner: Cursor::new(Vec::new()), + limit, + } + } + + fn limit_error() -> io::Error { + io::Error::new( + io::ErrorKind::InvalidData, + "RSS frame exceeds its byte limit", + ) + } +} + +impl Write for BoundedBuffer { + fn write(&mut self, bytes: &[u8]) -> io::Result { + let end = self.inner.position().checked_add(bytes.len() as u64); + if end.is_none_or(|end| end > self.limit as u64) { + return Err(Self::limit_error()); + } + self.inner.write(bytes) + } + + fn flush(&mut self) -> io::Result<()> { + self.inner.flush() + } +} + +impl Seek for BoundedBuffer { + fn seek(&mut self, position: SeekFrom) -> io::Result { + let previous = self.inner.position(); + let next = self.inner.seek(position)?; + if next > self.limit as u64 { + self.inner.set_position(previous); + return Err(Self::limit_error()); + } + Ok(next) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + + use arrow::array::Int64Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + + use crate::{read_ipc_compressed, CompressionCodec}; + + type CapturedFrames = Arc)>>>; + + #[derive(Clone, Default)] + struct RecordingPusher(CapturedFrames); + + impl PartitionPusher for RecordingPusher { + fn push_partition_data(&self, pid: usize, frame: &[u8]) -> Result<()> { + self.0.lock().unwrap().push((pid, frame.to_vec())); + Ok(()) + } + } + + fn batch(values: Vec) -> RecordBatch { + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])), + vec![Arc::new(Int64Array::from(values))], + ) + .unwrap() + } + + fn metrics() -> ShufflePartitionerMetrics { + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0) + } + + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call ZSTD_createCCtx. + fn rss_partition_writer_routes_complete_frames() { + let first = batch(vec![1, 2]); + let second = batch(vec![3]); + for codec in [ + CompressionCodec::None, + CompressionCodec::Lz4Frame, + CompressionCodec::Snappy, + CompressionCodec::Zstd(1), + ] { + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = ShuffleBlockWriter::try_new(first.schema().as_ref(), codec).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 3, 1024 * 1024).unwrap(); + let metrics = metrics(); + writer + .write(1, &mut [Ok(second.clone())].into_iter(), &metrics) + .unwrap(); + writer + .finish_partition(0, &mut [Ok(first.clone())].into_iter(), &metrics) + .unwrap(); + writer + .finish_partition(1, &mut std::iter::empty(), &metrics) + .unwrap(); + writer + .finish_partition(2, &mut std::iter::empty(), &metrics) + .unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = captured.lock().unwrap(); + assert_eq!(frames.len(), 2); + for ((pid, frame), (expected_pid, expected)) in + frames.iter().zip([(1, &second), (0, &first)]) + { + assert_eq!(*pid, expected_pid); + let declared = u64::from_le_bytes(frame[..8].try_into().unwrap()); + assert_eq!(declared + 8, frame.len() as u64); + assert_eq!(u64::from_le_bytes(frame[8..16].try_into().unwrap()), 1); + assert_eq!(&read_ipc_compressed(&frame[16..]).unwrap(), expected); + } + } + } + + #[test] + fn rss_partition_writer_checks_lifecycle_and_empty_batches() { + let input = batch(vec![]); + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + for (partitions, limit) in [ + (0, 1024), + (i32::MAX as usize + 1, 1024), + (2, 19), + (2, i32::MAX as usize + 1), + ] { + assert!(RssPartitionWriter::try_new( + pusher.clone(), + encoder.clone(), + partitions, + limit, + ) + .is_err()); + } + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 2, 1024).unwrap(); + let metrics = metrics(); + assert!(writer.finish_all(&metrics).is_err()); + assert!(writer + .finish_partition(1, &mut std::iter::empty(), &metrics) + .is_err()); + writer + .finish_partition(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap(); + assert!(writer + .finish_partition(0, &mut std::iter::empty(), &metrics) + .is_err()); + writer + .finish_partition(1, &mut std::iter::empty(), &metrics) + .unwrap(); + writer.finish_all(&metrics).unwrap(); + assert!(writer.finish_all(&metrics).is_err()); + assert!(captured.lock().unwrap().is_empty()); + } + + #[test] + fn rss_partition_writer_rejects_oversized_frames_before_push() { + let input = batch(vec![1]); + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, 20).unwrap(); + let metrics = metrics(); + let error = writer + .write(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap_err(); + assert!(error + .to_string() + .contains("RSS frame exceeds its byte limit")); + assert!(captured.lock().unwrap().is_empty()); + assert!(writer + .finish_partition(0, &mut std::iter::empty(), &metrics) + .is_err()); + } + + #[test] + fn rss_partition_writer_preserves_transport_errors() { + struct FailingPusher; + impl PartitionPusher for FailingPusher { + fn push_partition_data(&self, _pid: usize, _frame: &[u8]) -> Result<()> { + Err(DataFusionError::External(Box::new(io::Error::other( + "push failed", + )))) + } + } + let input = batch(vec![1]); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(FailingPusher, encoder, 1, 1024).unwrap(); + let metrics = metrics(); + let error = writer + .write(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap_err(); + let DataFusionError::External(source) = error else { + panic!("transport error was wrapped or stringified"); + }; + assert_eq!( + source.downcast_ref::().unwrap().to_string(), + "push failed" + ); + } + + #[test] + fn rss_partition_writer_stays_failed_after_input_error() { + let input = batch(vec![1]); + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, 1024).unwrap(); + let metrics = metrics(); + let error = DataFusionError::External(Box::new(io::Error::other("input failed"))); + let error = writer + .write( + 0, + &mut [Ok(input.clone()), Err(error)].into_iter(), + &metrics, + ) + .unwrap_err(); + let DataFusionError::External(source) = error else { + panic!("input error was wrapped or stringified"); + }; + assert_eq!( + source.downcast_ref::().unwrap().to_string(), + "input failed" + ); + assert_eq!(captured.lock().unwrap().len(), 1); + assert!(writer + .write(0, &mut [Ok(input)].into_iter(), &metrics) + .is_err()); + assert!(writer + .finish_partition(0, &mut std::iter::empty(), &metrics) + .is_err()); + assert!(writer.finish_all(&metrics).is_err()); + assert_eq!(captured.lock().unwrap().len(), 1); + } + + #[test] + fn rss_partition_writer_buffer_bounds_seek_and_write() { + let mut buffer = BoundedBuffer::new(4); + buffer.write_all(&[1, 2, 3, 4]).unwrap(); + assert!(buffer.write_all(&[5]).is_err()); + assert!(buffer.seek(SeekFrom::Start(5)).is_err()); + buffer.seek(SeekFrom::Start(0)).unwrap(); + buffer.write_all(&[9]).unwrap(); + assert_eq!(buffer.inner.into_inner(), vec![9, 2, 3, 4]); + } +} 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..35a59d243e8 --- /dev/null +++ b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java @@ -0,0 +1,43 @@ +/* + * 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; + +/** Task-owned callback for sending complete native Comet frames to a remote shuffle service. */ +@FunctionalInterface +public interface ShufflePartitionPusher { + + /** + * Accepts one complete, independently decodable Comet frame for an output partition. + * + *

The implementation must own any bytes it needs after this method returns. Success means + * acceptance, not remote map commit. Transport headers must not be included in the return value. + * Partial acceptance, cancellation, and supersession must fail rather than return a short count. + * The task owner must drain accepted pushes and commit or abort the attempt separately. + * + * @param partitionId the zero-based output partition + * @param bytes the complete encoded frame + * @param length the number of valid bytes in {@code bytes} + * @return exactly {@code length} on success + * @throws IOException if the frame cannot be accepted + */ + int pushPartitionData(int partitionId, byte[] bytes, int length) throws IOException; +} From 895e82b759ce878c9c09cd4f0d5d1557101b9d28 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:16:44 +0000 Subject: [PATCH 02/12] feat: add partition writer to shuffle plans --- native/core/src/execution/planner.rs | 208 +++++++++++++++++++- native/proto/src/proto/operator.proto | 6 + native/proto/src/proto/partitioning.proto | 20 ++ native/proto/tests/shuffle_writer_compat.rs | 131 ++++++++++++ 4 files changed, 363 insertions(+), 2 deletions(-) create mode 100644 native/proto/tests/shuffle_writer_compat.rs diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index ea5f1ef9238..46bc1f365cf 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1806,6 +1806,9 @@ impl PhysicalPlanner { ))) } OpStruct::ShuffleWriter(writer) => { + // Validate the destination before planning the child. In particular, an RSS or + // unknown destination must never fall through to the legacy local-file writer. + let (output_data_file, output_index_file) = local_shuffle_output_paths(writer)?; assert_eq!(children.len(), 1); let (scans, shuffle_scans, child) = self.create_plan(&children[0], inputs, partition_count)?; @@ -1842,8 +1845,8 @@ impl PhysicalPlanner { writer_input, partitioning, codec, - writer.output_data_file.clone(), - writer.output_index_file.clone(), + output_data_file.to_owned(), + output_index_file.to_owned(), writer.tracing_enabled, write_buffer_size, max_buffer_bytes, @@ -3788,6 +3791,43 @@ pub(crate) fn convert_spark_types_to_arrow_schema( arrow_schema } +/// Resolve the currently supported local destination without silently accepting remote plans. +fn local_shuffle_output_paths( + writer: &spark_operator::ShuffleWriter, +) -> Result<(&str, &str), ExecutionError> { + use datafusion_comet_proto::spark_partitioning::partition_writer::PartitionWriterStruct; + + let Some(partition_writer) = &writer.partition_writer else { + return Ok((&writer.output_data_file, &writer.output_index_file)); + }; + + match &partition_writer.partition_writer_struct { + Some(PartitionWriterStruct::LocalPartitionWriter(local)) => { + if local.output_data_file.is_empty() || local.output_index_file.is_empty() { + return Err(GeneralError( + "Local shuffle partition writer requires data and index paths".to_string(), + )); + } + if (!writer.output_data_file.is_empty() + && writer.output_data_file != local.output_data_file) + || (!writer.output_index_file.is_empty() + && writer.output_index_file != local.output_index_file) + { + return Err(GeneralError( + "Shuffle partition writer conflicts with legacy local paths".to_string(), + )); + } + Ok((&local.output_data_file, &local.output_index_file)) + } + Some(PartitionWriterStruct::RssPartitionWriter(_)) => Err(GeneralError( + "RSS shuffle partition writer is not supported yet".to_string(), + )), + None => Err(GeneralError( + "Shuffle partition writer has no recognized destination".to_string(), + )), + } +} + /// Wrap `child` in a `SchemaAlignExec` when its output drifts from what Spark catalyst /// declared. See . fn align_shuffle_writer_input( @@ -4634,6 +4674,170 @@ mod tests { }; use datafusion_comet_spark_expr::EvalMode; + mod shuffle_writer_destination { + use super::create_scan; + use crate::execution::operators::InputBatch; + use crate::execution::planner::{local_shuffle_output_paths, PhysicalPlanner}; + use datafusion::physical_plan::common::collect; + use datafusion::prelude::SessionContext; + use datafusion_comet_proto::spark_operator::{operator::OpStruct, Operator, ShuffleWriter}; + use datafusion_comet_proto::spark_partitioning::{ + partition_writer::PartitionWriterStruct, partitioning::PartitioningStruct, + LocalPartitionWriter, PartitionWriter, Partitioning, RssPartitionWriter, + SinglePartition, + }; + use prost::Message; + use tempfile::TempDir; + + fn local_destination(data: &str, index: &str) -> PartitionWriter { + PartitionWriter { + partition_writer_struct: Some(PartitionWriterStruct::LocalPartitionWriter( + LocalPartitionWriter { + output_data_file: data.to_string(), + output_index_file: index.to_string(), + }, + )), + } + } + + fn legacy_writer() -> ShuffleWriter { + ShuffleWriter { + output_data_file: "data".to_string(), + output_index_file: "index".to_string(), + ..Default::default() + } + } + + fn assert_rejected_before_child(writer: ShuffleWriter, expected: &str) { + let op = Operator { + // An invalid child ensures destination validation happens before child planning. + children: vec![Operator::default()], + op_struct: Some(OpStruct::ShuffleWriter(writer)), + ..Default::default() + }; + let err = PhysicalPlanner::default() + .create_plan(&op, &mut vec![], 1) + .err() + .expect("invalid destination must not produce a plan"); + assert!(err.to_string().contains(expected), "{err}"); + } + + #[test] + fn resolves_legacy_explicit_and_dual_written_local_paths() { + let legacy = legacy_writer(); + assert_eq!( + local_shuffle_output_paths(&legacy).unwrap(), + ("data", "index") + ); + + let explicit = ShuffleWriter { + partition_writer: Some(local_destination("data", "index")), + ..Default::default() + }; + assert_eq!( + local_shuffle_output_paths(&explicit).unwrap(), + ("data", "index") + ); + + let dual_written = ShuffleWriter { + partition_writer: explicit.partition_writer, + ..legacy + }; + assert_eq!( + local_shuffle_output_paths(&dual_written).unwrap(), + ("data", "index") + ); + } + + #[test] + fn rejects_conflicting_or_incomplete_local_paths() { + for (data, index) in [("other", "index"), ("data", "other")] { + let writer = ShuffleWriter { + partition_writer: Some(local_destination(data, index)), + ..legacy_writer() + }; + assert_rejected_before_child(writer, "conflicts with legacy local paths"); + } + for (data, index) in [("", "index"), ("data", ""), ("", "")] { + let writer = ShuffleWriter { + partition_writer: Some(local_destination(data, index)), + ..Default::default() + }; + assert_rejected_before_child(writer, "requires data and index paths"); + } + } + + #[test] + fn rejects_rss_even_when_legacy_paths_are_present() { + for handle in [0, 7, -1] { + for mut writer in [ShuffleWriter::default(), legacy_writer()] { + writer.partition_writer = Some(PartitionWriter { + partition_writer_struct: Some(PartitionWriterStruct::RssPartitionWriter( + RssPartitionWriter { + rss_partition_pusher: handle, + }, + )), + }); + assert_rejected_before_child( + writer, + "RSS shuffle partition writer is not supported", + ); + } + } + } + + #[test] + fn rejects_empty_and_unknown_destinations() { + for bytes in [b"\x12\x00".as_slice(), b"\x12\x02\x1a\x00".as_slice()] { + let decoded = ShuffleWriter::decode(bytes).unwrap(); + let writer = ShuffleWriter { + partition_writer: decoded.partition_writer, + ..legacy_writer() + }; + assert_rejected_before_child(writer, "no recognized destination"); + } + } + + #[tokio::test] + async fn local_destinations_execute_the_existing_file_writer() { + let dir = TempDir::new().unwrap(); + for mode in ["legacy", "explicit", "dual"] { + let data = dir.path().join(format!("{mode}.data")); + let index = dir.path().join(format!("{mode}.index")); + let data_path = data.to_str().unwrap(); + let index_path = index.to_str().unwrap(); + let writer = ShuffleWriter { + partitioning: Some(Partitioning { + partitioning_struct: Some(PartitioningStruct::SinglePartition( + SinglePartition {}, + )), + }), + partition_writer: (mode != "legacy") + .then(|| local_destination(data_path, index_path)), + output_data_file: if mode == "explicit" { "" } else { data_path }.to_string(), + output_index_file: if mode == "explicit" { "" } else { index_path }.to_string(), + write_buffer_size: 1024, + ..Default::default() + }; + let op = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::ShuffleWriter(writer)), + ..Default::default() + }; + let planner = PhysicalPlanner::default(); + let (mut scans, _, plan) = planner.create_plan(&op, &mut vec![], 1).unwrap(); + scans[0].set_input_batch(InputBatch::EOF); + let stream = plan + .native_plan + .execute(0, SessionContext::new().task_ctx()) + .unwrap(); + assert!(collect(stream).await.unwrap().is_empty()); + assert!(std::fs::read(data).unwrap().is_empty()); + assert_eq!(std::fs::read(index).unwrap(), vec![0; 16]); + } + } + } + #[test] fn test_unpack_dictionary_primitive() { let op_scan = Operator { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ec9448709cf..aaf36b33328 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -687,6 +687,12 @@ enum CompressionCodec { message ShuffleWriter { spark.spark_partitioning.Partitioning partitioning = 1; + // Absent means legacy local output using fields 3 and 4. An explicit local writer may + // duplicate those paths for older native libraries, but the paths must agree. + // Do not emit RSS until the native library's RSS protocol support has been verified. + // RSS must not populate the legacy local paths, and is not executable yet. + spark.spark_partitioning.PartitionWriter partition_writer = 2; + // Retained for compatibility with existing JVM and native plans. string output_data_file = 3; string output_index_file = 4; CompressionCodec codec = 5; diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index e70b8264f02..bb3acef7ccc 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -58,3 +58,23 @@ message RoundRobinPartition { // Maximum number of columns to hash. 0 means no limit (hash all columns). int32 max_hash_columns = 2; } + +// Storage destination for already-partitioned shuffle batches. +// A present message must select a recognized writer; an empty or unknown selection is an error. +message PartitionWriter { + oneof partition_writer_struct { + LocalPartitionWriter local_partition_writer = 1; + RssPartitionWriter rss_partition_writer = 2; + } +} + +message LocalPartitionWriter { + string output_data_file = 1; + string output_index_file = 2; +} + +message RssPartitionWriter { + // Opaque, task-scoped pusher handle. This schema does not make the handle usable: + // RSS plans remain unsupported until capability negotiation and safe handle ownership exist. + int64 rss_partition_pusher = 1; +} diff --git a/native/proto/tests/shuffle_writer_compat.rs b/native/proto/tests/shuffle_writer_compat.rs new file mode 100644 index 00000000000..15c6c0f0f5f --- /dev/null +++ b/native/proto/tests/shuffle_writer_compat.rs @@ -0,0 +1,131 @@ +// 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 datafusion_comet_proto::spark_operator::ShuffleWriter; +use datafusion_comet_proto::spark_partitioning::{ + partition_writer::PartitionWriterStruct, partitioning::PartitioningStruct, + LocalPartitionWriter, PartitionWriter, Partitioning, RssPartitionWriter, SinglePartition, +}; +use prost::Message; + +// The original ShuffleWriter schema, frozen independently of the generated current schema. +#[derive(Clone, PartialEq, Message)] +struct LegacyShuffleWriter { + #[prost(message, optional, tag = "1")] + partitioning: Option, + #[prost(string, tag = "3")] + output_data_file: String, + #[prost(string, tag = "4")] + output_index_file: String, +} + +fn single_partition() -> Option { + Some(Partitioning { + partitioning_struct: Some(PartitioningStruct::SinglePartition(SinglePartition {})), + }) +} + +fn local_writer() -> PartitionWriter { + PartitionWriter { + partition_writer_struct: Some(PartitionWriterStruct::LocalPartitionWriter( + LocalPartitionWriter { + output_data_file: "data".to_string(), + output_index_file: "index".to_string(), + }, + )), + } +} + +#[test] +fn legacy_shuffle_writer_wire_format_is_unchanged() { + let legacy = LegacyShuffleWriter { + partitioning: single_partition(), + output_data_file: "data".to_string(), + output_index_file: "index".to_string(), + }; + let bytes = legacy.encode_to_vec(); + assert_eq!(bytes, b"\x0a\x02\x12\x00\x1a\x04data\x22\x05index"); + + let decoded = ShuffleWriter::decode(bytes.as_slice()).unwrap(); + assert_eq!(decoded.partitioning, legacy.partitioning); + assert!(decoded.partition_writer.is_none()); + assert_eq!(decoded.output_data_file, legacy.output_data_file); + assert_eq!(decoded.output_index_file, legacy.output_index_file); + assert_eq!(decoded.encode_to_vec(), bytes); +} + +#[test] +fn dual_written_local_destination_is_readable_by_legacy_decoder() { + let writer = ShuffleWriter { + partitioning: single_partition(), + partition_writer: Some(local_writer()), + output_data_file: "data".to_string(), + output_index_file: "index".to_string(), + ..Default::default() + }; + let bytes = writer.encode_to_vec(); + assert_eq!(ShuffleWriter::decode(bytes.as_slice()).unwrap(), writer); + + let legacy = LegacyShuffleWriter::decode(bytes.as_slice()).unwrap(); + assert_eq!(legacy.partitioning, writer.partitioning); + assert_eq!(legacy.output_data_file, writer.output_data_file); + assert_eq!(legacy.output_index_file, writer.output_index_file); +} + +#[test] +fn partition_writer_uses_the_existing_rss_prototype_tags() { + let local = ShuffleWriter { + partition_writer: Some(local_writer()), + ..Default::default() + }; + // ShuffleWriter field 2 -> local variant 1 -> data/index fields 1/2. + let local_bytes = b"\x12\x0f\x0a\x0d\x0a\x04data\x12\x05index"; + assert_eq!(local.encode_to_vec(), local_bytes); + assert_eq!( + ShuffleWriter::decode(local_bytes.as_slice()).unwrap(), + local + ); + + let rss = ShuffleWriter { + partition_writer: Some(PartitionWriter { + partition_writer_struct: Some(PartitionWriterStruct::RssPartitionWriter( + RssPartitionWriter { + rss_partition_pusher: 7, + }, + )), + }), + ..Default::default() + }; + // ShuffleWriter field 2 -> RSS variant 2 -> opaque handle field 1. + let rss_bytes = b"\x12\x04\x12\x02\x08\x07"; + assert_eq!(rss.encode_to_vec(), rss_bytes); + assert_eq!(ShuffleWriter::decode(rss_bytes.as_slice()).unwrap(), rss); + + // Older native libraries cannot distinguish RSS from an absent destination. A producer + // therefore needs a capability handshake before it is allowed to emit an RSS plan. + let legacy = LegacyShuffleWriter::decode(rss_bytes.as_slice()).unwrap(); + assert!(legacy.output_data_file.is_empty()); + assert!(legacy.output_index_file.is_empty()); +} + +#[test] +fn unknown_destination_stays_distinct_from_legacy_local() { + // A future field 3 in the PartitionWriter oneof is unknown to this decoder. + let decoded = ShuffleWriter::decode(b"\x12\x02\x1a\x00".as_slice()).unwrap(); + let destination = decoded.partition_writer.expect("field 2 must stay present"); + assert!(destination.partition_writer_struct.is_none()); +} From e05541a37263d6af42edbd053c93846546e776aa Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:16:46 +0000 Subject: [PATCH 03/12] feat: add destination-aware native shuffle execution --- native/shuffle/src/lib.rs | 2 + native/shuffle/src/shuffle_destination.rs | 68 ++ native/shuffle/src/shuffle_writer.rs | 191 ++++-- .../src/writers/rss/rss_partition_writer.rs | 30 +- native/shuffle/tests/shuffle_destinations.rs | 594 ++++++++++++++++++ 5 files changed, 818 insertions(+), 67 deletions(-) create mode 100644 native/shuffle/src/shuffle_destination.rs create mode 100644 native/shuffle/tests/shuffle_destinations.rs diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index a7b43fabbaa..43b351d31ca 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -20,6 +20,7 @@ pub mod ipc; pub(crate) mod metrics; pub(crate) mod partitioners; mod schema_align; +mod shuffle_destination; mod shuffle_writer; mod spark_crc32c_hasher; pub mod spark_unsafe; @@ -28,5 +29,6 @@ pub(crate) mod writers; pub use comet_partitioning::CometPartitioning; pub use ipc::read_ipc_compressed; pub use schema_align::SchemaAlignExec; +pub use shuffle_destination::ShuffleDestination; pub use shuffle_writer::ShuffleWriterExec; pub use writers::{CompressionCodec, PartitionPusher, RssPartitionWriter, ShuffleBlockWriter}; diff --git a/native/shuffle/src/shuffle_destination.rs b/native/shuffle/src/shuffle_destination.rs new file mode 100644 index 00000000000..0b74a3c70ba --- /dev/null +++ b/native/shuffle/src/shuffle_destination.rs @@ -0,0 +1,68 @@ +// 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 std::fmt::{self, Debug, Formatter}; +use std::sync::Arc; + +use crate::PartitionPusher; + +/// Resolved output destination for a native shuffle writer. +/// +/// Unlike the protobuf destination, RSS carries an owned pusher, not a numeric handle. +/// Resolving and validating task-owned handles belongs to the caller. Cloning an RSS +/// destination shares the same pusher; a different task attempt needs its own pusher. +/// Holding it keeps its Rust state alive, but does not provide remote commit or cancellation +/// semantics. +#[derive(Clone)] +pub enum ShuffleDestination { + /// One local data file and its partition-offset index. + Local { + output_data_file: String, + output_index_file: String, + /// Size of the local file write buffer in bytes. + write_buffer_size: usize, + }, + /// Complete encoded frames sent through a task-owned callback. + Rss { + pusher: Arc, + /// Maximum encoded bytes in one complete Comet frame. + max_frame_bytes: usize, + }, +} + +impl Debug for ShuffleDestination { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Local { + output_data_file, + output_index_file, + write_buffer_size, + } => f + .debug_struct("Local") + .field("output_data_file", output_data_file) + .field("output_index_file", output_index_file) + .field("write_buffer_size", write_buffer_size) + .finish(), + Self::Rss { + max_frame_bytes, .. + } => f + .debug_struct("Rss") + .field("max_frame_bytes", max_frame_bytes) + .finish_non_exhaustive(), + } + } +} diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 46d3b592bc9..c4b7bbcb53c 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -22,10 +22,12 @@ use crate::partitioners::{ EmptySchemaShufflePartitioner, MultiPartitionShuffleRepartitioner, ShufflePartitioner, SinglePartitionShufflePartitioner, }; -use crate::writers::LocalPartitionWriter; -use crate::{CometPartitioning, CompressionCodec, ShuffleBlockWriter}; +use crate::writers::{LocalPartitionWriter, PartitionWriter}; +use crate::{ + CometPartitioning, CompressionCodec, PartitionPusher, RssPartitionWriter, ShuffleBlockWriter, + ShuffleDestination, +}; use async_trait::async_trait; -use datafusion::common::exec_datafusion_err; use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::EmptyRecordBatchStream; @@ -54,10 +56,8 @@ pub struct ShuffleWriterExec { input: Arc, /// Partitioning scheme to use partitioning: CometPartitioning, - /// Output data file path - output_data_file: String, - /// Output index file path - output_index_file: String, + /// Resolved storage destination + shuffle_destination: ShuffleDestination, /// Metrics metrics: ExecutionPlanMetricsSet, /// Cache for expensive-to-compute plan properties @@ -65,14 +65,12 @@ pub struct ShuffleWriterExec { /// The compression codec to use when compressing shuffle blocks codec: CompressionCodec, tracing_enabled: bool, - /// Size of the write buffer in bytes - write_buffer_size: usize, /// Maximum bytes buffered in memory before spilling; `None` disables the limit max_buffer_bytes: Option, } impl ShuffleWriterExec { - /// Create a new ShuffleWriterExec + /// Create a local-file shuffle writer. Existing callers do not need to select a destination. #[allow(clippy::too_many_arguments)] pub fn try_new( input: Arc, @@ -84,6 +82,44 @@ impl ShuffleWriterExec { write_buffer_size: usize, max_buffer_bytes: Option, ) -> Result { + Self::try_new_with_destination( + input, + partitioning, + codec, + ShuffleDestination::Local { + output_data_file, + output_index_file, + write_buffer_size, + }, + tracing_enabled, + max_buffer_bytes, + ) + } + + /// Create a shuffle writer with an already-resolved storage destination. + /// + /// The RSS pusher must belong to the calling task attempt. This constructor does not + /// resolve protobuf handles or perform a remote backend's task-level commit protocol. + /// `max_buffer_bytes` controls the multi-partition buffer for either destination. + /// Local file buffering and RSS frame limits are configured by `shuffle_destination`. + pub fn try_new_with_destination( + input: Arc, + partitioning: CometPartitioning, + codec: CompressionCodec, + shuffle_destination: ShuffleDestination, + tracing_enabled: bool, + max_buffer_bytes: Option, + ) -> Result { + if let ShuffleDestination::Rss { + max_frame_bytes, .. + } = &shuffle_destination + { + RssPartitionWriter::>::validate_options( + partitioning.partition_count(), + *max_frame_bytes, + )?; + } + let cache = Arc::new(PlanProperties::new( EquivalenceProperties::new(Arc::clone(&input.schema())), Partitioning::UnknownPartitioning(1), @@ -95,12 +131,10 @@ impl ShuffleWriterExec { input, partitioning, metrics: ExecutionPlanMetricsSet::new(), - output_data_file, - output_index_file, + shuffle_destination, cache, codec, tracing_enabled, - write_buffer_size, max_buffer_bytes, }) } @@ -149,14 +183,12 @@ impl ExecutionPlan for ShuffleWriterExec { children: Vec>, ) -> Result> { match children.len() { - 1 => Ok(Arc::new(ShuffleWriterExec::try_new( + 1 => Ok(Arc::new(ShuffleWriterExec::try_new_with_destination( Arc::clone(&children[0]), self.partitioning.clone(), self.codec.clone(), - self.output_data_file.clone(), - self.output_index_file.clone(), + self.shuffle_destination.clone(), self.tracing_enabled, - self.write_buffer_size, self.max_buffer_bytes, )?)), _ => panic!("ShuffleWriterExec wrong number of children"), @@ -179,14 +211,12 @@ impl ExecutionPlan for ShuffleWriterExec { futures::stream::once(external_shuffle( input, partition, - self.output_data_file.clone(), - self.output_index_file.clone(), + self.shuffle_destination.clone(), self.partitioning.clone(), metrics, context, self.codec.clone(), self.tracing_enabled, - self.write_buffer_size, self.max_buffer_bytes, )) .try_flatten(), @@ -198,46 +228,109 @@ impl ExecutionPlan for ShuffleWriterExec { async fn external_shuffle( mut input: SendableRecordBatchStream, partition: usize, - output_data_file: String, - output_index_file: String, + shuffle_destination: ShuffleDestination, partitioning: CometPartitioning, metrics: ShufflePartitionerMetrics, context: Arc, codec: CompressionCodec, tracing_enabled: bool, - write_buffer_size: usize, max_buffer_bytes: Option, ) -> Result { let schema = input.schema(); - let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?; - let local_partition_writer = LocalPartitionWriter::try_new( - output_data_file, - output_index_file, - shuffle_block_writer, - partitioning.partition_count(), - context.session_config().batch_size(), - write_buffer_size, - context.runtime_env(), - )?; - - let mut repartitioner: Box = match &partitioning { + let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec)?; + let mut repartitioner = match shuffle_destination { + ShuffleDestination::Local { + output_data_file, + output_index_file, + write_buffer_size, + } => { + let writer = LocalPartitionWriter::try_new( + output_data_file, + output_index_file, + shuffle_block_writer, + partitioning.partition_count(), + context.session_config().batch_size(), + write_buffer_size, + context.runtime_env(), + )?; + create_repartitioner( + partition, + writer, + Arc::clone(&schema), + partitioning, + metrics, + &context, + tracing_enabled, + max_buffer_bytes, + )? + } + ShuffleDestination::Rss { + pusher, + max_frame_bytes, + } => { + let writer = RssPartitionWriter::try_new( + pusher, + shuffle_block_writer, + partitioning.partition_count(), + max_frame_bytes, + )?; + create_repartitioner( + partition, + writer, + Arc::clone(&schema), + partitioning, + metrics, + &context, + tracing_enabled, + max_buffer_bytes, + )? + } + }; + + while let Some(batch) = input.next().await { + // Await insertion before pulling the next input batch, which may reuse its buffers. + // Preserve typed errors, including the original Java callback exception, unchanged. + repartitioner.insert_batch(batch?).await?; + } + + repartitioner.shuffle_write()?; + + // Writer-side finalization is not a remote map commit. The task owner handles that. + // Shuffle writers always have empty output. + Ok(Box::pin(EmptyRecordBatchStream::new(schema)) as SendableRecordBatchStream) +} + +/// Keep the existing partitioning algorithms shared by both storage backends. PartitionWriter +/// has generic methods, so choose its concrete type before erasing the partitioner's type. +#[allow(clippy::too_many_arguments)] +fn create_repartitioner( + partition: usize, + partition_writer: W, + schema: SchemaRef, + partitioning: CometPartitioning, + metrics: ShufflePartitionerMetrics, + context: &TaskContext, + tracing_enabled: bool, + max_buffer_bytes: Option, +) -> Result> { + Ok(match &partitioning { _ if schema.fields().is_empty() => { log::debug!("found empty schema, overriding {partitioning:?} partitioning with EmptySchemaShufflePartitioner"); Box::new(EmptySchemaShufflePartitioner::try_new( - local_partition_writer, - Arc::clone(&schema), + partition_writer, + schema, partitioning.partition_count(), metrics, )?) } any if any.partition_count() == 1 => Box::new(SinglePartitionShufflePartitioner::new( - local_partition_writer, + partition_writer, metrics, )), _ => Box::new(MultiPartitionShuffleRepartitioner::try_new( partition, - local_partition_writer, + partition_writer, partitioning, metrics, context.runtime_env(), @@ -245,25 +338,7 @@ async fn external_shuffle( tracing_enabled, max_buffer_bytes, )?), - }; - - while let Some(batch) = input.next().await { - // Await the repartitioner to insert the batch and shuffle the rows - // into the corresponding partition buffer. - // Otherwise, pull the next batch from the input stream might overwrite the - // current batch in the repartitioner. - repartitioner - .insert_batch(batch?) - .await - .map_err(|err| exec_datafusion_err!("Error inserting batch: {err}"))?; - } - - repartitioner - .shuffle_write() - .map_err(|err| exec_datafusion_err!("Error in shuffle write: {err}"))?; - - // shuffle writer always has empty output - Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))) as SendableRecordBatchStream) + }) } #[cfg(test)] diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs index b4feaef5199..b487c9741e2 100644 --- a/native/shuffle/src/writers/rss/rss_partition_writer.rs +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -16,6 +16,7 @@ // under the License. use std::io::{self, Cursor, Seek, SeekFrom, Write}; +use std::sync::Arc; use arrow::record_batch::RecordBatch; use datafusion::common::{DataFusionError, Result}; @@ -34,6 +35,12 @@ pub trait PartitionPusher: Send + Sync { fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()>; } +impl PartitionPusher for Arc

{ + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { + self.as_ref().push_partition_data(partition_id, frame) + } +} + impl PartitionPusher for JavaShufflePartitionPusher { fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { JavaShufflePartitionPusher::push_partition_data(self, partition_id, frame) @@ -64,6 +71,19 @@ impl RssPartitionWriter

{ num_partitions: usize, max_frame_bytes: usize, ) -> Result { + Self::validate_options(num_partitions, max_frame_bytes)?; + Ok(Self { + pusher, + block_writer, + num_partitions, + max_frame_bytes, + next_partition: 0, + finished: false, + failed: false, + }) + } + + pub(crate) fn validate_options(num_partitions: usize, max_frame_bytes: usize) -> Result<()> { if num_partitions == 0 || num_partitions > i32::MAX as usize { return Err(DataFusionError::Configuration( "Invalid RSS partition count".into(), @@ -74,15 +94,7 @@ impl RssPartitionWriter

{ "Invalid RSS frame byte limit".into(), )); } - Ok(Self { - pusher, - block_writer, - num_partitions, - max_frame_bytes, - next_partition: 0, - finished: false, - failed: false, - }) + Ok(()) } /// Encodes and synchronously submits one batch. A failed writer cannot be reused. diff --git a/native/shuffle/tests/shuffle_destinations.rs b/native/shuffle/tests/shuffle_destinations.rs new file mode 100644 index 00000000000..c19e55c8f50 --- /dev/null +++ b/native/shuffle/tests/shuffle_destinations.rs @@ -0,0 +1,594 @@ +// 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 std::error::Error; +use std::fmt; +use std::path::Path; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, RecordBatchOptions}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::row::{RowConverter, SortField}; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::config::SessionConfig; +use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_plan::common::collect; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_comet_shuffle::{ + read_ipc_compressed, CometPartitioning, CompressionCodec, PartitionPusher, ShuffleDestination, + ShuffleWriterExec, +}; +use datafusion_comet_spark_expr::murmur3::create_murmur3_hashes; + +const MAX_FRAME_BYTES: usize = 1024 * 1024; +const RANGE_BOUNDS: [i64; 3] = [17, 41, 63]; + +#[derive(Default)] +struct RecordingPusher(Mutex)>>); + +impl PartitionPusher for RecordingPusher { + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { + self.0.lock().unwrap().push((partition_id, frame.to_vec())); + Ok(()) + } +} + +impl RecordingPusher { + fn decoded(&self, num_partitions: usize) -> Vec> { + let mut partitions = vec![vec![]; num_partitions]; + for (partition_id, frame) in self.0.lock().unwrap().iter() { + assert!(*partition_id < num_partitions); + partitions[*partition_id].push(decode_frame(frame)); + } + partitions + } +} + +fn int_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) +} + +fn batch(values: impl IntoIterator) -> RecordBatch { + RecordBatch::try_new( + int_schema(), + vec![Arc::new(Int64Array::from_iter_values(values))], + ) + .unwrap() +} + +fn empty_schema_batch(num_rows: usize) -> RecordBatch { + RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + ) + .unwrap() +} + +fn source(schema: SchemaRef, batches: &[RecordBatch]) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[batches.to_vec()], schema, None).unwrap(), + ))) +} + +fn hash_partitioning(num_partitions: usize) -> CometPartitioning { + CometPartitioning::Hash(vec![Arc::new(Column::new("id", 0))], num_partitions) +} + +fn partitionings() -> [CometPartitioning; 4] { + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( + "id", 0, + )))]) + .unwrap(); + let converter = RowConverter::new(vec![SortField::new(DataType::Int64)]).unwrap(); + let bounds: ArrayRef = Arc::new(Int64Array::from(RANGE_BOUNDS.to_vec())); + let bounds = converter + .convert_columns(&[bounds]) + .unwrap() + .iter() + .map(|row| row.owned()) + .collect(); + [ + CometPartitioning::SinglePartition, + hash_partitioning(4), + CometPartitioning::RoundRobin(4, 0), + CometPartitioning::RangePartitioning(ordering, 4, Arc::new(converter), bounds), + ] +} + +fn codecs() -> [CompressionCodec; 4] { + [ + CompressionCodec::None, + CompressionCodec::Lz4Frame, + CompressionCodec::Snappy, + CompressionCodec::Zstd(1), + ] +} + +fn writer( + input: Arc, + partitioning: CometPartitioning, + codec: CompressionCodec, + shuffle_destination: ShuffleDestination, + max_buffer_bytes: Option, +) -> ShuffleWriterExec { + ShuffleWriterExec::try_new_with_destination( + input, + partitioning, + codec, + shuffle_destination, + false, + max_buffer_bytes, + ) + .unwrap() +} + +fn rss(pusher: Arc) -> ShuffleDestination { + ShuffleDestination::Rss { + pusher, + max_frame_bytes: MAX_FRAME_BYTES, + } +} + +async fn execute(plan: &dyn ExecutionPlan) -> Result<()> { + let context = SessionContext::new_with_config(SessionConfig::new().with_batch_size(32)); + execute_with_context(plan, &context).await +} + +async fn execute_with_context(plan: &dyn ExecutionPlan, context: &SessionContext) -> Result<()> { + let output = collect(plan.execute(0, context.task_ctx())?).await?; + assert!(output.is_empty(), "shuffle is a sink operator"); + Ok(()) +} + +fn paths(dir: &Path, tag: &str) -> (String, String) { + ( + dir.join(format!("{tag}.data")).to_str().unwrap().to_owned(), + dir.join(format!("{tag}.index")) + .to_str() + .unwrap() + .to_owned(), + ) +} + +fn local(output_paths: &(String, String)) -> ShuffleDestination { + ShuffleDestination::Local { + output_data_file: output_paths.0.clone(), + output_index_file: output_paths.1.clone(), + write_buffer_size: 4096, + } +} + +fn decode_frame(frame: &[u8]) -> RecordBatch { + assert!(frame.len() >= 20, "missing Comet frame header"); + let declared_length = u64::from_le_bytes(frame[..8].try_into().unwrap()); + assert_eq!(declared_length + 8, frame.len() as u64); + let batch = read_ipc_compressed(&frame[16..]).unwrap(); + let field_count = u64::from_le_bytes(frame[8..16].try_into().unwrap()); + assert_eq!(field_count, batch.num_columns() as u64); + batch +} + +fn decode_frames(mut bytes: &[u8]) -> Vec { + let mut batches = vec![]; + while !bytes.is_empty() { + let length = u64::from_le_bytes(bytes[..8].try_into().unwrap()) as usize + 8; + let (frame, rest) = bytes.split_at(length); + batches.push(decode_frame(frame)); + bytes = rest; + } + batches +} + +fn read_local(output_paths: &(String, String), num_partitions: usize) -> Vec> { + let data = std::fs::read(&output_paths.0).unwrap(); + let index = std::fs::read(&output_paths.1).unwrap(); + assert_eq!(index.len(), (num_partitions + 1) * 8); + let offsets: Vec = index + .chunks_exact(8) + .map(|bytes| u64::from_le_bytes(bytes.try_into().unwrap()) as usize) + .collect(); + assert_eq!(offsets[0], 0); + assert_eq!(offsets[num_partitions], data.len()); + offsets + .windows(2) + .map(|range| { + assert!(range[0] <= range[1]); + decode_frames(&data[range[0]..range[1]]) + }) + .collect() +} + +fn values(partitions: &[Vec]) -> Vec> { + partitions + .iter() + .map(|batches| { + let mut values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + // Repartitioning does not promise output order or identical IPC batch boundaries. + values.sort_unstable(); + values + }) + .collect() +} + +fn expected_values(partitioning: &CometPartitioning, batches: &[RecordBatch]) -> Vec> { + let num_partitions = partitioning.partition_count(); + let mut expected = vec![vec![]; num_partitions]; + for batch in batches { + let mut hashes = vec![42; batch.num_rows()]; + create_murmur3_hashes(batch.columns(), &mut hashes).unwrap(); + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for (value, hash) in column.values().iter().zip(hashes) { + let partition_id = match partitioning { + CometPartitioning::SinglePartition => 0, + CometPartitioning::Hash(_, _) | CometPartitioning::RoundRobin(_, _) => { + (hash as i32).rem_euclid(num_partitions as i32) as usize + } + CometPartitioning::RangePartitioning(_, _, _, _) => { + RANGE_BOUNDS.partition_point(|bound| bound <= value) + } + }; + expected[partition_id].push(*value); + } + } + for partition in &mut expected { + partition.sort_unstable(); + } + expected +} + +#[tokio::test] +#[cfg_attr(miri, ignore)] // Miri cannot call the compression libraries. +async fn destinations_preserve_local_bytes_and_partition_rows() { + let batches = vec![batch(0..31), batch(31..79), batch([]), batch(79..95)]; + let input = source(int_schema(), &batches); + let dir = tempfile::tempdir().unwrap(); + + for (codec_id, codec) in codecs().into_iter().enumerate() { + for (partitioning_id, partitioning) in partitionings().into_iter().enumerate() { + let num_partitions = partitioning.partition_count(); + let expected = expected_values(&partitioning, &batches); + let pusher = Arc::new(RecordingPusher::default()); + let remote = writer( + Arc::clone(&input), + partitioning.clone(), + codec.clone(), + rss(pusher.clone()), + None, + ); + execute(&remote).await.unwrap(); + assert_eq!(values(&pusher.decoded(num_partitions)), expected); + + let tag = format!("{codec_id}-{partitioning_id}"); + let legacy_paths = paths(dir.path(), &format!("legacy-{tag}")); + let explicit_paths = paths(dir.path(), &format!("explicit-{tag}")); + let legacy = ShuffleWriterExec::try_new( + Arc::clone(&input), + partitioning.clone(), + codec.clone(), + legacy_paths.0.clone(), + legacy_paths.1.clone(), + false, + 4096, + None, + ) + .unwrap(); + let explicit = writer( + Arc::clone(&input), + partitioning, + codec.clone(), + local(&explicit_paths), + None, + ); + execute(&legacy).await.unwrap(); + execute(&explicit).await.unwrap(); + assert_eq!( + std::fs::read(&legacy_paths.0).unwrap(), + std::fs::read(&explicit_paths.0).unwrap() + ); + assert_eq!( + std::fs::read(&legacy_paths.1).unwrap(), + std::fs::read(&explicit_paths.1).unwrap() + ); + assert_eq!( + values(&read_local(&explicit_paths, num_partitions)), + expected + ); + } + } +} + +#[tokio::test] +async fn rss_buffer_limit_spills_without_changing_partition_rows() { + let batches: Vec<_> = (0..16).map(|i| batch(i * 16..(i + 1) * 16)).collect(); + let input = source(int_schema(), &batches); + let partitioning = hash_partitioning(4); + let expected = expected_values(&partitioning, &batches); + // RSS flushing must not ask DataFusion to create a local spill file. + let runtime = RuntimeEnvBuilder::new() + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build() + .unwrap(); + let context = SessionContext::new_with_config_rt( + SessionConfig::new().with_batch_size(32), + Arc::new(runtime), + ); + + for max_buffer_bytes in [None, Some(1)] { + let pusher = Arc::new(RecordingPusher::default()); + let exec = writer( + Arc::clone(&input), + partitioning.clone(), + CompressionCodec::None, + rss(pusher.clone()), + max_buffer_bytes, + ); + execute_with_context(&exec, &context).await.unwrap(); + assert_eq!(values(&pusher.decoded(4)), expected); + let spills = exec.metrics().unwrap().spill_count().unwrap(); + if max_buffer_bytes.is_some() { + assert!(spills > 0, "the test must exercise insertion-time flushing"); + } else { + assert_eq!(spills, 0); + } + } +} + +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn rss_empty_schema_preserves_row_counts() { + let batches = vec![ + empty_schema_batch(7), + empty_schema_batch(0), + empty_schema_batch(11), + ]; + for codec in codecs() { + for partitioning in [ + CometPartitioning::SinglePartition, + CometPartitioning::RoundRobin(4, 0), + ] { + let num_partitions = partitioning.partition_count(); + let pusher = Arc::new(RecordingPusher::default()); + let exec = writer( + source(batches[0].schema(), &batches), + partitioning, + codec.clone(), + rss(pusher.clone()), + None, + ); + execute(&exec).await.unwrap(); + let decoded = pusher.decoded(num_partitions); + assert_eq!(decoded[0].len(), 1); + assert_eq!(decoded[0][0].num_columns(), 0); + assert_eq!(decoded[0][0].num_rows(), 18); + assert!(decoded[1..].iter().all(Vec::is_empty)); + } + } +} + +#[tokio::test] +async fn empty_input_produces_no_remote_frames_or_local_data() { + let dir = tempfile::tempdir().unwrap(); + for (schema_id, schema) in [int_schema(), Arc::new(Schema::empty())] + .into_iter() + .enumerate() + { + for (partitioning_id, partitioning) in [ + CometPartitioning::SinglePartition, + CometPartitioning::RoundRobin(4, 0), + ] + .into_iter() + .enumerate() + { + // A missing batch and a present zero-row batch must behave identically. + for (batch_id, batches) in [vec![], vec![RecordBatch::new_empty(Arc::clone(&schema))]] + .into_iter() + .enumerate() + { + let num_partitions = partitioning.partition_count(); + let input = source(Arc::clone(&schema), &batches); + let pusher = Arc::new(RecordingPusher::default()); + execute(&writer( + Arc::clone(&input), + partitioning.clone(), + CompressionCodec::None, + rss(pusher.clone()), + None, + )) + .await + .unwrap(); + assert!(pusher.0.lock().unwrap().is_empty()); + + let output_paths = paths( + dir.path(), + &format!("empty-{schema_id}-{partitioning_id}-{batch_id}"), + ); + execute(&writer( + input, + partitioning.clone(), + CompressionCodec::None, + local(&output_paths), + None, + )) + .await + .unwrap(); + assert!(read_local(&output_paths, num_partitions) + .iter() + .all(Vec::is_empty)); + } + } + } +} + +#[tokio::test] +async fn replacing_children_keeps_the_selected_destination() { + let input = source(int_schema(), &[batch([1, 2])]); + let replacement = source(int_schema(), &[batch([8, 9, 10])]); + let pusher = Arc::new(RecordingPusher::default()); + let dir = tempfile::tempdir().unwrap(); + let output_paths = paths(dir.path(), "replacement"); + + for shuffle_destination in [rss(pusher.clone()), local(&output_paths)] { + let original = Arc::new(writer( + Arc::clone(&input), + CometPartitioning::SinglePartition, + CompressionCodec::None, + shuffle_destination, + None, + )); + let rewritten = original + .with_new_children(vec![Arc::clone(&replacement)]) + .unwrap(); + execute(rewritten.as_ref()).await.unwrap(); + } + assert_eq!(values(&pusher.decoded(1)), vec![vec![8, 9, 10]]); + assert_eq!(values(&read_local(&output_paths, 1)), vec![vec![8, 9, 10]]); +} + +#[test] +fn rss_configuration_is_rejected_before_execution() { + let input = source(int_schema(), &[]); + let pusher = Arc::new(RecordingPusher::default()); + for (num_partitions, max_frame_bytes) in [ + (0, MAX_FRAME_BYTES), + (i32::MAX as usize + 1, MAX_FRAME_BYTES), + (1, 0), + (1, 19), + (1, i32::MAX as usize + 1), + ] { + let result = ShuffleWriterExec::try_new_with_destination( + Arc::clone(&input), + CometPartitioning::RoundRobin(num_partitions, 0), + CompressionCodec::None, + ShuffleDestination::Rss { + pusher: pusher.clone(), + max_frame_bytes, + }, + false, + None, + ); + assert!(matches!(result, Err(DataFusionError::Configuration(_)))); + } + assert!(pusher.0.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn oversized_frame_fails_before_calling_the_pusher() { + let pusher = Arc::new(RecordingPusher::default()); + let exec = writer( + source(int_schema(), &[batch(0..64)]), + CometPartitioning::SinglePartition, + CompressionCodec::None, + ShuffleDestination::Rss { + pusher: pusher.clone(), + max_frame_bytes: 20, + }, + None, + ); + let error = execute(&exec).await.unwrap_err(); + assert!(error + .to_string() + .contains("RSS frame exceeds its byte limit")); + assert!(pusher.0.lock().unwrap().is_empty()); +} + +#[derive(Debug)] +struct CallbackFailure(Arc<()>); + +impl fmt::Display for CallbackFailure { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("test callback failure") + } +} + +impl Error for CallbackFailure {} + +struct FailingPusher { + marker: Arc<()>, + calls: AtomicUsize, +} + +impl PartitionPusher for FailingPusher { + fn push_partition_data(&self, _partition_id: usize, _frame: &[u8]) -> Result<()> { + self.calls.fetch_add(1, Ordering::Relaxed); + Err(DataFusionError::External(Box::new(CallbackFailure( + Arc::clone(&self.marker), + )))) + } +} + +#[tokio::test] +async fn callback_errors_keep_their_type_and_identity() { + for (partitioning, batches, max_buffer_bytes) in [ + // Single-partition writes happen while inserting the input batch. + (CometPartitioning::SinglePartition, vec![batch(0..64)], None), + // Multi-partition output is written at finalization unless a spill is forced. + (hash_partitioning(4), vec![batch(0..64)], None), + (hash_partitioning(4), vec![batch(0..64)], Some(1)), + // Empty-schema output is also written during finalization. + ( + CometPartitioning::RoundRobin(4, 0), + vec![empty_schema_batch(3)], + None, + ), + ] { + let marker = Arc::new(()); + let pusher = Arc::new(FailingPusher { + marker: Arc::clone(&marker), + calls: AtomicUsize::new(0), + }); + let exec = writer( + source(batches[0].schema(), &batches), + partitioning, + CompressionCodec::None, + rss(pusher.clone()), + max_buffer_bytes, + ); + let error = execute(&exec).await.unwrap_err(); + let DataFusionError::External(error) = error else { + panic!("callback error was wrapped or stringified: {error:?}"); + }; + let callback_error = error.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(&callback_error.0, &marker)); + assert_eq!(pusher.calls.load(Ordering::Relaxed), 1); + } +} From fe9c819359def8fe96f7340c144e2eae593e9fea Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:17:46 +0000 Subject: [PATCH 04/12] feat: bind task-owned Java pushers to native shuffle planning --- native/core/src/execution/jni_api.rs | 102 +++- native/core/src/execution/planner.rs | 251 +++++++++- .../src/execution/rss_planner_jni_tests.rs | 465 ++++++++++++++++++ native/proto/src/proto/operator.proto | 2 +- native/proto/src/proto/partitioning.proto | 4 +- .../src/writers/rss/rss_partition_writer.rs | 2 +- .../main/scala/org/apache/comet/Native.scala | 22 + 7 files changed, 824 insertions(+), 24 deletions(-) create mode 100644 native/core/src/execution/rss_planner_jni_tests.rs diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 2933e7569c6..b33251f81cd 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -21,8 +21,11 @@ use super::{serde, utils::SparkArrowConvert}; use crate::{ errors::{try_unwrap_or_throw, CometError, CometResult}, execution::{ - metrics::utils::update_comet_metric, planner::PhysicalPlanner, serde::to_arrow_datatype, - shuffle::spark_unsafe::row::process_sorted_row_partition, sort::RdxSort, + metrics::utils::update_comet_metric, + planner::{PhysicalPlanner, RegisteredShufflePusher}, + serde::to_arrow_datatype, + shuffle::spark_unsafe::row::process_sorted_row_partition, + sort::RdxSort, }, jvm_bridge::JVMClasses, }; @@ -41,7 +44,9 @@ use datafusion::{ physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream}, prelude::{SessionConfig, SessionContext}, }; +use datafusion_comet_jni_bridge::shuffle_partition_pusher::JavaShufflePartitionPusher; use datafusion_comet_proto::spark_operator::Operator; +use datafusion_comet_shuffle::PartitionPusher; use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl}; use datafusion_spark::function::array::array_contains::SparkArrayContains; use datafusion_spark::function::array::repeat::SparkArrayRepeat; @@ -120,6 +125,10 @@ use std::sync::OnceLock; #[cfg(feature = "jemalloc")] use tikv_jemalloc_ctl::{epoch, stats}; +#[cfg(test)] +#[path = "rss_planner_jni_tests.rs"] +mod rss_planner_jni_tests; + static TOKIO_RUNTIME: Mutex> = Mutex::new(None); #[cfg(feature = "jemalloc")] @@ -394,6 +403,9 @@ struct ExecutionContext { pub class_loader: Option>>>, /// Removes this context's tracing memory-pool entry on every exit path. memory_pool_registration: Option, + /// Optional RSS callback registered by this task before the physical plan is built. + /// The callback's global JNI references are released with this execution context. + pub rss_pusher: Option, } /// Accept serialized query plan and return the address of the native query plan. @@ -583,6 +595,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( task_context, class_loader, memory_pool_registration, + rss_pusher: None, }); Ok(Box::into_raw(exec_context) as i64) @@ -590,6 +603,88 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( }) } +/// Register an RSS callback on the native context that owns its Spark task attempt. +/// +/// Registration must happen after `createPlan` and before the first `executePlan`. The handle is +/// an opaque task-local identifier, not a native address or a process-wide registry key. +/// +/// # Safety +/// The execution context address must refer to a live native plan returned by `createPlan`. +#[no_mangle] +pub unsafe extern "system" fn Java_org_apache_comet_Native_registerRssPartitionPusher( + e: EnvUnowned, + _class: JClass, + exec_context: jlong, + handle: jlong, + object: JObject, + num_partitions: jint, + max_frame_bytes: jint, +) { + try_unwrap_or_throw(&e, |env| { + if exec_context == 0 { + return Err(CometError::NullPointer( + "Comet execution context is null".into(), + )); + } + + register_rss_partition_pusher( + env, + get_execution_context(exec_context), + handle, + &object, + num_partitions, + max_frame_bytes, + ) + }) +} + +fn register_rss_partition_pusher( + env: &mut Env, + exec_context: &mut ExecutionContext, + handle: i64, + object: &JObject<'_>, + num_partitions: jint, + max_frame_bytes: jint, +) -> CometResult<()> { + if handle <= 0 { + return Err(CometError::Config( + "RSS partition pusher handle must be positive".into(), + )); + } + if num_partitions <= 0 { + return Err(CometError::Config("Invalid RSS partition count".into())); + } + if max_frame_bytes <= 0 { + return Err(CometError::Config("Invalid RSS frame byte limit".into())); + } + if exec_context.root_op.is_some() { + return Err(CometError::Config( + "RSS partition pusher must be registered before execution".into(), + )); + } + if exec_context.rss_pusher.is_some() { + return Err(CometError::Config( + "RSS partition pusher is already registered for this task".into(), + )); + } + + let num_partitions = num_partitions as usize; + let max_frame_bytes = max_frame_bytes as usize; + let pusher: Arc = Arc::new(JavaShufflePartitionPusher::try_new( + env, + object, + num_partitions, + max_frame_bytes, + )?); + exec_context.rss_pusher = Some(RegisteredShufflePusher { + handle, + num_partitions, + max_frame_bytes, + pusher, + }); + Ok(()) +} + /// Configure DataFusion session context. fn prepare_datafusion_session_context( batch_size: usize, @@ -830,7 +925,8 @@ 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_rss_pusher(exec_context.rss_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 46bc1f365cf..e66008d78c0 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -44,7 +44,7 @@ use crate::execution::{ planner::expression_registry::ExpressionRegistry, planner::operator_registry::OperatorRegistry, serde::to_arrow_datatype, - shuffle::{SchemaAlignExec, ShuffleWriterExec}, + shuffle::{PartitionPusher, SchemaAlignExec, ShuffleDestination, ShuffleWriterExec}, }; use crate::jvm_bridge::{jni_call, JVMClasses}; use arrow::compute::CastOptions; @@ -244,6 +244,15 @@ pub struct BinaryExprOptions { pub const TEST_EXEC_CONTEXT_ID: i64 = -1; +/// An opaque RSS callback registration owned by one native task execution context. +#[derive(Clone)] +pub(crate) struct RegisteredShufflePusher { + pub(crate) handle: i64, + pub(crate) num_partitions: usize, + pub(crate) max_frame_bytes: usize, + pub(crate) pusher: Arc, +} + /// The query planner for converting Spark query plans to DataFusion query plans. pub struct PhysicalPlanner { // The execution context id of this planner. @@ -263,6 +272,8 @@ pub struct PhysicalPlanner { /// `ExecutionContext`; see that struct for the propagation rationale. `None` when no driving /// Spark task is available. class_loader: Option>>>, + /// Callback explicitly registered on this task's native execution context. + rss_pusher: Option, } impl Default for PhysicalPlanner { @@ -281,6 +292,7 @@ impl PhysicalPlanner { sql_text_pool: vec![], task_context: None, class_loader: None, + rss_pusher: None, } } @@ -381,6 +393,12 @@ impl PhysicalPlanner { self } + /// Attach the callback owned by this planner's Spark task attempt, if one was registered. + pub(crate) fn with_rss_pusher(mut self, pusher: Option) -> Self { + self.rss_pusher = pusher; + self + } + /// Return session context of this planner. pub fn session_ctx(&self) -> &Arc { &self.session_ctx @@ -1806,9 +1824,59 @@ impl PhysicalPlanner { ))) } OpStruct::ShuffleWriter(writer) => { - // Validate the destination before planning the child. In particular, an RSS or - // unknown destination must never fall through to the legacy local-file writer. - let (output_data_file, output_index_file) = local_shuffle_output_paths(writer)?; + use datafusion_comet_proto::spark_partitioning::partition_writer::PartitionWriterStruct; + + // Resolve task ownership before planning the child so malformed or unregistered + // RSS destinations can never fall through to legacy local-file output. + let rss_registration = match writer + .partition_writer + .as_ref() + .and_then(|partition_writer| partition_writer.partition_writer_struct.as_ref()) + { + Some(PartitionWriterStruct::RssPartitionWriter(rss)) => { + if !writer.output_data_file.is_empty() + || !writer.output_index_file.is_empty() + { + return Err(GeneralError( + "RSS shuffle partition writer must not specify legacy local paths" + .to_string(), + )); + } + if rss.rss_partition_pusher <= 0 { + return Err(GeneralError( + "RSS partition pusher handle must be positive".to_string(), + )); + } + let registered = self.rss_pusher.as_ref().ok_or_else(|| { + GeneralError( + "RSS partition pusher is not registered for this task".to_string(), + ) + })?; + if registered.handle != rss.rss_partition_pusher { + return Err(GeneralError( + "RSS partition pusher handle does not belong to this task" + .to_string(), + )); + } + Some(registered) + } + _ => None, + }; + + let shuffle_destination = if let Some(registered) = rss_registration { + ShuffleDestination::Rss { + pusher: Arc::clone(®istered.pusher), + max_frame_bytes: registered.max_frame_bytes, + } + } else { + let (output_data_file, output_index_file) = local_shuffle_output_paths(writer)?; + ShuffleDestination::Local { + output_data_file: output_data_file.to_owned(), + output_index_file: output_index_file.to_owned(), + write_buffer_size: writer.write_buffer_size as usize, + } + }; + assert_eq!(children.len(), 1); let (scans, shuffle_scans, child) = self.create_plan(&children[0], inputs, partition_count)?; @@ -1823,6 +1891,16 @@ impl PhysicalPlanner { writer_input.schema(), )?; + if let Some(registered) = rss_registration { + let partition_count = partitioning.partition_count(); + if registered.num_partitions != partition_count { + return Err(GeneralError(format!( + "RSS partition count {} does not match shuffle partition count {}", + registered.num_partitions, partition_count + ))); + } + } + let codec = match writer.codec.try_into() { Ok(SparkCompressionCodec::None) => Ok(CompressionCodec::None), Ok(SparkCompressionCodec::Snappy) => Ok(CompressionCodec::Snappy), @@ -1836,19 +1914,16 @@ impl PhysicalPlanner { ))), }?; - let write_buffer_size = writer.write_buffer_size as usize; // Zero on the wire means the limit is disabled; normalize it here so the writer // only ever sees a real limit or none at all. let max_buffer_bytes = (writer.max_buffer_bytes > 0).then_some(writer.max_buffer_bytes as usize); - let shuffle_writer = Arc::new(ShuffleWriterExec::try_new( + let shuffle_writer = Arc::new(ShuffleWriterExec::try_new_with_destination( writer_input, partitioning, codec, - output_data_file.to_owned(), - output_index_file.to_owned(), + shuffle_destination, writer.tracing_enabled, - write_buffer_size, max_buffer_bytes, )?); @@ -4677,7 +4752,12 @@ mod tests { mod shuffle_writer_destination { use super::create_scan; use crate::execution::operators::InputBatch; - use crate::execution::planner::{local_shuffle_output_paths, PhysicalPlanner}; + use crate::execution::planner::{ + local_shuffle_output_paths, PhysicalPlanner, RegisteredShufflePusher, + }; + use crate::execution::shuffle::PartitionPusher; + use arrow::array::Int32Array; + use datafusion::error::Result as DataFusionResult; use datafusion::physical_plan::common::collect; use datafusion::prelude::SessionContext; use datafusion_comet_proto::spark_operator::{operator::OpStruct, Operator, ShuffleWriter}; @@ -4686,9 +4766,49 @@ mod tests { LocalPartitionWriter, PartitionWriter, Partitioning, RssPartitionWriter, SinglePartition, }; + use futures::{poll, StreamExt}; use prost::Message; + use std::sync::{Arc, Mutex}; + use std::task::Poll; use tempfile::TempDir; + type RecordedFrames = Arc)>>>; + + struct RecordingPusher { + frames: RecordedFrames, + } + + impl PartitionPusher for RecordingPusher { + fn push_partition_data( + &self, + partition_id: usize, + frame: &[u8], + ) -> DataFusionResult<()> { + self.frames + .lock() + .unwrap() + .push((partition_id, frame.to_vec())); + Ok(()) + } + } + + fn registered_pusher( + handle: i64, + num_partitions: usize, + max_frame_bytes: usize, + ) -> (RegisteredShufflePusher, RecordedFrames) { + let frames = Arc::new(Mutex::new(vec![])); + let registration = RegisteredShufflePusher { + handle, + num_partitions, + max_frame_bytes, + pusher: Arc::new(RecordingPusher { + frames: Arc::clone(&frames), + }), + }; + (registration, frames) + } + fn local_destination(data: &str, index: &str) -> PartitionWriter { PartitionWriter { partition_writer_struct: Some(PartitionWriterStruct::LocalPartitionWriter( @@ -4708,6 +4828,24 @@ mod tests { } } + fn rss_writer(handle: i64) -> ShuffleWriter { + ShuffleWriter { + partitioning: Some(Partitioning { + partitioning_struct: Some(PartitioningStruct::SinglePartition( + SinglePartition {}, + )), + }), + partition_writer: Some(PartitionWriter { + partition_writer_struct: Some(PartitionWriterStruct::RssPartitionWriter( + RssPartitionWriter { + rss_partition_pusher: handle, + }, + )), + }), + ..Default::default() + } + } + fn assert_rejected_before_child(writer: ShuffleWriter, expected: &str) { let op = Operator { // An invalid child ensures destination validation happens before child planning. @@ -4717,8 +4855,7 @@ mod tests { }; let err = PhysicalPlanner::default() .create_plan(&op, &mut vec![], 1) - .err() - .expect("invalid destination must not produce a plan"); + .expect_err("invalid destination must not produce a plan"); assert!(err.to_string().contains(expected), "{err}"); } @@ -4768,7 +4905,7 @@ mod tests { } #[test] - fn rejects_rss_even_when_legacy_paths_are_present() { + fn rejects_unregistered_or_invalid_rss_before_planning_child() { for handle in [0, 7, -1] { for mut writer in [ShuffleWriter::default(), legacy_writer()] { writer.partition_writer = Some(PartitionWriter { @@ -4778,12 +4915,92 @@ mod tests { }, )), }); - assert_rejected_before_child( - writer, - "RSS shuffle partition writer is not supported", - ); + let expected = if !writer.output_data_file.is_empty() + || !writer.output_index_file.is_empty() + { + "must not specify legacy local paths" + } else if handle <= 0 { + "handle must be positive" + } else { + "not registered for this task" + }; + assert_rejected_before_child(writer, expected); + } + } + } + + #[test] + fn rejects_other_tasks_rss_handle_before_planning_child() { + let (registration, frames) = registered_pusher(7, 1, 1024); + let op = Operator { + children: vec![Operator::default()], + op_struct: Some(OpStruct::ShuffleWriter(rss_writer(8))), + ..Default::default() + }; + let error = PhysicalPlanner::default() + .with_rss_pusher(Some(registration)) + .create_plan(&op, &mut vec![], 1) + .expect_err("a different task's RSS pusher must be rejected"); + assert!(error.to_string().contains("does not belong to this task")); + assert!(frames.lock().unwrap().is_empty()); + } + + #[test] + fn rejects_incompatible_rss_partition_counts_and_frame_limits() { + for (registration, expected) in [ + (registered_pusher(7, 2, 1024).0, "partition count"), + (registered_pusher(7, 1, 19).0, "frame byte limit"), + ] { + let op = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::ShuffleWriter(rss_writer(7))), + ..Default::default() + }; + let error = PhysicalPlanner::default() + .with_rss_pusher(Some(registration)) + .create_plan(&op, &mut vec![], 1) + .expect_err("incompatible RSS options must be rejected"); + assert!(error.to_string().contains(expected), "{error}"); + } + } + + #[tokio::test] + async fn registered_rss_destination_pushes_complete_partition_frames() { + let (registration, frames) = registered_pusher(7, 1, 1024 * 1024); + let op = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::ShuffleWriter(rss_writer(7))), + ..Default::default() + }; + let planner = PhysicalPlanner::default().with_rss_pusher(Some(registration)); + let (mut scans, _, plan) = planner.create_plan(&op, &mut vec![], 1).unwrap(); + scans[0].set_input_batch(InputBatch::Batch( + vec![Arc::new(Int32Array::from(vec![11, 29, 47]))], + 3, + )); + + let mut stream = plan + .native_plan + .execute(0, SessionContext::new().task_ctx()) + .unwrap(); + let mut eof_sent = false; + loop { + match poll!(stream.next()) { + Poll::Ready(Some(Ok(_))) => panic!("RSS writer must not produce batches"), + Poll::Ready(Some(Err(error))) => panic!("RSS writer failed: {error}"), + Poll::Ready(None) => break, + Poll::Pending if !eof_sent => { + scans[0].set_input_batch(InputBatch::EOF); + eof_sent = true; + } + Poll::Pending => tokio::task::yield_now().await, } } + + let frames = frames.lock().unwrap(); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].0, 0); + assert!(frames[0].1.len() >= 20); } #[test] diff --git a/native/core/src/execution/rss_planner_jni_tests.rs b/native/core/src/execution/rss_planner_jni_tests.rs new file mode 100644 index 00000000000..ed191c6de73 --- /dev/null +++ b/native/core/src/execution/rss_planner_jni_tests.rs @@ -0,0 +1,465 @@ +// 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 std::path::PathBuf; +use std::process::Command; +use std::sync::Arc; +use std::task::Poll; +use std::time::{Duration, Instant}; + +use arrow::array::Int32Array; +use datafusion::common::DataFusionError; +use datafusion::prelude::SessionContext; +use datafusion_comet_proto::spark_expression; +use datafusion_comet_proto::spark_operator::{operator::OpStruct, Operator, Scan, ShuffleWriter}; +use datafusion_comet_proto::spark_partitioning::{ + partition_writer::PartitionWriterStruct, partitioning::PartitioningStruct, PartitionWriter, + Partitioning, RssPartitionWriter, SinglePartition, +}; +use futures::{poll, StreamExt}; +use jni::objects::{Global, JByteArray, JObject, JValue}; +use jni::{Env, EnvUnowned, InitArgsBuilder, JNIVersion, JavaVM}; + +use super::{ + parse_memory_pool_config, register_rss_partition_pusher, ExecutionContext, + Java_org_apache_comet_Native_registerRssPartitionPusher, +}; +use crate::errors::{CometError, CometResult}; +use crate::execution::operators::InputBatch; +use crate::execution::planner::{PhysicalPlanner, RegisteredShufflePusher}; +use crate::execution::shuffle::read_ipc_compressed; + +const FIRST_HANDLE: i64 = 41; +const SECOND_HANDLE: i64 = 42; +const TINY_FRAME_HANDLE: i64 = 43; +const MAX_FRAME_BYTES: i32 = 1024 * 1024; + +fn rss_plan(handle: i64) -> Operator { + Operator { + children: vec![Operator { + op_struct: Some(OpStruct::Scan(Scan { + fields: vec![spark_expression::DataType { + type_id: 3, + type_info: None, + }], + source: "task-owned RSS pusher test".to_string(), + })), + ..Default::default() + }], + op_struct: Some(OpStruct::ShuffleWriter(ShuffleWriter { + partitioning: Some(Partitioning { + partitioning_struct: Some(PartitioningStruct::SinglePartition(SinglePartition {})), + }), + partition_writer: Some(PartitionWriter { + partition_writer_struct: Some(PartitionWriterStruct::RssPartitionWriter( + RssPartitionWriter { + rss_partition_pusher: handle, + }, + )), + }), + write_buffer_size: 1024, + ..Default::default() + })), + ..Default::default() + } +} + +fn execution_context( + env: &mut Env, + spark_plan: Operator, + task_attempt_id: i64, +) -> CometResult { + let metrics = env.new_object(jni::jni_str!("java/lang/Object"), jni::jni_sig!("()V"), &[])?; + + Ok(ExecutionContext { + id: task_attempt_id, + task_attempt_id, + spark_plan, + partition_count: 1, + root_op: None, + scans: vec![], + shuffle_scans: vec![], + input_sources: vec![], + stream: None, + batch_receiver: None, + metrics: Arc::new(env.new_global_ref(&metrics)?), + metrics_update_interval: None, + metrics_last_update_time: Instant::now(), + poll_count_since_metrics_check: 0, + plan_creation_time: Duration::ZERO, + session_ctx: Arc::new(SessionContext::new()), + debug_native: false, + explain_native: false, + memory_pool_config: parse_memory_pool_config(false, "unbounded".to_string(), 0, 0)?, + tracing_enabled: false, + rust_thread_id: 0, + tracing_memory_metric_name: String::new(), + tracing_event_name: String::new(), + task_context: None, + rss_pusher: None, + }) +} + +fn assert_rejected_plan(plan: &Operator, registration: Option) { + let rejected = PhysicalPlanner::default() + .with_rss_pusher(registration) + .create_plan(plan, &mut vec![], 1); + assert!( + rejected.is_err(), + "invalid task-owned RSS plan was accepted" + ); +} + +fn execute_registered_plan( + context: &ExecutionContext, + values: &[i32], +) -> Result<(), DataFusionError> { + let planner = PhysicalPlanner::new(Arc::clone(&context.session_ctx), 0) + .with_rss_pusher(context.rss_pusher.clone()); + let (mut scans, _, plan) = planner + .create_plan(&context.spark_plan, &mut vec![], 1) + .expect("registered RSS plan should be executable"); + scans[0].set_input_batch(InputBatch::Batch( + vec![Arc::new(Int32Array::from(values.to_vec()))], + values.len(), + )); + + let mut stream = plan + .native_plan + .execute(0, context.session_ctx.task_ctx())?; + tokio::runtime::Runtime::new().unwrap().block_on(async { + let mut eof_sent = false; + loop { + match poll!(stream.next()) { + Poll::Ready(Some(Ok(_))) => panic!("RSS shuffle writer must not produce batches"), + Poll::Ready(Some(Err(error))) => return Err(error), + Poll::Ready(None) => return Ok(()), + Poll::Pending if !eof_sent => { + scans[0].set_input_batch(InputBatch::EOF); + eof_sent = true; + } + Poll::Pending => tokio::task::yield_now().await, + } + } + }) +} + +fn calls(vm: &JavaVM, fixture: &Global>) -> i32 { + vm.attach_current_thread(|env| -> jni::errors::Result { + env.get_field(fixture, jni::jni_str!("calls"), jni::jni_sig!("I"))? + .i() + }) + .unwrap() +} + +#[test] +#[cfg_attr(miri, ignore)] // Miri cannot launch a JVM. +fn task_owned_rss_plan_pushes_complete_frames_to_java() { + let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let classes = tempfile::tempdir().unwrap(); + let javac = std::env::var_os("JAVA_HOME") + .map(|java_home| PathBuf::from(java_home).join("bin/javac")) + .unwrap_or_else(|| PathBuf::from("javac")); + let output = + Command::new(javac) + .arg("-d") + .arg(classes.path()) + .arg(manifest.join( + "../../spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java", + )) + .arg(manifest.join("../jni-bridge/tests/java/RecordingShufflePartitionPusher.java")) + .output() + .expect("compile the task-owned RSS Java callback fixture"); + assert!( + output.status.success(), + "javac: {}", + String::from_utf8_lossy(&output.stderr) + ); + + let vm = JavaVM::new( + InitArgsBuilder::new() + .version(JNIVersion::V1_8) + .option("-Xcheck:jni") + .option(format!("-Djava.class.path={}", classes.path().display())) + .build() + .unwrap(), + ) + .unwrap(); + + let (first_context, second_context, tiny_context, first_fixture, second_fixture, tiny_fixture) = + vm.attach_current_thread(|env| -> CometResult<_> { + let first = env.new_object( + jni::jni_str!("org/apache/comet/shuffle/RecordingShufflePartitionPusher"), + jni::jni_sig!("()V"), + &[], + )?; + let second = env.new_object( + jni::jni_str!("org/apache/comet/shuffle/RecordingShufflePartitionPusher"), + jni::jni_sig!("()V"), + &[], + )?; + let tiny = env.new_object( + jni::jni_str!("org/apache/comet/shuffle/RecordingShufflePartitionPusher"), + jni::jni_sig!("()V"), + &[], + )?; + let wrong_type = + env.new_object(jni::jni_str!("java/lang/Object"), jni::jni_sig!("()V"), &[])?; + + let mut first_context = execution_context(env, rss_plan(FIRST_HANDLE), 101)?; + for (handle, partitions, max_bytes) in [ + (0, 1, MAX_FRAME_BYTES), + (-1, 1, MAX_FRAME_BYTES), + (FIRST_HANDLE, 0, MAX_FRAME_BYTES), + (FIRST_HANDLE, -1, MAX_FRAME_BYTES), + (FIRST_HANDLE, 1, 0), + (FIRST_HANDLE, 1, -1), + ] { + assert!(register_rss_partition_pusher( + env, + &mut first_context, + handle, + &first, + partitions, + max_bytes, + ) + .is_err()); + assert!(first_context.rss_pusher.is_none()); + } + assert!(register_rss_partition_pusher( + env, + &mut first_context, + FIRST_HANDLE, + &JObject::null(), + 1, + MAX_FRAME_BYTES, + ) + .is_err()); + assert!(register_rss_partition_pusher( + env, + &mut first_context, + FIRST_HANDLE, + &wrong_type, + 1, + MAX_FRAME_BYTES, + ) + .is_err()); + assert!(first_context.rss_pusher.is_none()); + + let registration_class = env.find_class(jni::jni_str!("java/lang/Object"))?; + let callback = env.new_local_ref(&first)?; + // SAFETY: this is the current thread's live JNI attachment; a null context must be + // rejected at the exported boundary without dereferencing it. + unsafe { + Java_org_apache_comet_Native_registerRssPartitionPusher( + EnvUnowned::from_raw(env.get_raw()), + registration_class, + 0, + FIRST_HANDLE, + callback, + 1, + MAX_FRAME_BYTES, + ); + } + let null_context_exception = env + .exception_occurred() + .expect("null execution context must raise a Java exception"); + env.exception_clear(); + let null_pointer_exception = + env.find_class(jni::jni_str!("java/lang/NullPointerException"))?; + assert!(env.is_instance_of(&null_context_exception, &null_pointer_exception)?); + assert!(first_context.rss_pusher.is_none()); + + let registration_class = env.find_class(jni::jni_str!("java/lang/Object"))?; + let callback = env.new_local_ref(&first)?; + let context_address = &mut first_context as *mut ExecutionContext as i64; + // SAFETY: the JNI attachment, callback, and execution context all remain live for + // the synchronous exported registration call. + unsafe { + Java_org_apache_comet_Native_registerRssPartitionPusher( + EnvUnowned::from_raw(env.get_raw()), + registration_class, + context_address, + FIRST_HANDLE, + callback, + 1, + MAX_FRAME_BYTES, + ); + } + assert!( + !env.exception_check(), + "valid RSS pusher registration raised a Java exception" + ); + assert!(register_rss_partition_pusher( + env, + &mut first_context, + SECOND_HANDLE, + &second, + 1, + MAX_FRAME_BYTES, + ) + .is_err()); + assert_eq!( + first_context.rss_pusher.as_ref().unwrap().handle, + FIRST_HANDLE + ); + + let mut second_context = execution_context(env, rss_plan(SECOND_HANDLE), 102)?; + register_rss_partition_pusher( + env, + &mut second_context, + SECOND_HANDLE, + &second, + 1, + MAX_FRAME_BYTES, + )?; + + // The JNI bridge accepts a positive byte limit, but the destination writer rejects + // anything too small to contain the mandatory 20-byte Comet frame header. + let mut tiny_context = execution_context(env, rss_plan(TINY_FRAME_HANDLE), 103)?; + register_rss_partition_pusher(env, &mut tiny_context, TINY_FRAME_HANDLE, &tiny, 1, 19)?; + + Ok(( + first_context, + second_context, + tiny_context, + env.new_global_ref(&first)?, + env.new_global_ref(&second)?, + env.new_global_ref(&tiny)?, + )) + }) + .unwrap(); + + assert_rejected_plan(&first_context.spark_plan, None); + assert_rejected_plan(&first_context.spark_plan, second_context.rss_pusher.clone()); + assert_rejected_plan(&second_context.spark_plan, first_context.rss_pusher.clone()); + for invalid_handle in [0, -1, SECOND_HANDLE] { + assert_rejected_plan(&rss_plan(invalid_handle), first_context.rss_pusher.clone()); + } + + let mut wrong_partition_count = first_context.rss_pusher.clone().unwrap(); + wrong_partition_count.num_partitions = 2; + assert_rejected_plan(&first_context.spark_plan, Some(wrong_partition_count)); + assert_rejected_plan(&tiny_context.spark_plan, tiny_context.rss_pusher.clone()); + + // Registration cannot alter an execution context after its physical plan is initialized. + let (_, _, planned_root) = PhysicalPlanner::new(Arc::clone(&first_context.session_ctx), 0) + .with_rss_pusher(first_context.rss_pusher.clone()) + .create_plan(&first_context.spark_plan, &mut vec![], 1) + .unwrap(); + let mut late_context = vm + .attach_current_thread(|env| execution_context(env, rss_plan(FIRST_HANDLE), 104)) + .unwrap(); + late_context.root_op = Some(planned_root); + vm.attach_current_thread(|env| -> CometResult<()> { + let error = register_rss_partition_pusher( + env, + &mut late_context, + FIRST_HANDLE, + first_fixture.as_obj(), + 1, + MAX_FRAME_BYTES, + ) + .expect_err("registration after physical planning must be rejected"); + assert!(error.to_string().contains("before execution"), "{error}"); + assert!(late_context.rss_pusher.is_none()); + Ok(()) + }) + .unwrap(); + + assert_eq!(calls(&vm, &first_fixture), 0); + assert_eq!(calls(&vm, &second_fixture), 0); + assert_eq!(calls(&vm, &tiny_fixture), 0); + + let expected_values = [11, 29, 47]; + execute_registered_plan(&first_context, &expected_values).unwrap(); + assert_eq!(calls(&vm, &first_fixture), 1); + assert_eq!(calls(&vm, &second_fixture), 0); + assert_eq!(calls(&vm, &tiny_fixture), 0); + + let frame = vm + .attach_current_thread(|env| -> jni::errors::Result> { + assert_eq!( + env.get_field( + &first_fixture, + jni::jni_str!("partitionId"), + jni::jni_sig!("I"), + )? + .i()?, + 0 + ); + let bytes = env + .get_field( + &first_fixture, + jni::jni_str!("lastBytes"), + jni::jni_sig!("[B"), + )? + .l()?; + // SAFETY: the fixture field has descriptor byte[] and into_raw transfers its local ref. + let bytes = unsafe { JByteArray::from_raw(env, bytes.into_raw()) }; + env.convert_byte_array(&bytes) + }) + .unwrap(); + assert!(frame.len() >= 20); + assert_eq!( + u64::from_le_bytes(frame[..8].try_into().unwrap()) + 8, + frame.len() as u64 + ); + assert_eq!(u64::from_le_bytes(frame[8..16].try_into().unwrap()), 1); + let decoded = read_ipc_compressed(&frame[16..]).unwrap(); + let decoded_values = decoded + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(decoded_values.values().as_ref(), expected_values.as_slice()); + + // A Java callback failure must survive RSS writing as the original throwable, not as a + // flattened string or a newly constructed generic native exception. + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + env.set_field( + &second_fixture, + jni::jni_str!("failureMode"), + jni::jni_sig!("I"), + JValue::Int(1), + ) + }) + .unwrap(); + let error = execute_registered_plan(&second_context, &expected_values).unwrap_err(); + let DataFusionError::External(source) = error else { + panic!("Java callback failure lost its typed DataFusion wrapper: {error}"); + }; + let Some(CometError::JavaException { throwable, .. }) = source.downcast_ref::() + else { + panic!("Java callback failure lost its original throwable: {source}"); + }; + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + let expected = env + .get_field( + &second_fixture, + jni::jni_str!("failure"), + jni::jni_sig!("Ljava/io/IOException;"), + )? + .l()?; + assert!(env.is_same_object(&expected, throwable)?); + Ok(()) + }) + .unwrap(); + assert_eq!(calls(&vm, &first_fixture), 1); + assert_eq!(calls(&vm, &second_fixture), 1); + assert_eq!(calls(&vm, &tiny_fixture), 0); +} diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index aaf36b33328..e7b331555a7 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -690,7 +690,7 @@ message ShuffleWriter { // Absent means legacy local output using fields 3 and 4. An explicit local writer may // duplicate those paths for older native libraries, but the paths must agree. // Do not emit RSS until the native library's RSS protocol support has been verified. - // RSS must not populate the legacy local paths, and is not executable yet. + // RSS must not populate the legacy local paths and requires a task-owned pusher registration. spark.spark_partitioning.PartitionWriter partition_writer = 2; // Retained for compatibility with existing JVM and native plans. string output_data_file = 3; diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index bb3acef7ccc..5309d2ebd99 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -74,7 +74,7 @@ message LocalPartitionWriter { } message RssPartitionWriter { - // Opaque, task-scoped pusher handle. This schema does not make the handle usable: - // RSS plans remain unsupported until capability negotiation and safe handle ownership exist. + // Opaque, task-scoped pusher handle. Native planning resolves this handle only against + // the pusher registered on the owning execution context. int64 rss_partition_pusher = 1; } diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs index b487c9741e2..357e2ef819c 100644 --- a/native/shuffle/src/writers/rss/rss_partition_writer.rs +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -49,7 +49,7 @@ impl PartitionPusher for JavaShufflePartitionPusher { } /// Encodes already-partitioned batches using the existing Comet format and sends one frame per -/// callback. This foundation is not selected by the native planner yet. +/// callback. The native planner selects it only for a registered task-owned pusher. /// /// There are no retained per-reducer buffers. The byte limit caps encoded output, not Arrow's /// encoding scratch space or the backend's asynchronous memory. Production admission, row diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 51b3e5e41be..8da705c7710 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,27 @@ class Native extends NativeBase { classLoader: ClassLoader): Long // scalastyle:on + /** + * Register a task-owned remote shuffle callback before the native plan is first executed. + * + * @param plan + * the native execution context returned by createPlan + * @param handle + * the positive, opaque handle serialized in this task's remote shuffle plan + * @param pusher + * the callback that accepts complete encoded shuffle frames + * @param numPartitions + * the number of remote shuffle output partitions + * @param maxFrameBytes + * the maximum encoded bytes accepted for one complete shuffle frame + */ + @native def registerRssPartitionPusher( + plan: Long, + handle: Long, + pusher: ShufflePartitionPusher, + numPartitions: Int, + maxFrameBytes: Int): Unit + /** * Execute a native query plan based on given input Arrow arrays. * From b08351b2ce44a54cef2655a9b7289319ecb85af6 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:18:14 +0000 Subject: [PATCH 05/12] feat: add Celeborn shuffle manager and partition pusher --- .../CelebornShufflePartitionPusher.java | 183 +++++++++ .../comet/CometSparkSessionExtensions.scala | 22 +- .../CelebornShufflePusherFactory.scala | 98 +++++ .../shuffle/CometCelebornShuffleManager.scala | 148 +++++++ .../shuffle/CometShuffleExchangeExec.scala | 10 +- .../CometSparkSessionExtensionsSuite.scala | 60 ++- .../CelebornShufflePartitionPusherSuite.scala | 368 ++++++++++++++++++ .../CometCelebornShuffleManagerSuite.scala | 227 +++++++++++ 8 files changed, 1103 insertions(+), 13 deletions(-) create mode 100644 spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java create mode 100644 spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala create mode 100644 spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java new file mode 100644 index 00000000000..61cc5abe326 --- /dev/null +++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java @@ -0,0 +1,183 @@ +/* + * 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; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; + +/** Sends complete native Comet frames through an existing, task-scoped Celeborn shuffle client. */ +public final class CelebornShufflePartitionPusher implements ShufflePartitionPusher { + + // Celeborn prefixes every accepted payload with four transport-level integers. + private static final int CELEBORN_BATCH_HEADER_BYTES = 4 * Integer.BYTES; + + private final Object shuffleClient; + private final Method pushOrMergeData; + private final int shuffleId; + private final int mapId; + private final int encodedAttemptId; + private final int numMappers; + private final int numPartitions; + + /** + * Binds the existing Celeborn client and all shuffle identity to one map task. + * + *

The client is deliberately accepted as an {@link Object}: Celeborn is provided by the Spark + * executor, and Comet must remain usable when the optional Celeborn client is absent. Its public + * raw-push method is resolved once so an unsupported client fails before any data is written. + * + * @param shuffleClient the existing Celeborn {@code ShuffleClientImpl} for this application + * @param shuffleId the Celeborn shuffle ID, which can differ from the Spark shuffle ID + * @param mapId the logical map-partition index + * @param encodedAttemptId the combined Spark stage and task attempt number + * @param numMappers the number of map tasks in the shuffle + * @param numPartitions the number of output reduce partitions + */ + public CelebornShufflePartitionPusher( + Object shuffleClient, + int shuffleId, + int mapId, + int encodedAttemptId, + int numMappers, + int numPartitions) { + if (shuffleClient == null) { + throw new IllegalArgumentException("Celeborn shuffle client must not be null"); + } + if (shuffleId < 0) { + throw new IllegalArgumentException("Celeborn shuffle ID must not be negative"); + } + if (mapId < 0) { + throw new IllegalArgumentException("Celeborn map ID must not be negative"); + } + if (encodedAttemptId < 0) { + throw new IllegalArgumentException("Celeborn encoded attempt ID must not be negative"); + } + if (numMappers <= 0) { + throw new IllegalArgumentException("Celeborn mapper count must be positive"); + } + if (mapId >= numMappers) { + throw new IllegalArgumentException("Celeborn map ID is outside the mapper count"); + } + if (numPartitions <= 0) { + throw new IllegalArgumentException("Celeborn partition count must be positive"); + } + + Method method; + try { + method = + shuffleClient + .getClass() + .getMethod( + "pushOrMergeData", + int.class, + int.class, + int.class, + int.class, + byte[].class, + int.class, + int.class, + int.class, + int.class, + boolean.class, + boolean.class); + } catch (NoSuchMethodException | SecurityException e) { + throw new IllegalArgumentException( + "Celeborn shuffle client does not provide the required public raw-push API", e); + } + if (method.getReturnType() != int.class || Modifier.isStatic(method.getModifiers())) { + throw new IllegalArgumentException( + "Celeborn raw-push API must be an instance method returning an int"); + } + + this.shuffleClient = shuffleClient; + this.pushOrMergeData = method; + this.shuffleId = shuffleId; + this.mapId = mapId; + this.encodedAttemptId = encodedAttemptId; + this.numMappers = numMappers; + this.numPartitions = numPartitions; + } + + @Override + public int pushPartitionData(int partitionId, byte[] bytes, int length) throws IOException { + if (partitionId < 0 || partitionId >= numPartitions) { + throw new IOException("Celeborn output partition is outside this task's partition count"); + } + if (bytes == null) { + throw new IOException("Celeborn shuffle frame must not be null"); + } + if (length <= 0 || length > bytes.length) { + throw new IOException("Celeborn shuffle frame length must be within the supplied bytes"); + } + if (length > Integer.MAX_VALUE - CELEBORN_BATCH_HEADER_BYTES) { + throw new IOException("Celeborn shuffle frame and transport header exceed the byte limit"); + } + + final int accepted; + try { + // doPush=true sends this complete frame immediately. skipCompress=true preserves the + // existing Comet frame, which already owns its compression and framing format. + accepted = + (int) + pushOrMergeData.invoke( + shuffleClient, + shuffleId, + mapId, + encodedAttemptId, + partitionId, + bytes, + 0, + length, + numMappers, + numPartitions, + true, + true); + } catch (IllegalAccessException e) { + throw new IOException("Cannot invoke the public Celeborn raw-push API", e); + } catch (InvocationTargetException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) { + throw (IOException) cause; + } + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new IOException("Celeborn raw shuffle push failed", cause); + } + + int expected = length + CELEBORN_BATCH_HEADER_BYTES; + if (accepted != expected) { + throw new IOException( + "Celeborn raw shuffle push accepted " + + accepted + + " bytes; expected " + + expected + + " including its transport header"); + } + + // ShufflePartitionPusher reports Comet payload bytes, never Celeborn transport bytes. + return length; + } +} diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 4dc36194331..6232b330f25 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.{SparkSession, SparkSessionExtensions} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.{TreeNode, TreeNodeTag} import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleManager} import org.apache.spark.sql.execution._ import org.apache.spark.sql.internal.SQLConf @@ -121,6 +122,8 @@ class CometSparkSessionExtensions object CometSparkSessionExtensions extends Logging { lazy val isBigEndian: Boolean = ByteOrder.nativeOrder().equals(ByteOrder.BIG_ENDIAN) + private val SHUFFLE_MANAGER_KEY = "spark.shuffle.manager" + /** * Checks whether Comet extension should be loaded for Spark. */ @@ -137,7 +140,8 @@ object CometSparkSessionExtensions extends Logging { if (COMET_SHUFFLE_ENABLED.get(conf) && !isCometShuffleManagerEnabled(conf)) { logWarning( "Comet extension is disabled because spark.shuffle.manager is not set to " + - "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager. " + + s"${classOf[CometShuffleManager].getName} or " + + s"${classOf[CometCelebornShuffleManager].getName}. " + "Comet provides limited benefit without its shuffle manager. " + s"Set ${COMET_SHUFFLE_ENABLED.key}=false to keep Comet enabled with " + "Spark's default shuffle manager.") @@ -173,16 +177,18 @@ object CometSparkSessionExtensions extends Logging { } } - // Check whether Comet shuffle is enabled: - // 1. `COMET_SHUFFLE_ENABLED` is true - // 2. `spark.shuffle.manager` is set to `CometShuffleManager` - // 3. Off-heap memory is enabled || Spark/Comet unit testing + // The Celeborn manager is valid for loading Comet, but its native shuffle writer and reader + // are not wired yet. Keep exchanges on Celeborn's existing Spark shuffle path until they are. def isCometShuffleEnabled(conf: SQLConf): Boolean = - COMET_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) + COMET_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) && + conf.getConfString(SHUFFLE_MANAGER_KEY) != classOf[CometCelebornShuffleManager].getName def isCometShuffleManagerEnabled(conf: SQLConf): Boolean = { - conf.contains("spark.shuffle.manager") && conf.getConfString("spark.shuffle.manager") == - "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" + conf.contains(SHUFFLE_MANAGER_KEY) && { + val manager = conf.getConfString(SHUFFLE_MANAGER_KEY) + manager == classOf[CometShuffleManager].getName || + manager == classOf[CometCelebornShuffleManager].getName + } } def isCometScan(op: SparkPlan): Boolean = { diff --git a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala new file mode 100644 index 00000000000..25cb3383cce --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala @@ -0,0 +1,98 @@ +/* + * 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 org.apache.spark.{SparkConf, TaskContext} + +/** Creates task-owned Celeborn pushers using the application's existing Spark configuration. */ +object CelebornShufflePusherFactory { + + private val CELEBORN_ENABLED_KEY = "spark.comet.celeborn.enabled" + private val SHUFFLE_MANAGER_KEY = "spark.shuffle.manager" + private val SHUFFLE_DATA_IO_KEY = "spark.shuffle.sort.io.plugin.class" + private val CELEBORN_MASTER_ENDPOINTS_KEY = "spark.celeborn.master.endpoints" + + private val CELEBORN_SHUFFLE_MANAGER = + "org.apache.spark.shuffle.celeborn.SparkShuffleManager" + private val COMET_CELEBORN_SHUFFLE_MANAGER = + "org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager" + private val CELEBORN_SHUFFLE_DATA_IO = + "org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO" + + private val MAX_STAGE_ATTEMPTS = 1 << 15 + private val MAX_TASK_ATTEMPTS = 1 << 16 + + /** Detects resolved Celeborn configuration while honoring an explicit application opt-out. */ + def isEnabled(conf: SparkConf): Boolean = { + conf.getBoolean(CELEBORN_ENABLED_KEY, true) && + (conf + .getOption(SHUFFLE_MANAGER_KEY) + .exists(manager => + manager == CELEBORN_SHUFFLE_MANAGER || manager == COMET_CELEBORN_SHUFFLE_MANAGER) || + conf.getOption(SHUFFLE_DATA_IO_KEY).contains(CELEBORN_SHUFFLE_DATA_IO) || + conf.getOption(CELEBORN_MASTER_ENDPOINTS_KEY).exists(_.trim.nonEmpty)) + } + + /** Match Celeborn's stage/task attempt packing without depending on its Spark client jar. */ + private[shuffle] def encodeAttemptNumber(stageAttempt: Int, taskAttempt: Int): Int = { + require( + stageAttempt >= 0 && stageAttempt < MAX_STAGE_ATTEMPTS, + s"Celeborn stage attempt must be between 0 and ${MAX_STAGE_ATTEMPTS - 1}: " + + stageAttempt) + require( + taskAttempt >= 0 && taskAttempt < MAX_TASK_ATTEMPTS, + s"Celeborn task attempt must be between 0 and ${MAX_TASK_ATTEMPTS - 1}: " + + taskAttempt) + + (stageAttempt << 16) | taskAttempt + } + + /** + * Bind an existing Celeborn client to one Spark map-task attempt. + * + * The caller supplies the already-resolved Celeborn shuffle ID; it may differ from Spark's + * shuffle ID after a stage retry. Task metadata is captured here because native callbacks can + * execute on threads where Spark's thread-local TaskContext is unavailable. + */ + def create( + conf: SparkConf, + client: AnyRef, + celebornShuffleId: Int, + numMappers: Int, + numPartitions: Int, + taskContext: TaskContext): Option[ShufflePartitionPusher] = { + if (!isEnabled(conf)) { + None + } else { + val mapId = taskContext.partitionId() + val stageAttempt = taskContext.stageAttemptNumber() + val taskAttempt = taskContext.attemptNumber() + val encodedAttempt = encodeAttemptNumber(stageAttempt, taskAttempt) + Some( + new CelebornShufflePartitionPusher( + client, + celebornShuffleId, + mapId, + encodedAttempt, + numMappers, + numPartitions)) + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala new file mode 100644 index 00000000000..d4d92cbe39f --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -0,0 +1,148 @@ +/* + * 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.spark.sql.comet.execution.shuffle + +import java.lang.reflect.InvocationTargetException + +import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext} +import org.apache.spark.shuffle.{ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} + +import org.apache.comet.util.ClassLoaders + +/** + * Lets Comet execution coexist with the application's existing Celeborn shuffle manager. + * + * Ordinary Spark shuffle dependencies are owned entirely by Celeborn. Native and JVM Comet + * shuffle dependencies remain unsupported until their remote writer, reader, and task lifecycle + * are integrated; rejecting them prevents either local-disk fallback or handing a native Comet + * input iterator to Celeborn's row writer. + * + * Celeborn is loaded reflectively because its client is an optional, application-provided + * dependency rather than part of Comet's compile-time or runtime distribution. + */ +class CometCelebornShuffleManager private[shuffle] ( + conf: SparkConf, + isDriver: Boolean, + backendFactory: (SparkConf, Boolean) => ShuffleManager) + extends ShuffleManager { + + /** Constructor selected by Spark for driver and executor shuffle managers. */ + def this(conf: SparkConf, isDriver: Boolean) = + this(conf, isDriver, CometCelebornShuffleManager.createBackend) + + private val celebornManager = Option(backendFactory(conf, isDriver)).getOrElse { + throw new IllegalStateException("Celeborn Spark shuffle manager factory returned null") + } + + override def registerShuffle[K, V, C]( + shuffleId: Int, + dependency: ShuffleDependency[K, V, C]): ShuffleHandle = { + dependency match { + case _: CometShuffleDependency[_, _, _] => rejectCometShuffle() + case _ => celebornManager.registerShuffle(shuffleId, dependency) + } + } + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { + rejectCometHandle(handle) + celebornManager.getWriter(handle, mapId, context, metrics) + } + + override def getReader[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + rejectCometHandle(handle) + celebornManager.getReader( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) + } + + override def shuffleBlockResolver: ShuffleBlockResolver = + celebornManager.shuffleBlockResolver + + override def unregisterShuffle(shuffleId: Int): Boolean = + celebornManager.unregisterShuffle(shuffleId) + + override def stop(): Unit = celebornManager.stop() + + private def rejectCometHandle(handle: ShuffleHandle): Unit = handle match { + case _: CometNativeShuffleHandle[_, _] => rejectCometShuffle() + case _: CometBypassMergeSortShuffleHandle[_, _] => rejectCometShuffle() + case _: CometSerializedShuffleHandle[_, _] => rejectCometShuffle() + case _ => + } + + private def rejectCometShuffle(): Nothing = { + throw new UnsupportedOperationException( + "Comet shuffle over Celeborn is not supported yet; its remote writer, reader, " + + "and task lifecycle must be integrated before Comet shuffle can be enabled") + } +} + +private[shuffle] object CometCelebornShuffleManager { + + private val CELEBORN_MANAGER_CLASS = "org.apache.spark.shuffle.celeborn.SparkShuffleManager" + + private[shuffle] def createBackend(conf: SparkConf, isDriver: Boolean): ShuffleManager = { + try { + val managerClass = ClassLoaders.loadClass(CELEBORN_MANAGER_CLASS) + if (!classOf[ShuffleManager].isAssignableFrom(managerClass)) { + throw new IllegalStateException( + "Celeborn Spark shuffle manager does not implement ShuffleManager: " + + CELEBORN_MANAGER_CLASS) + } + + val constructor = managerClass.getConstructor(classOf[SparkConf], java.lang.Boolean.TYPE) + constructor.newInstance(conf, Boolean.box(isDriver)).asInstanceOf[ShuffleManager] + } catch { + case error: ClassNotFoundException => + throw new IllegalStateException( + s"Celeborn Spark shuffle manager is not available: $CELEBORN_MANAGER_CLASS. " + + "Ensure the Celeborn Spark client is present on the application classpath", + error) + case error: InvocationTargetException => + throw new IllegalStateException( + s"Could not initialize Celeborn Spark shuffle manager: $CELEBORN_MANAGER_CLASS", + Option(error.getCause).getOrElse(error)) + case error: ReflectiveOperationException => + throw new IllegalStateException( + s"Could not construct Celeborn Spark shuffle manager: $CELEBORN_MANAGER_CLASS", + error) + case error: LinkageError => + throw new IllegalStateException( + s"Could not load Celeborn Spark shuffle manager: $CELEBORN_MANAGER_CLASS", + error) + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 32fe282c955..778c3013172 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -51,7 +51,7 @@ import com.google.common.base.Objects import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometConf.{COMET_SHUFFLE_ENABLED, COMET_SHUFFLE_MODE} -import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isCometShuffleManagerEnabled, withFallbackReasons} +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isCometShuffleEnabled, isCometShuffleManagerEnabled, withFallbackReasons} import org.apache.comet.serde.{Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.{CometTypeShim, ShimCometShuffleExchangeExec} @@ -676,7 +676,13 @@ object CometShuffleExchangeExec if (!COMET_SHUFFLE_ENABLED.get(op.conf)) { Some(s"Comet shuffle is not enabled: ${COMET_SHUFFLE_ENABLED.key} is not enabled") } else if (!isCometShuffleManagerEnabled(op.conf)) { - Some(s"spark.shuffle.manager is not set to ${classOf[CometShuffleManager].getName}") + Some( + "spark.shuffle.manager is not set to " + + s"${classOf[CometShuffleManager].getName} or " + + classOf[CometCelebornShuffleManager].getName) + } else if (!isCometShuffleEnabled(op.conf)) { + Some( + "Celeborn-backed Comet shuffle is unavailable until its native writer and reader are wired") } else { None } diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 846fd633a3e..d4a2586b7a4 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -21,6 +21,9 @@ package org.apache.comet import org.apache.spark.SparkConf import org.apache.spark.sql._ +import org.apache.spark.sql.catalyst.plans.physical.SinglePartition +import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleExchangeExec, CometShuffleManager} +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.internal.SQLConf class CometSparkSessionExtensionsSuite extends CometTestBase { @@ -53,7 +56,7 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { NativeBase.setLoaded(true) } - test("isCometLoaded requires CometShuffleManager when shuffle.enabled=true") { + test("isCometLoaded requires a supported Comet shuffle manager when shuffle.enabled=true") { val conf = new SQLConf conf.setConfString(CometConf.COMET_ENABLED.key, "true") @@ -66,10 +69,61 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { // shuffle.enabled=true with the Comet shuffle manager registered: Comet should load. conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") + conf.setConfString("spark.shuffle.manager", classOf[CometShuffleManager].getName) + assert(isCometLoaded(conf)) + assert(isCometShuffleManagerEnabled(conf)) + assert(isCometShuffleEnabled(conf)) + } + + test("Celeborn manager loads Comet without enabling its unfinished native shuffle transport") { + val conf = new SQLConf + conf.setConfString(CometConf.COMET_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") + conf.setConfString("spark.shuffle.manager", classOf[CometCelebornShuffleManager].getName) + + assert(isCometShuffleManagerEnabled(conf)) + assert(isCometLoaded(conf)) + assert(!isCometShuffleEnabled(conf)) + } + + test("Celeborn manager leaves shuffle exchanges on its existing Spark shuffle path") { + val child = spark.emptyDataFrame.queryExecution.executedPlan + val session = spark.newSession() + val conf = session.sessionState.conf + conf.setConfString(CometConf.COMET_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") + conf.setConfString("spark.shuffle.manager", classOf[CometCelebornShuffleManager].getName) + + val previousActiveSession = SparkSession.getActiveSession + try { + SparkSession.setActiveSession(session) + val shuffle = ShuffleExchangeExec(SinglePartition, child) + + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).isEmpty) + assert( + shuffle + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + .exists(_.contains("Celeborn-backed Comet shuffle is unavailable"))) + } finally { + previousActiveSession match { + case Some(previousSession) => SparkSession.setActiveSession(previousSession) + case None => SparkSession.clearActiveSession() + } + } + } + + test("stock Celeborn manager does not satisfy Comet shuffle-manager requirements") { + val conf = new SQLConf + conf.setConfString(CometConf.COMET_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") conf.setConfString( "spark.shuffle.manager", - "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") - assert(isCometLoaded(conf)) + "org.apache.spark.shuffle.celeborn.SparkShuffleManager") + + assert(!isCometShuffleManagerEnabled(conf)) + assert(!isCometShuffleEnabled(conf)) + assert(!isCometLoaded(conf)) } test("Arrow properties") { diff --git a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala new file mode 100644 index 00000000000..abc89b1e2c8 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala @@ -0,0 +1,368 @@ +/* + * 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 +import java.util.concurrent.atomic.AtomicReference + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.{SparkConf, TaskContext} + +/** Public so the production adapter can invoke its method through ordinary Java reflection. */ +final class RecordingCelebornShuffleClient { + + @volatile var acceptedBytes: Option[Int] = None + @volatile var failure: Throwable = _ + @volatile var observedTaskContext: TaskContext = _ + @volatile var lastPush: RecordedCelebornPush = _ + + @throws[IOException] + def pushOrMergeData( + shuffleId: Int, + mapId: Int, + attemptId: Int, + partitionId: Int, + bytes: Array[Byte], + offset: Int, + length: Int, + numMappers: Int, + numPartitions: Int, + doPush: Boolean, + skipCompress: Boolean): Int = { + observedTaskContext = TaskContext.get() + lastPush = RecordedCelebornPush( + shuffleId, + mapId, + attemptId, + partitionId, + bytes, + offset, + length, + numMappers, + numPartitions, + doPush, + skipCompress) + + if (failure != null) { + throw failure + } + + acceptedBytes.getOrElse(length + 16) + } +} + +final case class RecordedCelebornPush( + shuffleId: Int, + mapId: Int, + attemptId: Int, + partitionId: Int, + bytes: Array[Byte], + offset: Int, + length: Int, + numMappers: Int, + numPartitions: Int, + doPush: Boolean, + skipCompress: Boolean) + +/** The method exists, but its return type is incompatible with Celeborn's raw-push contract. */ +final class WrongReturnTypeCelebornShuffleClient { + def pushOrMergeData( + shuffleId: Int, + mapId: Int, + attemptId: Int, + partitionId: Int, + bytes: Array[Byte], + offset: Int, + length: Int, + numMappers: Int, + numPartitions: Int, + doPush: Boolean, + skipCompress: Boolean): Long = length.toLong +} + +class CelebornShufflePartitionPusherSuite extends AnyFunSuite { + + private val managerKey = "spark.shuffle.manager" + private val managerClass = "org.apache.spark.shuffle.celeborn.SparkShuffleManager" + private val compositeManagerClass = + "org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager" + private val pluginKey = "spark.shuffle.sort.io.plugin.class" + private val pluginClass = "org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO" + private val endpointsKey = "spark.celeborn.master.endpoints" + private val enabledKey = "spark.comet.celeborn.enabled" + + private def enabledConf: SparkConf = + new SparkConf(false).set(endpointsKey, "celeborn-master:9097") + + /** Spark scopes empty() to private[spark] in Scala, but its JVM companion method is public. */ + private def emptyTaskContext(): TaskContext = + TaskContext.getClass.getMethod("empty").invoke(TaskContext).asInstanceOf[TaskContext] + + private def pusher(client: AnyRef): CelebornShufflePartitionPusher = + new CelebornShufflePartitionPusher(client, 19, 3, (4 << 16) | 7, 12, 9) + + test("raw Celeborn push forwards captured task metadata and preserves Comet frame bytes") { + val client = new RecordingCelebornShuffleClient + val bytes = Array[Byte](1, 2, 3, 4) + + assert(pusher(client).pushPartitionData(6, bytes, 3) == 3) + + val push = client.lastPush + assert(push.shuffleId == 19) + assert(push.mapId == 3) + assert(push.attemptId == ((4 << 16) | 7)) + assert(push.partitionId == 6) + assert(push.bytes eq bytes) + assert(push.offset == 0) + assert(push.length == 3) + assert(push.numMappers == 12) + assert(push.numPartitions == 9) + assert(push.doPush) + assert(push.skipCompress) + } + + test("raw Celeborn push requires exactly the payload plus its transport header") { + val bytes = Array[Byte](1, 2, 3) + + Seq(0, bytes.length, bytes.length + 15, bytes.length + 17, -1).foreach { accepted => + val client = new RecordingCelebornShuffleClient + client.acceptedBytes = Some(accepted) + + val error = intercept[IOException] { + pusher(client).pushPartitionData(0, bytes, bytes.length) + } + + assert(error.getMessage.contains(accepted.toString)) + assert(error.getMessage.contains((bytes.length + 16).toString)) + } + } + + test("raw Celeborn push preserves the exact client IOException") { + val client = new RecordingCelebornShuffleClient + val expected = new IOException("worker rejected the shuffle frame") + client.failure = expected + + val actual = intercept[IOException] { + pusher(client).pushPartitionData(0, Array[Byte](7), 1) + } + + assert(actual eq expected) + } + + test("raw Celeborn push preserves unchecked client failures") { + val client = new RecordingCelebornShuffleClient + val expected = new IllegalStateException("client was closed") + client.failure = expected + + val actual = intercept[IllegalStateException] { + pusher(client).pushPartitionData(0, Array[Byte](7), 1) + } + + assert(actual eq expected) + } + + test("adapter rejects a missing client or incompatible Celeborn raw-push API") { + intercept[IllegalArgumentException] { + new CelebornShufflePartitionPusher(null, 0, 0, 0, 1, 1) + } + + val missing = intercept[IllegalArgumentException] { + new CelebornShufflePartitionPusher(new Object, 0, 0, 0, 1, 1) + } + assert(missing.getMessage.contains("raw-push")) + + intercept[IllegalArgumentException] { + new CelebornShufflePartitionPusher(new WrongReturnTypeCelebornShuffleClient, 0, 0, 0, 1, 1) + } + } + + test("adapter rejects invalid shuffle identity, mapper identity, and partition counts") { + val client = new RecordingCelebornShuffleClient + + Seq( + (-1, 0, 0, 1, 1), + (0, -1, 0, 1, 1), + (0, 0, -1, 1, 1), + (0, 0, 0, 0, 1), + (0, 1, 0, 1, 1), + (0, 0, 0, 1, 0)).foreach { case (shuffleId, mapId, attemptId, numMappers, numPartitions) => + intercept[IllegalArgumentException] { + new CelebornShufflePartitionPusher( + client, + shuffleId, + mapId, + attemptId, + numMappers, + numPartitions) + } + } + } + + test("adapter rejects invalid output partitions and incomplete Comet frames") { + val client = new RecordingCelebornShuffleClient + val adapter = pusher(client) + + Seq(-1, 9).foreach { partitionId => + intercept[IOException] { + adapter.pushPartitionData(partitionId, Array[Byte](1), 1) + } + } + + intercept[IOException] { + adapter.pushPartitionData(0, null, 1) + } + + Seq(-1, 0, 2).foreach { length => + intercept[IOException] { + adapter.pushPartitionData(0, Array[Byte](1), length) + } + } + + assert(client.lastPush == null) + } + + test("factory recognizes the existing Celeborn shuffle manager") { + val conf = new SparkConf(false).set(managerKey, managerClass) + + assert(CelebornShufflePusherFactory.isEnabled(conf)) + + conf.set(managerKey, "org.apache.spark.shuffle.sort.SortShuffleManager") + assert(!CelebornShufflePusherFactory.isEnabled(conf)) + } + + test("factory recognizes the composite Comet and Celeborn shuffle manager") { + val conf = new SparkConf(false).set(managerKey, compositeManagerClass) + + assert(CelebornShufflePusherFactory.isEnabled(conf)) + + conf.set(enabledKey, "false") + assert(!CelebornShufflePusherFactory.isEnabled(conf)) + } + + test("factory recognizes the existing Celeborn shuffle data IO plugin") { + val conf = new SparkConf(false).set(pluginKey, pluginClass) + + assert(CelebornShufflePusherFactory.isEnabled(conf)) + + conf.set(pluginKey, "org.apache.spark.shuffle.sort.io.LocalDiskShuffleDataIO") + assert(!CelebornShufflePusherFactory.isEnabled(conf)) + } + + test("factory recognizes nonblank existing Celeborn master endpoints") { + assert(CelebornShufflePusherFactory.isEnabled(enabledConf)) + + val conf = new SparkConf(false).set(endpointsKey, " ") + assert(!CelebornShufflePusherFactory.isEnabled(conf)) + } + + test("application enable flag alone does not invent an unconfigured Celeborn backend") { + val conf = new SparkConf(false).set(enabledKey, "true") + + assert(!CelebornShufflePusherFactory.isEnabled(conf)) + } + + test("explicit application opt-out overrides every existing Celeborn selection") { + Seq( + managerKey -> managerClass, + managerKey -> compositeManagerClass, + pluginKey -> pluginClass, + endpointsKey -> "celeborn-master:9097").foreach { case (key, value) => + val conf = new SparkConf(false).set(key, value).set(enabledKey, "false") + + assert(!CelebornShufflePusherFactory.isEnabled(conf)) + } + } + + test("disabled factory does not inspect a client, task context, or task metadata") { + val conf = new SparkConf(false) + + assert( + CelebornShufflePusherFactory + .create( + conf, + null, + celebornShuffleId = -1, + numMappers = 0, + numPartitions = 0, + taskContext = null) + .isEmpty) + } + + test("enabled factory binds an existing client to captured Spark task metadata") { + val client = new RecordingCelebornShuffleClient + val taskContext = emptyTaskContext() + val adapter = CelebornShufflePusherFactory + .create( + enabledConf, + client, + celebornShuffleId = 27, + numMappers = 8, + numPartitions = 4, + taskContext = taskContext) + .get + + assert(adapter.pushPartitionData(2, Array[Byte](3, 4), 2) == 2) + assert(client.lastPush.shuffleId == 27) + assert(client.lastPush.mapId == taskContext.partitionId()) + assert(client.lastPush.attemptId == 0) + assert(client.lastPush.numMappers == 8) + assert(client.lastPush.numPartitions == 4) + } + + test("factory encodes the stage and task attempt without loading the Celeborn client jar") { + assert(CelebornShufflePusherFactory.encodeAttemptNumber(4, 7) == ((4 << 16) | 7)) + assert(CelebornShufflePusherFactory.encodeAttemptNumber(0, 0) == 0) + assert(CelebornShufflePusherFactory.encodeAttemptNumber(32767, 65535) == Int.MaxValue) + } + + test("factory rejects attempts that cannot be represented as a nonnegative Celeborn ID") { + Seq((-1, 0), (32768, 0), (0, -1), (0, 65536)).foreach { case (stageAttempt, taskAttempt) => + intercept[IllegalArgumentException] { + CelebornShufflePusherFactory.encodeAttemptNumber(stageAttempt, taskAttempt) + } + } + } + + test("a captured task-owned pusher works on a worker without Spark's thread-local context") { + val client = new RecordingCelebornShuffleClient + val adapter = CelebornShufflePusherFactory + .create(enabledConf, client, 11, 3, 2, emptyTaskContext()) + .get + val failure = new AtomicReference[Throwable]() + + val worker = new Thread(() => { + try { + assert(TaskContext.get() == null) + assert(adapter.pushPartitionData(1, Array[Byte](9), 1) == 1) + } catch { + case error: Throwable => failure.set(error) + } + }) + + worker.start() + worker.join(5000) + + assert(!worker.isAlive, "worker thread did not finish") + assert(failure.get() == null) + assert(client.observedTaskContext == null) + assert(client.lastPush.partitionId == 1) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala new file mode 100644 index 00000000000..4393a0070c3 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala @@ -0,0 +1,227 @@ +/* + * 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.spark.sql.comet.execution.shuffle + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext} +import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} + +class CometCelebornShuffleManagerSuite extends AnyFunSuite { + + private class RecordingShuffleManager extends ShuffleManager { + + val returnedHandle: ShuffleHandle = + new BaseShuffleHandle[Any, Any, Any](31, null) + + var registration: Option[(Int, ShuffleDependency[_, _, _])] = None + var writerCall: Option[(ShuffleHandle, Long)] = None + var readerCall: Option[(ShuffleHandle, Int, Int, Int, Int)] = None + var unregisteredShuffleId: Option[Int] = None + var unregisterResult = true + var resolverReads = 0 + var stopped = false + var registrationFailure: RuntimeException = _ + + override def registerShuffle[K, V, C]( + shuffleId: Int, + dependency: ShuffleDependency[K, V, C]): ShuffleHandle = { + registration = Some((shuffleId, dependency)) + if (registrationFailure != null) { + throw registrationFailure + } + returnedHandle + } + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { + writerCall = Some((handle, mapId)) + null + } + + override def getReader[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + readerCall = Some((handle, startMapIndex, endMapIndex, startPartition, endPartition)) + null + } + + override def shuffleBlockResolver: ShuffleBlockResolver = { + resolverReads += 1 + null + } + + override def unregisterShuffle(shuffleId: Int): Boolean = { + unregisteredShuffleId = Some(shuffleId) + unregisterResult + } + + override def stop(): Unit = stopped = true + } + + private def manager( + backend: ShuffleManager, + conf: SparkConf = new SparkConf(false), + isDriver: Boolean = true): CometCelebornShuffleManager = { + new CometCelebornShuffleManager(conf, isDriver, (_, _) => backend) + } + + test("manager forwards the existing Spark configuration and driver identity") { + val conf = new SparkConf(false) + .set("spark.celeborn.master.endpoints", "existing-master:9097") + val backend = new RecordingShuffleManager + var observedConf: SparkConf = null + var observedIsDriver = true + + new CometCelebornShuffleManager( + conf, + false, + (actualConf, actualIsDriver) => { + observedConf = actualConf + observedIsDriver = actualIsDriver + backend + }) + + assert(observedConf eq conf) + assert(!observedIsDriver) + assert(conf.get("spark.celeborn.master.endpoints") == "existing-master:9097") + } + + test("ordinary shuffle registration preserves the existing Celeborn handle and fallback") { + val backend = new RecordingShuffleManager + val composite = manager(backend) + val dependency = null.asInstanceOf[ShuffleDependency[Any, Any, Any]] + + val handle = composite.registerShuffle(31, dependency) + + assert(handle eq backend.returnedHandle) + assert(backend.registration.contains((31, dependency))) + } + + test("ordinary map writers and reduce readers are owned by the existing Celeborn manager") { + val backend = new RecordingShuffleManager + val composite = manager(backend) + val handle = backend.returnedHandle + + assert(composite.getWriter[Any, Any](handle, 93L, null, null) == null) + assert(backend.writerCall.contains((handle, 93L))) + + assert(composite.getReader[Any, Any](handle, 2, 8, 3, 7, null, null) == null) + assert(backend.readerCall.contains((handle, 2, 8, 3, 7))) + } + + test("the inherited all-mapper reader also delegates to the existing Celeborn manager") { + val backend = new RecordingShuffleManager + val composite = manager(backend) + val handle = backend.returnedHandle + + assert(composite.getReader[Any, Any](handle, 4, 9, null, null) == null) + assert(backend.readerCall.contains((handle, 0, Int.MaxValue, 4, 9))) + } + + test("resolver, shuffle cleanup, and shutdown preserve existing Celeborn behavior") { + val backend = new RecordingShuffleManager + backend.unregisterResult = false + val composite = manager(backend) + + assert(composite.shuffleBlockResolver == null) + assert(backend.resolverReads == 1) + + assert(!composite.unregisterShuffle(17)) + assert(backend.unregisteredShuffleId.contains(17)) + + composite.stop() + assert(backend.stopped) + } + + test("backend registration failures propagate without local Comet fallback") { + val backend = new RecordingShuffleManager + val expected = new IllegalStateException("remote shuffle registration failed") + backend.registrationFailure = expected + val composite = manager(backend) + + val actual = intercept[IllegalStateException] { + composite.registerShuffle[Any, Any, Any](31, null) + } + + assert(actual eq expected) + assert(backend.writerCall.isEmpty) + } + + test("native and JVM Comet handles never reach the stock Celeborn row path") { + val backend = new RecordingShuffleManager + val composite = manager(backend) + val unsupportedHandles = Seq[ShuffleHandle]( + new CometNativeShuffleHandle[Any, Any](41, null), + new CometBypassMergeSortShuffleHandle[Any, Any](42, null), + new CometSerializedShuffleHandle[Any, Any](43, null)) + + unsupportedHandles.foreach { handle => + val writerFailure = intercept[UnsupportedOperationException] { + composite.getWriter[Any, Any](handle, 0L, null, null) + } + assert(writerFailure.getMessage.contains("Comet shuffle over Celeborn is not supported")) + + val readerFailure = intercept[UnsupportedOperationException] { + composite.getReader[Any, Any](handle, 0, 1, 0, 1, null, null) + } + assert(readerFailure.getMessage.contains("Comet shuffle over Celeborn is not supported")) + } + + assert(backend.writerCall.isEmpty) + assert(backend.readerCall.isEmpty) + } + + test("a missing delegated backend fails closed without selecting local shuffle") { + val error = intercept[IllegalStateException] { + new CometCelebornShuffleManager(new SparkConf(false), true, (_, _) => null) + } + + assert(error.getMessage.contains("factory returned null")) + } + + test("the public constructor rejects an application without the Celeborn client") { + val celebornManagerAvailable = + try { + getClass.getClassLoader.loadClass("org.apache.spark.shuffle.celeborn.SparkShuffleManager") + true + } catch { + case _: ClassNotFoundException => false + } + + if (celebornManagerAvailable) { + cancel("The optional Celeborn Spark client is present on this test classpath") + } + + val error = intercept[IllegalStateException] { + new CometCelebornShuffleManager(new SparkConf(false), true) + } + + assert(error.getMessage.contains("Celeborn Spark shuffle manager is not available")) + } +} From db51df0a87367645d2711ef6aec5edf84f51ca93 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:19:09 +0000 Subject: [PATCH 06/12] feat: complete Celeborn native shuffle map-side push lifecycle --- native/Cargo.lock | 1 + .../src/shuffle_partition_pusher.rs | 95 +- .../java/RecordingShufflePartitionPusher.java | 21 + .../tests/shuffle_partition_pusher.rs | 44 +- native/shuffle/Cargo.toml | 1 + .../src/writers/rss/rss_partition_writer.rs | 1533 ++++++++++++++++- .../CelebornShufflePartitionPusher.java | 794 ++++++++- .../shuffle/ExecutorShufflePushAdmission.java | 93 + .../comet/shuffle/ShufflePartitionPusher.java | 11 + .../scala/org/apache/comet/CometConf.scala | 26 + .../org/apache/comet/CometExecIterator.scala | 37 +- .../CelebornShufflePusherFactory.scala | 295 +++- .../shuffle/CometCelebornShuffleManager.scala | 645 ++++++- .../shuffle/CometNativeShuffleWriter.scala | 289 +++- .../shuffle/RecordingCelebornSparkUtils.java | 39 + .../CelebornShufflePartitionPusherSuite.scala | 1280 +++++++++++++- ...ometCelebornNativeShuffleWriterSuite.scala | 620 +++++++ .../CometCelebornShuffleManagerSuite.scala | 351 +++- 18 files changed, 6029 insertions(+), 146 deletions(-) create mode 100644 spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java create mode 100644 spark/src/test/java/org/apache/comet/shuffle/RecordingCelebornSparkUtils.java create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala diff --git a/native/Cargo.lock b/native/Cargo.lock index 2aa88ee591c..58e8976d32d 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2041,6 +2041,7 @@ name = "datafusion-comet-shuffle" version = "1.1.0" dependencies = [ "arrow", + "arrow-select", "async-trait", "bytes", "clap", diff --git a/native/jni-bridge/src/shuffle_partition_pusher.rs b/native/jni-bridge/src/shuffle_partition_pusher.rs index d5bf744acfe..66d456c5e9d 100644 --- a/native/jni-bridge/src/shuffle_partition_pusher.rs +++ b/native/jni-bridge/src/shuffle_partition_pusher.rs @@ -42,6 +42,8 @@ struct Callback { // The class must stay alive for as long as the cached method ID. _class: Global>, push_partition_data: JMethodID, + reserve_partition_data: JMethodID, + release_partition_data_reservation: JMethodID, } impl JavaShufflePartitionPusher { @@ -76,18 +78,81 @@ impl JavaShufflePartitionPusher { jni::jni_str!("pushPartitionData"), jni::jni_sig!("(I[BI)I"), )?; + let reserve_partition_data = env.get_method_id( + &class, + jni::jni_str!("reservePartitionData"), + jni::jni_sig!("(I)V"), + )?; + let release_partition_data_reservation = env.get_method_id( + &class, + jni::jni_str!("releasePartitionDataReservation"), + jni::jni_sig!("()V"), + )?; Ok(Self { inner: Arc::new(Callback { vm: env.get_java_vm()?, object: env.new_global_ref(object)?, _class: env.new_global_ref(&class)?, push_partition_data, + reserve_partition_data, + release_partition_data_reservation, }), num_partitions, max_push_bytes, }) } + /// Reserve JVM-owned admission before the caller allocates an encoded native frame. + /// + /// Arrow's uncompressed IPC scratch space can exceed the configured compressed-frame limit. + /// The JVM callback validates this reservation against its executor-wide admission budget; + /// only actual submitted frames are constrained by `max_push_bytes`. + pub fn reserve_partition_data(&self, reservation_bytes: usize) -> CometResult<()> { + if reservation_bytes == 0 { + return Err(CometError::Config( + "RSS reservation size must be positive".into(), + )); + } + let reservation_bytes = i32::try_from(reservation_bytes) + .map_err(|_| CometError::Config("RSS reservation size exceeds jint".into()))?; + + self.inner + .vm + .attach_current_thread(|env| -> jni::errors::Result<()> { + let args = [JValue::Int(reservation_bytes).as_jni()]; + // SAFETY: try_new cached the exact interface method and retains its class. + unsafe { + env.call_method_unchecked( + &self.inner.object, + self.inner.reserve_partition_data, + ReturnType::Primitive(Primitive::Void), + &args, + )?; + } + Ok(()) + }) + .map_err(Self::callback_error) + } + + /// Return a reservation whose encoded frame was never submitted to the Java callback. + pub fn release_partition_data_reservation(&self) -> CometResult<()> { + self.inner + .vm + .attach_current_thread(|env| -> jni::errors::Result<()> { + // SAFETY: try_new cached the exact interface method and retains its class. + unsafe { + env.call_method_unchecked( + &self.inner.object, + self.inner.release_partition_data_reservation, + ReturnType::Primitive(Primitive::Void), + &[], + )?; + } + Ok(()) + }) + .map_err(Self::callback_error) + } + /// Copies a caller-validated complete frame into Java-owned memory. /// /// The configured bound is checked before allocating the Java array. This only bounds this @@ -135,19 +200,7 @@ impl JavaShufflePartitionPusher { .i() } }) - .map_err(|source| match source { - JniError::CaughtJavaException { - exception, - name, - msg, - .. - } => CometError::JavaException { - class: name, - msg, - throwable: exception, - }, - source => CometError::JNI { source }, - })?; + .map_err(Self::callback_error)?; if accepted != length { return Err(CometError::Internal(format!( @@ -156,4 +209,20 @@ impl JavaShufflePartitionPusher { } Ok(()) } + + fn callback_error(source: JniError) -> CometError { + match source { + JniError::CaughtJavaException { + exception, + name, + msg, + .. + } => CometError::JavaException { + class: name, + msg, + throwable: exception, + }, + source => CometError::JNI { source }, + } + } } diff --git a/native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java b/native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java index 85427a571d9..b60c750f8bc 100644 --- a/native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java +++ b/native/jni-bridge/tests/java/RecordingShufflePartitionPusher.java @@ -27,15 +27,36 @@ public final class RecordingShufflePartitionPusher implements ShufflePartitionPu public int partitionId; public int adjustment; public int failureMode; + public int reservationCalls; + public int reservationReleases; + public int reservedBytes; + public boolean reservedBeforePush; public byte[] lastBytes; public final IOException failure = new IOException("recorded push failure"); + @Override + public void reservePartitionData(int maxLength) throws IOException { + reservationCalls++; + if (failureMode == 2) { + throw failure; + } + reservedBytes = maxLength; + } + + @Override + public void releasePartitionDataReservation() { + reservationReleases++; + reservedBytes = 0; + } + @Override public int pushPartitionData(int partitionId, byte[] bytes, int length) throws IOException { calls++; if (failureMode != 0) { throw failure; } + reservedBeforePush = reservedBytes >= length; + reservedBytes = 0; this.partitionId = partitionId; lastBytes = Arrays.copyOf(bytes, length); return length + adjustment; diff --git a/native/jni-bridge/tests/shuffle_partition_pusher.rs b/native/jni-bridge/tests/shuffle_partition_pusher.rs index 9a7b3a38f76..1a1b5a147b7 100644 --- a/native/jni-bridge/tests/shuffle_partition_pusher.rs +++ b/native/jni-bridge/tests/shuffle_partition_pusher.rs @@ -119,10 +119,13 @@ fn rss_jni_callback_contract() { // Global references and the cached method remain valid on another attached native thread. let worker_pusher = pusher.clone(); let worker_frame = frame.clone(); - std::thread::spawn(move || worker_pusher.push_partition_data(1, &worker_frame)) - .join() - .unwrap() - .unwrap(); + std::thread::spawn(move || { + worker_pusher.reserve_partition_data(worker_frame.len())?; + worker_pusher.push_partition_data(1, &worker_frame) + }) + .join() + .unwrap() + .unwrap(); vm.attach_current_thread(|env| -> jni::errors::Result<()> { assert_eq!( @@ -135,6 +138,22 @@ fn rss_jni_callback_contract() { .i()?, 1 ); + assert_eq!( + env.get_field( + &fixture, + jni::jni_str!("reservationCalls"), + jni::jni_sig!("I"), + )? + .i()?, + 1 + ); + assert!(env + .get_field( + &fixture, + jni::jni_str!("reservedBeforePush"), + jni::jni_sig!("Z"), + )? + .z()?); let bytes = env .get_field(&fixture, jni::jni_str!("lastBytes"), jni::jni_sig!("[B"))? .l()?; @@ -146,6 +165,23 @@ fn rss_jni_callback_contract() { .unwrap(); assert!(pusher.push_partition_data(2, &frame).is_err()); + assert!(pusher.reserve_partition_data(0).is_err()); + assert!(pusher + .reserve_partition_data(i32::MAX as usize + 1) + .is_err()); + pusher.reserve_partition_data(frame.len() + 1).unwrap(); + vm.attach_current_thread(|env| -> jni::errors::Result<()> { + assert_eq!( + env.get_field(&fixture, jni::jni_str!("reservedBytes"), jni::jni_sig!("I"))? + .i()?, + frame.len() as i32 + 1 + ); + Ok(()) + }) + .unwrap(); + pusher.release_partition_data_reservation().unwrap(); + pusher.reserve_partition_data(frame.len()).unwrap(); + pusher.release_partition_data_reservation().unwrap(); assert!(pusher.push_partition_data(0, &[]).is_err()); assert!(pusher .push_partition_data(0, &vec![0; frame.len() + 1]) diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml index 58d82d24803..ca02085a5cc 100644 --- a/native/shuffle/Cargo.toml +++ b/native/shuffle/Cargo.toml @@ -30,6 +30,7 @@ publish = false [dependencies] arrow = { workspace = true } +arrow-select = "58.4.0" async-trait = { workspace = true } bytes = { workspace = true } clap = { version = "4", features = ["derive"], optional = true } diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs index 357e2ef819c..406003601dc 100644 --- a/native/shuffle/src/writers/rss/rss_partition_writer.rs +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -18,7 +18,14 @@ use std::io::{self, Cursor, Seek, SeekFrom, Write}; use std::sync::Arc; +use arrow::array::cast::AsArray; +use arrow::array::{ + Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, MapArray, StructArray, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::DataType; use arrow::record_batch::RecordBatch; +use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion::common::{DataFusionError, Result}; use datafusion::physical_plan::metrics::Time; use datafusion_comet_jni_bridge::shuffle_partition_pusher::JavaShufflePartitionPusher; @@ -32,16 +39,43 @@ use crate::ShuffleBlockWriter; /// Acceptance is not a remote map commit. Implementations own asynchronous admission, retry, /// cancellation, and commit; they must not split a frame into independently interleavable pushes. pub trait PartitionPusher: Send + Sync { + /// Reserve executor-wide capacity for encoding scratch and every overlapping frame copy. + fn reserve_partition_data(&self, _reservation_bytes: usize) -> Result<()> { + Ok(()) + } + + /// Return an unsubmitted encoding reservation. + fn release_partition_data_reservation(&self) -> Result<()> { + Ok(()) + } + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()>; } impl PartitionPusher for Arc

{ + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + self.as_ref().reserve_partition_data(reservation_bytes) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + self.as_ref().release_partition_data_reservation() + } + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { self.as_ref().push_partition_data(partition_id, frame) } } impl PartitionPusher for JavaShufflePartitionPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + JavaShufflePartitionPusher::reserve_partition_data(self, reservation_bytes) + .map_err(Into::into) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + JavaShufflePartitionPusher::release_partition_data_reservation(self).map_err(Into::into) + } + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { JavaShufflePartitionPusher::push_partition_data(self, partition_id, frame) .map_err(Into::into) @@ -51,9 +85,12 @@ impl PartitionPusher for JavaShufflePartitionPusher { /// Encodes already-partitioned batches using the existing Comet format and sends one frame per /// callback. The native planner selects it only for a registered task-owned pusher. /// -/// There are no retained per-reducer buffers. The byte limit caps encoded output, not Arrow's -/// encoding scratch space or the backend's asynchronous memory. Production admission, row -/// splitting, integrity, and map-commit handling must be supplied before enabling remote plans. +/// There are no retained per-reducer buffers. Oversized batches are split at row boundaries into +/// independently decodable frames before Arrow can allocate an oversized IPC body. A single row +/// with an oversized uncompressed body may still fit after compression; its complete scratch +/// budget is reserved before encoding, and an oversized encoded frame remains an error. Every +/// reservation also covers the simultaneously live native frame, JNI array, and backend request +/// payload. The backend owns asynchronous transport admission. pub struct RssPartitionWriter { pusher: P, block_writer: ShuffleBlockWriter, @@ -107,15 +144,7 @@ impl RssPartitionWriter

{ ) -> Result<()> { let result = (|| { self.check_partition(partition_id)?; - let mut output = BoundedBuffer::new(self.max_frame_bytes); - self.block_writer - .write_batch(batch, &mut output, encode_time)?; - let frame = output.inner.get_ref(); - if !frame.is_empty() { - let _timer = write_time.timer(); - self.pusher.push_partition_data(partition_id, frame)?; - } - Ok(()) + self.write_batch_within_limit(partition_id, batch, encode_time, write_time) })(); if result.is_err() { self.failed = true; @@ -123,6 +152,571 @@ impl RssPartitionWriter

{ result } + fn write_batch_within_limit( + &mut self, + partition_id: usize, + batch: &RecordBatch, + encode_time: &Time, + write_time: &Time, + ) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + + // Arrow materializes full dictionary remaps, copied values, and the IPC body before + // writing anything to our bounded destination. Estimate those allocations without + // compacting first; ordinary oversized batches can still be split before admission. + let original_ipc_data_size = Self::estimated_pre_compaction_ipc_data_size(batch)?; + let compaction_scratch = Self::estimated_compaction_scratch(batch)?; + if compaction_scratch == 0 + && original_ipc_data_size > self.max_frame_bytes + && batch.num_rows() > 1 + { + return self.write_split_batch(partition_id, batch, encode_time, write_time); + } + + // The encoded native frame remains live while JNI copies it into a Java array and the + // backend copies that array into its asynchronous request. Acquire all three copies + // together before encoding: growing admission after any allocation can deadlock other + // concurrently admitted encoders. The JVM accounts for its additional request header. + let overlapping_frame_bytes = self.max_frame_bytes.checked_mul(3).ok_or_else(|| { + DataFusionError::Configuration("RSS frame copy reservation exceeds usize".into()) + })?; + let reservation_bytes = + overlapping_frame_bytes.max(original_ipc_data_size.saturating_add(compaction_scratch)); + self.pusher.reserve_partition_data(reservation_bytes)?; + + let compacted = match Self::compact_dictionary_columns(batch) { + Ok(compacted) => compacted, + Err(error) => { + self.pusher.release_partition_data_reservation()?; + return Err(error); + } + }; + let compacted_batch = compacted.as_ref().unwrap_or(batch); + let ipc_data_size = match Self::estimated_ipc_data_size(compacted_batch) { + Ok(ipc_data_size) => ipc_data_size, + Err(error) => { + drop(compacted); + self.pusher.release_partition_data_reservation()?; + return Err(error); + } + }; + if ipc_data_size > self.max_frame_bytes && compacted_batch.num_rows() > 1 { + // Do not retain the compacted parent or its thread-owned reservation while children + // compete for executor admission. Re-compact each original child after it is admitted. + drop(compacted); + self.pusher.release_partition_data_reservation()?; + return self.write_split_batch(partition_id, batch, encode_time, write_time); + } + + let mut output = BoundedBuffer::new(self.max_frame_bytes); + if let Err(error) = self.block_writer.write_batch( + compacted_batch, + &mut output, + &mut self.compression_context, + encode_time, + ) { + let exceeded = output.exceeded; + drop(output); + drop(compacted); + self.pusher.release_partition_data_reservation()?; + if !exceeded { + return Err(error); + } + if batch.num_rows() <= 1 { + return Err(self.oversized_row()); + } + return self.write_split_batch(partition_id, batch, encode_time, write_time); + } + + // The JVM shrinks admission to three actual frame copies when claiming this reservation, + // then to the retained request only after submission. Release dictionary copies and + // normalization scratch before the smaller synchronous-copy claim becomes visible. + drop(compacted); + let frame = output.inner.get_ref(); + if !frame.is_empty() { + let _timer = write_time.timer(); + if let Err(error) = self.pusher.push_partition_data(partition_id, frame) { + drop(output); + self.pusher.release_partition_data_reservation()?; + return Err(error); + } + } else { + drop(output); + self.pusher.release_partition_data_reservation()?; + } + Ok(()) + } + + fn write_split_batch( + &mut self, + partition_id: usize, + batch: &RecordBatch, + encode_time: &Time, + write_time: &Time, + ) -> Result<()> { + let midpoint = batch.num_rows() / 2; + self.write_batch_within_limit( + partition_id, + &batch.slice(0, midpoint), + encode_time, + write_time, + )?; + self.write_batch_within_limit( + partition_id, + &batch.slice(midpoint, batch.num_rows() - midpoint), + encode_time, + write_time, + ) + } + + fn oversized_row(&self) -> DataFusionError { + DataFusionError::Execution(format!( + "RSS frame exceeds its byte limit: a single RSS row exceeds the {}-byte limit", + self.max_frame_bytes + )) + } + + fn estimated_ipc_data_size(batch: &RecordBatch) -> Result { + batch.columns().iter().try_fold(0usize, |total, column| { + let data = column.to_data(); + let logical = data.get_slice_memory_size()?; + let additional = Self::additional_ipc_data_size(&data); + // IPC aligns every buffer to 64 bytes and may synthesize a validity buffer even + // when an array has no nulls. Include that padding before charging one frame. + let padding = Self::ipc_buffer_count(&data).saturating_mul(64); + Ok(total + .saturating_add(logical) + .saturating_add(additional) + .saturating_add(padding)) + }) + } + + /// Estimate logical rows without charging unrelated list/map backing entries. + /// + /// `ArrayData::get_slice_memory_size` recurses through an entire list/map child even when the + /// parent offsets select only one entry. Their zero-copy child slices are sufficient here; + /// dictionary values remain fully charged because dense garbage collection traverses them. + fn estimated_pre_compaction_ipc_data_size(batch: &RecordBatch) -> Result { + batch.columns().iter().try_fold(0usize, |total, column| { + Ok(total.saturating_add(Self::estimated_live_array_ipc_data_size(column.as_ref())?)) + }) + } + + fn estimated_live_array_ipc_data_size(array: &dyn Array) -> Result { + let data = array.to_data(); + let child_logical = data.child_data().iter().try_fold(0usize, |total, child| { + Ok::<_, DataFusionError>(total.saturating_add(child.get_slice_memory_size()?)) + })?; + let logical = data.get_slice_memory_size()?.saturating_sub(child_logical); + let child_additional = data.child_data().iter().fold(0usize, |total, child| { + total.saturating_add(Self::additional_ipc_data_size(child)) + }); + let additional = Self::additional_ipc_data_size(&data).saturating_sub(child_additional); + let padding = data.buffers().len().saturating_add(1).saturating_mul(64); + let own = logical.saturating_add(additional).saturating_add(padding); + + let children = + if let Some(dictionary) = array.as_any_dictionary_opt() { + Self::estimated_live_array_ipc_data_size(dictionary.values().as_ref())? + } else { + match array.data_type() { + DataType::List(_) => { + let list = array.as_any().downcast_ref::().unwrap(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let live_values = list.values().slice(start, end - start); + Self::estimated_live_array_ipc_data_size(live_values.as_ref())? + } + DataType::LargeList(_) => { + let list = array.as_any().downcast_ref::().unwrap(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let live_values = list.values().slice(start, end - start); + Self::estimated_live_array_ipc_data_size(live_values.as_ref())? + } + DataType::FixedSizeList(_, _) => { + let list = array.as_any().downcast_ref::().unwrap(); + Self::estimated_live_array_ipc_data_size(list.values().as_ref())? + } + DataType::Struct(_) => { + let structure = array.as_any().downcast_ref::().unwrap(); + structure + .columns() + .iter() + .try_fold(0usize, |total, column| { + Ok::<_, DataFusionError>(total.saturating_add( + Self::estimated_live_array_ipc_data_size(column.as_ref())?, + )) + })? + } + DataType::Map(_, _) => { + let map = array.as_any().downcast_ref::().unwrap(); + let offsets = map.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let live_entries = map.entries().slice(start, end - start); + Self::estimated_live_array_ipc_data_size(&live_entries)? + } + _ => data.child_data().iter().try_fold(0usize, |total, child| { + let logical = child.get_slice_memory_size()?; + let additional = Self::additional_ipc_data_size(child); + let padding = Self::ipc_buffer_count(child).saturating_mul(64); + Ok::<_, DataFusionError>(total.saturating_add( + logical.saturating_add(additional).saturating_add(padding), + )) + })?, + } + }; + Ok(own.saturating_add(children)) + } + + fn additional_ipc_data_size(data: &arrow::array::ArrayData) -> usize { + // Arrow IPC writes a validity bitmap even when the original array has no null buffer. + let validity = if data.nulls().is_none() { + data.len().div_ceil(8) + } else { + 0 + }; + // ArrayData's slice estimate charges one offset per row, while IPC writes the additional + // terminating offset. View arrays instead serialize every shared backing data buffer. + let own = if matches!(data.data_type(), DataType::BinaryView | DataType::Utf8View) { + data.buffers() + .iter() + .skip(1) + .fold(0usize, |total, buffer| total.saturating_add(buffer.len())) + } else { + match data.data_type() { + DataType::Binary | DataType::Utf8 => { + let temporary = if data.buffer::(0)[0] != 0 { + data.len().saturating_add(1).saturating_mul(4) + } else { + 0 + }; + 4usize.saturating_add(temporary) + } + DataType::LargeBinary | DataType::LargeUtf8 => { + let temporary = if data.buffer::(0)[0] != 0 { + data.len().saturating_add(1).saturating_mul(8) + } else { + 0 + }; + 8usize.saturating_add(temporary) + } + DataType::List(_) | DataType::Map(_, _) => 4, + DataType::LargeList(_) => 8, + _ => 0, + } + }; + data.child_data() + .iter() + .fold(own.saturating_add(validity), |total, child| { + total.saturating_add(Self::additional_ipc_data_size(child)) + }) + } + + fn ipc_buffer_count(data: &arrow::array::ArrayData) -> usize { + data.child_data() + .iter() + .fold(data.buffers().len().saturating_add(1), |total, child| { + total.saturating_add(Self::ipc_buffer_count(child)) + }) + } + + /// Conservatively charge dense dictionary garbage collection and nested offset rebasing. + /// + /// This walks the original arrays without compacting, building occupancy masks, or allocating + /// per-row scratch. In particular, dense remaps scale with the *full* dictionary cardinality, + /// even when only one value is referenced by the current batch slice. + fn estimated_compaction_scratch(batch: &RecordBatch) -> Result { + batch.columns().iter().try_fold(0usize, |total, column| { + Ok(total.saturating_add(Self::estimated_array_compaction_scratch(column.as_ref())?)) + }) + } + + fn estimated_array_compaction_scratch(array: &dyn Array) -> Result { + if let Some(dictionary) = array.as_any_dictionary_opt() { + let key_width = match dictionary.keys().data_type() { + DataType::Int8 | DataType::UInt8 => 1, + DataType::Int16 | DataType::UInt16 => 2, + DataType::Int32 | DataType::UInt32 => 4, + DataType::Int64 | DataType::UInt64 => 8, + _ => unreachable!("Arrow dictionary keys must have an integer type"), + }; + let value_count = dictionary.values().len(); + let key_count = dictionary.keys().len(); + let occupancy = value_count.div_ceil(8).saturating_add(64); + let dense_remap = value_count.saturating_mul(key_width); + let copied_keys = key_count + .saturating_mul(key_width) + .saturating_add(key_count.div_ceil(8)) + .saturating_add(64); + let filter_indices = key_count + .min(value_count) + .saturating_mul(std::mem::size_of::<(usize, usize)>()); + let values = dictionary.values().to_data(); + let copied_values = values + .get_slice_memory_size()? + .saturating_add(Self::additional_ipc_data_size(&values)) + .saturating_add(Self::ipc_buffer_count(&values).saturating_mul(64)); + let nested = Self::estimated_array_compaction_scratch(dictionary.values().as_ref())?; + return Ok(occupancy + .saturating_add(dense_remap) + .saturating_add(copied_keys) + .saturating_add(filter_indices) + .saturating_add(copied_values) + .saturating_add(nested)); + } + + match array.data_type() { + DataType::List(_) => { + let list = array.as_any().downcast_ref::().unwrap(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let normalized = start != 0 || end != list.values().len(); + let rebased = if normalized { + list.len().saturating_add(1).saturating_mul(4) + } else { + 0 + }; + let live_values = list.values().slice(start, end - start); + Ok( + rebased.saturating_add(Self::estimated_array_compaction_scratch( + live_values.as_ref(), + )?), + ) + } + DataType::LargeList(_) => { + let list = array.as_any().downcast_ref::().unwrap(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let normalized = start != 0 || end != list.values().len(); + let rebased = if normalized { + list.len().saturating_add(1).saturating_mul(8) + } else { + 0 + }; + let live_values = list.values().slice(start, end - start); + Ok( + rebased.saturating_add(Self::estimated_array_compaction_scratch( + live_values.as_ref(), + )?), + ) + } + DataType::FixedSizeList(_, _) => { + let list = array.as_any().downcast_ref::().unwrap(); + Self::estimated_array_compaction_scratch(list.values().as_ref()) + } + DataType::Struct(_) => { + let structure = array.as_any().downcast_ref::().unwrap(); + structure + .columns() + .iter() + .try_fold(0usize, |total, column| { + Ok( + total.saturating_add(Self::estimated_array_compaction_scratch( + column.as_ref(), + )?), + ) + }) + } + DataType::Map(_, _) => { + let map = array.as_any().downcast_ref::().unwrap(); + let offsets = map.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let normalized = start != 0 || end != map.entries().len(); + let rebased = if normalized { + map.len().saturating_add(1).saturating_mul(4) + } else { + 0 + }; + let live_entries = map.entries().slice(start, end - start); + Ok( + rebased + .saturating_add(Self::estimated_array_compaction_scratch(&live_entries)?), + ) + } + _ => Ok(0), + } + } + + fn compact_dictionary_columns(batch: &RecordBatch) -> Result> { + let mut changed = false; + let columns = batch + .columns() + .iter() + .map(|column| -> Result { + let (compacted, column_changed) = Self::compact_array(column)?; + changed |= column_changed; + Ok(compacted) + }) + .collect::>>()?; + + if changed { + Ok(Some(RecordBatch::try_new(batch.schema(), columns)?)) + } else { + Ok(None) + } + } + + fn compact_array(array: &ArrayRef) -> Result<(ArrayRef, bool)> { + if let Some(dictionary) = array.as_any_dictionary_opt() { + let compacted = garbage_collect_any_dictionary(dictionary)?; + let compacted_dictionary = compacted.as_any_dictionary(); + let dictionary_changed = + compacted_dictionary.values().len() != dictionary.values().len(); + let (values, values_changed) = Self::compact_array(compacted_dictionary.values())?; + if values_changed { + return Ok((compacted_dictionary.with_values(values), true)); + } + return Ok(if dictionary_changed { + (compacted, true) + } else { + (Arc::clone(array), false) + }); + } + + match array.data_type() { + DataType::List(field) => { + let list = array.as_any().downcast_ref::().unwrap(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let live_values = list.values().slice(start, end - start); + let (values, child_changed) = Self::compact_array(&live_values)?; + let normalized = start != 0 || end != list.values().len(); + if !child_changed && !normalized { + return Ok((Arc::clone(array), false)); + } + let rebased = if start == 0 { + list.offsets().clone() + } else { + OffsetBuffer::new( + offsets + .iter() + .map(|offset| offset - offsets[0]) + .collect::>() + .into(), + ) + }; + let rebuilt = + ListArray::try_new(Arc::clone(field), rebased, values, list.nulls().cloned())?; + Ok((Arc::new(rebuilt), true)) + } + DataType::LargeList(field) => { + let list = array.as_any().downcast_ref::().unwrap(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let live_values = list.values().slice(start, end - start); + let (values, child_changed) = Self::compact_array(&live_values)?; + let normalized = start != 0 || end != list.values().len(); + if !child_changed && !normalized { + return Ok((Arc::clone(array), false)); + } + let rebased = if start == 0 { + list.offsets().clone() + } else { + OffsetBuffer::new( + offsets + .iter() + .map(|offset| offset - offsets[0]) + .collect::>() + .into(), + ) + }; + let rebuilt = LargeListArray::try_new( + Arc::clone(field), + rebased, + values, + list.nulls().cloned(), + )?; + Ok((Arc::new(rebuilt), true)) + } + DataType::FixedSizeList(field, size) => { + let list = array.as_any().downcast_ref::().unwrap(); + let (values, changed) = Self::compact_array(list.values())?; + if !changed { + return Ok((Arc::clone(array), false)); + } + let rebuilt = FixedSizeListArray::try_new_with_length( + Arc::clone(field), + *size, + values, + list.nulls().cloned(), + list.len(), + )?; + Ok((Arc::new(rebuilt), true)) + } + DataType::Struct(fields) => { + let structure = array.as_any().downcast_ref::().unwrap(); + let mut changed = false; + let columns = structure + .columns() + .iter() + .map(|column| { + let (compacted, child_changed) = Self::compact_array(column)?; + changed |= child_changed; + Ok(compacted) + }) + .collect::>>()?; + if !changed { + return Ok((Arc::clone(array), false)); + } + let rebuilt = StructArray::try_new_with_length( + fields.clone(), + columns, + structure.nulls().cloned(), + structure.len(), + )?; + Ok((Arc::new(rebuilt), true)) + } + DataType::Map(field, ordered) => { + let map = array.as_any().downcast_ref::().unwrap(); + let offsets = map.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[offsets.len() - 1] as usize; + let live_entries: ArrayRef = Arc::new(map.entries().slice(start, end - start)); + let (entries, child_changed) = Self::compact_array(&live_entries)?; + let normalized = start != 0 || end != map.entries().len(); + if !child_changed && !normalized { + return Ok((Arc::clone(array), false)); + } + let rebased = if start == 0 { + map.offsets().clone() + } else { + OffsetBuffer::new( + offsets + .iter() + .map(|offset| offset - offsets[0]) + .collect::>() + .into(), + ) + }; + let entries = entries + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let rebuilt = MapArray::try_new( + Arc::clone(field), + rebased, + entries, + map.nulls().cloned(), + *ordered, + )?; + Ok((Arc::new(rebuilt), true)) + } + _ => Ok((Arc::clone(array), false)), + } + } + fn check_partition(&self, partition_id: usize) -> Result<()> { if self.failed || self.finished { return Err(DataFusionError::Execution("RSS writer is closed".into())); @@ -214,6 +808,7 @@ impl PartitionWriter for RssPartitionWriter

{ struct BoundedBuffer { inner: Cursor>, limit: usize, + exceeded: bool, } impl BoundedBuffer { @@ -221,6 +816,7 @@ impl BoundedBuffer { Self { inner: Cursor::new(Vec::new()), limit, + exceeded: false, } } @@ -236,6 +832,7 @@ impl Write for BoundedBuffer { fn write(&mut self, bytes: &[u8]) -> io::Result { let end = self.inner.position().checked_add(bytes.len() as u64); if end.is_none_or(|end| end > self.limit as u64) { + self.exceeded = true; return Err(Self::limit_error()); } self.inner.write(bytes) @@ -252,6 +849,7 @@ impl Seek for BoundedBuffer { let next = self.inner.seek(position)?; if next > self.limit as u64 { self.inner.set_position(previous); + self.exceeded = true; return Err(Self::limit_error()); } Ok(next) @@ -263,8 +861,12 @@ mod tests { use super::*; use std::sync::{Arc, Mutex}; - use arrow::array::Int64Array; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::array::{ + BinaryArray, BooleanArray, DictionaryArray, Int32Array, Int64Array, StringArray, + StringViewArray, + }; + use arrow::compute::cast; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use crate::{read_ipc_compressed, CompressionCodec}; @@ -281,6 +883,41 @@ mod tests { } } + #[derive(Default)] + struct RecordedReservations { + active: Option, + reservations: Vec, + releases: usize, + frames: Vec>, + } + + #[derive(Clone, Default)] + struct ReservationRecordingPusher(Arc>); + + impl PartitionPusher for ReservationRecordingPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + let mut recorded = self.0.lock().unwrap(); + assert!(recorded.active.is_none()); + recorded.active = Some(reservation_bytes); + recorded.reservations.push(reservation_bytes); + Ok(()) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + let mut recorded = self.0.lock().unwrap(); + assert!(recorded.active.take().is_some()); + recorded.releases += 1; + Ok(()) + } + + fn push_partition_data(&self, _partition_id: usize, frame: &[u8]) -> Result<()> { + let mut recorded = self.0.lock().unwrap(); + assert!(recorded.active.take().unwrap() >= 3 * frame.len()); + recorded.frames.push(frame.to_vec()); + Ok(()) + } + } + fn batch(values: Vec) -> RecordBatch { RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])), @@ -289,6 +926,23 @@ mod tests { .unwrap() } + fn sparse_int32_dictionary_batch(keys: Vec, value_count: i32) -> RecordBatch { + let dictionary = DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::new(Int32Array::from_iter_values(0..value_count)), + ) + .unwrap(); + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)), + false, + )])), + vec![Arc::new(dictionary)], + ) + .unwrap() + } + fn metrics() -> ShufflePartitionerMetrics { ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0) } @@ -337,6 +991,85 @@ mod tests { } } + #[test] + fn rss_partition_writer_reserves_overlapping_frame_copies_before_encoding() { + let limit = 1_024; + let input = batch(vec![1, 2]); + let scratch = + RssPartitionWriter::::estimated_ipc_data_size(&input) + .unwrap(); + assert!(scratch < 3 * limit); + + let pusher = ReservationRecordingPusher::default(); + let recorded = Arc::clone(&pusher.0); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, limit).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input.clone())].into_iter(), &metrics) + .unwrap(); + + let recorded = recorded.lock().unwrap(); + assert!(recorded.active.is_none()); + assert_eq!(recorded.reservations, vec![3 * limit]); + assert_eq!(recorded.releases, 0); + assert_eq!(recorded.frames.len(), 1); + assert_eq!( + read_ipc_compressed(&recorded.frames[0][16..]).unwrap(), + input + ); + } + + #[test] + fn rss_partition_writer_rejects_unadmitted_frame_copies_before_encoding() { + struct RejectingCopiesPusher(Arc>>); + + impl PartitionPusher for RejectingCopiesPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + self.0.lock().unwrap().push(reservation_bytes); + Err(DataFusionError::External(Box::new(io::Error::other( + "overlapping frame copies were not admitted", + )))) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + panic!("a denied frame-copy reservation must not be released") + } + + fn push_partition_data(&self, _partition_id: usize, _frame: &[u8]) -> Result<()> { + panic!("a denied frame-copy reservation must prevent encoding and push") + } + } + + let limit = 1_024; + let input = batch(vec![1, 2]); + let reservations = Arc::new(Mutex::new(Vec::new())); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new( + RejectingCopiesPusher(Arc::clone(&reservations)), + encoder, + 1, + limit, + ) + .unwrap(); + let metrics = metrics(); + + let error = writer + .write(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap_err(); + let DataFusionError::External(source) = error else { + panic!("frame-copy admission failure was wrapped or stringified"); + }; + assert_eq!( + source.downcast_ref::().unwrap().to_string(), + "overlapping frame copies were not admitted" + ); + assert_eq!(*reservations.lock().unwrap(), vec![3 * limit]); + } + #[test] fn rss_partition_writer_checks_lifecycle_and_empty_batches() { let input = batch(vec![]); @@ -390,15 +1123,781 @@ mod tests { let error = writer .write(0, &mut [Ok(input)].into_iter(), &metrics) .unwrap_err(); - assert!(error - .to_string() - .contains("RSS frame exceeds its byte limit")); + assert!(error.to_string().contains("single RSS row exceeds")); assert!(captured.lock().unwrap().is_empty()); assert!(writer .finish_partition(0, &mut std::iter::empty(), &metrics) .is_err()); } + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call ZSTD_createCCtx. + fn rss_partition_writer_splits_oversized_batches_at_row_boundaries() { + let expected: Vec = (0i64..1024) + .map(|value| { + value + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407) + }) + .collect(); + let input = batch(expected.clone()); + for codec in [ + CompressionCodec::None, + CompressionCodec::Lz4Frame, + CompressionCodec::Snappy, + CompressionCodec::Zstd(1), + ] { + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = ShuffleBlockWriter::try_new(input.schema().as_ref(), codec).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, 1024).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input.clone())].into_iter(), &metrics) + .unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = captured.lock().unwrap(); + assert!(frames.len() > 1); + let mut actual = Vec::new(); + for (partition, frame) in frames.iter() { + assert_eq!(*partition, 0); + assert!(frame.len() <= 1024); + let decoded = read_ipc_compressed(&frame[16..]).unwrap(); + let values = decoded + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + actual.extend(values.iter().map(|value| value.unwrap())); + } + assert_eq!(actual, expected); + } + } + + #[test] + fn rss_partition_writer_splits_wide_batches_before_allocating_arrow_ipc() { + let primitive = batch((0..4_096).collect()); + let binary_schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Binary, + false, + )])); + let binary = RecordBatch::try_new( + Arc::clone(&binary_schema), + vec![Arc::new(BinaryArray::from_iter_values( + (0..128).map(|index| vec![index as u8; 512]), + ))], + ) + .unwrap(); + + for input in [primitive, binary] { + let limit = 2_048; + assert!( + RssPartitionWriter::::estimated_ipc_data_size(&input).unwrap() + > limit + ); + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None) + .unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, limit).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input.clone())].into_iter(), &metrics) + .unwrap(); + + let frames = captured.lock().unwrap(); + assert!(frames.len() > 1); + let decoded = frames + .iter() + .map(|(_, frame)| { + assert!(frame.len() <= limit); + read_ipc_compressed(&frame[16..]).unwrap() + }) + .collect::>(); + assert_eq!( + arrow_select::concat::concat_batches(&input.schema(), &decoded).unwrap(), + input + ); + } + } + + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call ZSTD_createCCtx. + fn rss_partition_writer_reserves_oversized_single_row_before_compressing_it() { + let limit = 2_048; + let payload = vec![b'a'; 32 * 1_024]; + let input = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Binary, + false, + )])), + vec![Arc::new(BinaryArray::from_iter_values( + [payload.as_slice()], + ))], + ) + .unwrap(); + let estimated_ipc_data_size = + RssPartitionWriter::::estimated_ipc_data_size(&input) + .unwrap(); + assert!(estimated_ipc_data_size > 3 * limit); + + for codec in [ + CompressionCodec::Lz4Frame, + CompressionCodec::Snappy, + CompressionCodec::Zstd(1), + ] { + let pusher = ReservationRecordingPusher::default(); + let recorded = Arc::clone(&pusher.0); + let encoder = ShuffleBlockWriter::try_new(input.schema().as_ref(), codec).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, limit).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input.clone())].into_iter(), &metrics) + .unwrap(); + writer.finish_all(&metrics).unwrap(); + + let recorded = recorded.lock().unwrap(); + assert!(recorded.active.is_none()); + assert_eq!(recorded.reservations, vec![estimated_ipc_data_size]); + assert_eq!(recorded.releases, 0); + assert_eq!(recorded.frames.len(), 1); + assert!(recorded.frames[0].len() <= limit); + assert_eq!( + read_ipc_compressed(&recorded.frames[0][16..]).unwrap(), + input + ); + } + } + + #[test] + fn rss_partition_writer_rejects_single_row_before_encoding_when_scratch_is_not_admitted() { + struct RejectingPusher(Arc>>); + + impl PartitionPusher for RejectingPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + self.0.lock().unwrap().push(reservation_bytes); + Err(DataFusionError::External(Box::new(io::Error::other( + "scratch admission denied", + )))) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + panic!("a denied scratch reservation must not be released") + } + + fn push_partition_data(&self, _partition_id: usize, _frame: &[u8]) -> Result<()> { + panic!("a denied scratch reservation must prevent encoding and push") + } + } + + let limit = 1_024; + let payload = vec![b'a'; 32 * 1_024]; + let input = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Binary, + false, + )])), + vec![Arc::new(BinaryArray::from_iter_values( + [payload.as_slice()], + ))], + ) + .unwrap(); + let estimated_ipc_data_size = + RssPartitionWriter::::estimated_ipc_data_size(&input).unwrap(); + assert!(estimated_ipc_data_size > limit); + + let reservations = Arc::new(Mutex::new(Vec::new())); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::Lz4Frame) + .unwrap(); + let mut writer = RssPartitionWriter::try_new( + RejectingPusher(Arc::clone(&reservations)), + encoder, + 1, + limit, + ) + .unwrap(); + let metrics = metrics(); + + let error = writer + .write(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap_err(); + let DataFusionError::External(source) = error else { + panic!("scratch admission failure was wrapped or stringified"); + }; + assert_eq!( + source.downcast_ref::().unwrap().to_string(), + "scratch admission denied" + ); + assert_eq!(*reservations.lock().unwrap(), vec![estimated_ipc_data_size]); + } + + #[test] + fn rss_partition_writer_releases_oversized_single_row_when_compression_cannot_fit() { + let limit = 1_024; + let payload = (0..16 * 1_024) + .scan(0x243f_6a88u32, |state, _| { + *state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + Some((*state >> 24) as u8) + }) + .collect::>(); + let input = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Binary, + false, + )])), + vec![Arc::new(BinaryArray::from_iter_values( + [payload.as_slice()], + ))], + ) + .unwrap(); + let estimated_ipc_data_size = + RssPartitionWriter::::estimated_ipc_data_size(&input) + .unwrap(); + assert!(estimated_ipc_data_size > limit); + + let pusher = ReservationRecordingPusher::default(); + let recorded = Arc::clone(&pusher.0); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::Lz4Frame) + .unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, limit).unwrap(); + let metrics = metrics(); + + let error = writer + .write(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap_err(); + assert!(error.to_string().contains("single RSS row exceeds")); + + let recorded = recorded.lock().unwrap(); + assert!(recorded.active.is_none()); + assert_eq!(recorded.reservations, vec![estimated_ipc_data_size]); + assert_eq!(recorded.releases, 1); + assert!(recorded.frames.is_empty()); + } + + #[test] + fn rss_partition_writer_reserves_dense_dictionary_remap_before_compaction() { + struct RejectingDictionaryPusher { + requests: Arc>>, + max_frame_bytes: usize, + } + + impl PartitionPusher for RejectingDictionaryPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + self.requests.lock().unwrap().push(reservation_bytes); + if reservation_bytes > self.max_frame_bytes { + return Err(DataFusionError::External(Box::new(io::Error::other( + "dictionary compaction admission denied", + )))); + } + Ok(()) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + panic!("a denied dictionary reservation must not be released") + } + + fn push_partition_data(&self, _partition_id: usize, _frame: &[u8]) -> Result<()> { + panic!("dictionary compaction must not run before admission") + } + } + + let value_count = 65_536; + let limit = 1_024; + let input = sparse_int32_dictionary_batch(vec![value_count / 2], value_count); + let original_ipc_size = + RssPartitionWriter::::estimated_ipc_data_size(&input) + .unwrap(); + let compaction_scratch = + RssPartitionWriter::::estimated_compaction_scratch(&input) + .unwrap(); + assert!(compaction_scratch >= value_count as usize * std::mem::size_of::()); + assert!(original_ipc_size > limit); + + let requests = Arc::new(Mutex::new(Vec::new())); + let pusher = RejectingDictionaryPusher { + requests: Arc::clone(&requests), + max_frame_bytes: limit, + }; + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, limit).unwrap(); + let metrics = metrics(); + + let error = writer + .write(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap_err(); + let DataFusionError::External(source) = error else { + panic!("dictionary admission failure was wrapped or stringified"); + }; + assert_eq!( + source.downcast_ref::().unwrap().to_string(), + "dictionary compaction admission denied" + ); + assert_eq!( + *requests.lock().unwrap(), + vec![original_ipc_size + compaction_scratch] + ); + } + + #[test] + fn rss_partition_writer_releases_dictionary_compaction_before_splitting() { + let value_count = 8_192; + let limit = 1_024; + let input = sparse_int32_dictionary_batch((0..128).collect(), value_count); + let pusher = ReservationRecordingPusher::default(); + let recorded = Arc::clone(&pusher.0); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, limit).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input.clone())].into_iter(), &metrics) + .unwrap(); + + let recorded = recorded.lock().unwrap(); + assert!(recorded.active.is_none()); + assert!(recorded.releases > 0); + assert!(recorded.frames.len() > 1); + assert_eq!( + recorded.reservations.len(), + recorded.releases + recorded.frames.len() + ); + assert!(recorded + .reservations + .iter() + .all(|bytes| *bytes >= value_count as usize * std::mem::size_of::())); + + let decoded = recorded + .frames + .iter() + .map(|frame| { + assert!(frame.len() <= limit); + read_ipc_compressed(&frame[16..]).unwrap() + }) + .collect::>(); + assert_eq!( + arrow_select::concat::concat_batches(&input.schema(), &decoded).unwrap(), + input + ); + } + + #[test] + fn rss_partition_writer_admits_only_live_entries_of_sliced_lists_and_maps() { + struct BoundedAdmissionPusher { + recorder: ReservationRecordingPusher, + capacity: usize, + } + + impl PartitionPusher for BoundedAdmissionPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + if reservation_bytes > self.capacity { + return Err(DataFusionError::External(Box::new(io::Error::other( + "unrelated nested backing exceeded admission", + )))); + } + self.recorder.reserve_partition_data(reservation_bytes) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + self.recorder.release_partition_data_reservation() + } + + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { + self.recorder.push_partition_data(partition_id, frame) + } + } + + let value_count = 16_384; + let values: ArrayRef = Arc::new(Int32Array::from_iter_values(0..value_count)); + let offsets = OffsetBuffer::new((0..=value_count).collect::>().into()); + let item = Arc::new(Field::new("item", DataType::Int32, false)); + let list = ListArray::try_new( + Arc::clone(&item), + offsets.clone(), + Arc::clone(&values), + None, + ) + .unwrap() + .slice(value_count as usize / 2, 1); + let list_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + DataType::List(item), + false, + )])), + vec![Arc::new(list)], + ) + .unwrap(); + + let fields = vec![ + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::new(Field::new("value", DataType::Int32, false)), + ]; + let entries = StructArray::try_new( + fields.clone().into(), + vec![Arc::clone(&values), Arc::clone(&values)], + None, + ) + .unwrap(); + let entry = Arc::new(Field::new( + "entries", + DataType::Struct(fields.into()), + false, + )); + let map = MapArray::try_new(Arc::clone(&entry), offsets, entries, None, false) + .unwrap() + .slice(value_count as usize / 2, 1); + let map_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "entries", + DataType::Map(entry, false), + false, + )])), + vec![Arc::new(map)], + ) + .unwrap(); + + let frame_limit = 2_048; + let admission_capacity = 3 * frame_limit; + for input in [list_batch, map_batch] { + let raw_backing_size = + RssPartitionWriter::::estimated_ipc_data_size(&input) + .unwrap(); + assert!(raw_backing_size > admission_capacity); + + let recorder = ReservationRecordingPusher::default(); + let recorded = Arc::clone(&recorder.0); + let pusher = BoundedAdmissionPusher { + recorder, + capacity: admission_capacity, + }; + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None) + .unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, frame_limit).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input.clone())].into_iter(), &metrics) + .unwrap(); + + let recorded = recorded.lock().unwrap(); + assert_eq!(recorded.reservations.len(), 1); + assert!(recorded.reservations[0] <= admission_capacity); + assert_eq!(recorded.frames.len(), 1); + assert_eq!( + read_ipc_compressed(&recorded.frames[0][16..]).unwrap(), + input + ); + } + } + + #[test] + fn rss_partition_writer_counts_validity_offsets_and_view_backing_buffers() { + let booleans: ArrayRef = Arc::new(BooleanArray::from(vec![true; 8_192])); + let boolean_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Boolean, + false, + )])), + vec![Arc::clone(&booleans)], + ) + .unwrap(); + let raw_boolean_size = booleans.to_data().get_slice_memory_size().unwrap(); + let estimated_boolean_size = + RssPartitionWriter::::estimated_ipc_data_size(&boolean_batch).unwrap(); + assert!(estimated_boolean_size >= raw_boolean_size + 1_024); + + let strings: ArrayRef = Arc::new(StringArray::from(vec!["x"; 8_192])); + let sliced_strings = strings.slice(1, 8_191); + let string_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8, + false, + )])), + vec![Arc::clone(&sliced_strings)], + ) + .unwrap(); + let raw_string_size = sliced_strings.to_data().get_slice_memory_size().unwrap(); + let estimated_string_size = + RssPartitionWriter::::estimated_ipc_data_size(&string_batch).unwrap(); + assert!(estimated_string_size >= raw_string_size + 8_192 * 4); + + let views: ArrayRef = Arc::new(StringViewArray::from(vec![ + "small".to_owned(), + "unused".repeat(1_024), + ])); + let sliced = views.slice(0, 1); + let view_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8View, + false, + )])), + vec![Arc::clone(&sliced)], + ) + .unwrap(); + assert!(sliced.to_data().get_slice_memory_size().unwrap() < 128); + assert!( + RssPartitionWriter::::estimated_ipc_data_size(&view_batch).unwrap() + > 6_000 + ); + } + + #[test] + fn rss_partition_writer_reserves_only_batches_that_are_small_enough_to_encode() { + #[derive(Default)] + struct ReservationState { + reserved: bool, + reservations: usize, + releases: usize, + pushes: usize, + } + + struct TrackingPusher(Arc>); + + impl PartitionPusher for TrackingPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + let mut state = self.0.lock().unwrap(); + assert!(!state.reserved); + assert_eq!(reservation_bytes, 3 * 1_024); + state.reserved = true; + state.reservations += 1; + Ok(()) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + let mut state = self.0.lock().unwrap(); + assert!(state.reserved); + state.reserved = false; + state.releases += 1; + Ok(()) + } + + fn push_partition_data(&self, _partition_id: usize, _frame: &[u8]) -> Result<()> { + let mut state = self.0.lock().unwrap(); + assert!(state.reserved); + state.reserved = false; + state.pushes += 1; + Ok(()) + } + } + + let input = batch((0..512).collect()); + let state = Arc::new(Mutex::new(ReservationState::default())); + let pusher = TrackingPusher(Arc::clone(&state)); + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, 1_024).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap(); + + let state = state.lock().unwrap(); + assert!(!state.reserved); + assert!(state.pushes > 1); + // Wide batches are now split before admission, so no reservation is wasted on an + // encoding attempt already known to exceed the Arrow IPC scratch-space budget. + assert_eq!(state.releases, 0); + assert_eq!(state.reservations, state.releases + state.pushes); + } + + #[test] + fn rss_partition_writer_compacts_dictionary_values_before_splitting_frames() { + let values: Vec = (0..256) + .map(|index| { + format!( + "value-{index:03}-{}", + (0..72) + .map(|offset| char::from(b'a' + ((index + offset) % 26) as u8)) + .collect::() + ) + }) + .collect(); + let expected: Vec = (0..96).map(|index| values[index * 2].clone()).collect(); + let dictionary = DictionaryArray::::try_new( + Int32Array::from((0..96).map(|index| index * 2).collect::>()), + Arc::new(StringArray::from(values.clone())), + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + )])); + let input = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(dictionary)]).unwrap(); + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, 2_048).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = captured.lock().unwrap(); + assert!(frames.len() > 1); + let mut actual = Vec::new(); + for (_, frame) in frames.iter() { + assert!(frame.len() <= 2_048); + let decoded = read_ipc_compressed(&frame[16..]).unwrap(); + let dictionary = decoded + .column(0) + .as_any() + .downcast_ref::>() + .unwrap(); + assert!(dictionary.values().len() < values.len()); + let plain = cast(dictionary, &DataType::Utf8).unwrap(); + let strings = plain.as_any().downcast_ref::().unwrap(); + actual.extend(strings.iter().map(|value| value.unwrap().to_owned())); + } + assert_eq!(actual, expected); + } + + #[test] + fn rss_partition_writer_compacts_an_oversized_dictionary_for_a_single_row() { + let mut values = vec!["small".to_owned()]; + values.extend((0..64).map(|index| format!("unused-{index}-{}", "x".repeat(256)))); + let dictionary = DictionaryArray::::try_new( + Int32Array::from(vec![0]), + Arc::new(StringArray::from(values)), + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(dictionary)]).unwrap(); + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, 1_024).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(batch)].into_iter(), &metrics) + .unwrap(); + + let frames = captured.lock().unwrap(); + assert_eq!(frames.len(), 1); + let decoded = read_ipc_compressed(&frames[0].1[16..]).unwrap(); + let dictionary = decoded + .column(0) + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(dictionary.values().len(), 1); + assert_eq!( + dictionary + .values() + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "small" + ); + } + + #[test] + fn rss_partition_writer_compacts_dictionaries_nested_in_sliced_lists_and_structs() { + let values = (0..128) + .map(|index| format!("value-{index:03}-{}", "x".repeat(128))) + .collect::>(); + let dictionary_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let dictionary = DictionaryArray::::try_new( + Int32Array::from((0..32).collect::>()), + Arc::new(StringArray::from(values.clone())), + ) + .unwrap(); + let item = Arc::new(Field::new("item", dictionary_type, false)); + let lists = ListArray::try_new( + Arc::clone(&item), + OffsetBuffer::new((0..=32).collect::>().into()), + Arc::new(dictionary), + None, + ) + .unwrap(); + let nested = Arc::new(Field::new("items", DataType::List(item), false)); + let structure = StructArray::try_new( + vec![Arc::clone(&nested)].into(), + vec![Arc::new(lists)], + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "nested", + DataType::Struct(vec![nested].into()), + false, + )])); + let input = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(structure)]) + .unwrap() + .slice(7, 16); + let pusher = RecordingPusher::default(); + let captured = Arc::clone(&pusher.0); + let encoder = ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::None).unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, 2_048).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input)].into_iter(), &metrics) + .unwrap(); + + let frames = captured.lock().unwrap(); + assert!(frames.len() > 1); + let mut actual = Vec::new(); + for (_, frame) in frames.iter() { + assert!(frame.len() <= 2_048); + let decoded = read_ipc_compressed(&frame[16..]).unwrap(); + let structure = decoded + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let lists = structure + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(lists.value_offsets()[0], 0); + let dictionary = lists + .values() + .as_any() + .downcast_ref::>() + .unwrap(); + assert!(dictionary.values().len() < values.len()); + let plain = cast(dictionary, &DataType::Utf8).unwrap(); + actual.extend( + plain + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|value| value.unwrap().to_owned()), + ); + } + assert_eq!(actual, values[7..23].to_vec()); + } + #[test] fn rss_partition_writer_preserves_transport_errors() { struct FailingPusher; diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java index 61cc5abe326..d5dcdd79edc 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java +++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java @@ -20,23 +20,116 @@ package org.apache.comet.shuffle; import java.io.IOException; +import java.lang.reflect.Field; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.lang.reflect.Modifier; +import java.lang.reflect.Proxy; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.IdentityHashMap; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLongArray; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.atomic.LongAdder; /** Sends complete native Comet frames through an existing, task-scoped Celeborn shuffle client. */ public final class CelebornShufflePartitionPusher implements ShufflePartitionPusher { // Celeborn prefixes every accepted payload with four transport-level integers. private static final int CELEBORN_BATCH_HEADER_BYTES = 4 * Integer.BYTES; + private static final int DEFAULT_MAX_IN_FLIGHT_BYTES = 256 * 1024 * 1024; + private static final long COMPLETION_RECONCILIATION_INTERVAL_MILLIS = 10; + private static final StackWalker COMPLETION_CALL_STACK = StackWalker.getInstance(); + // One daemon only samples counters; it owns no client/configuration and stops retaining idle + // maps. + private static final ScheduledThreadPoolExecutor COMPLETION_RECONCILER = + createCompletionReconciler(); private final Object shuffleClient; private final Method pushOrMergeData; + private final Method mapperEnd; + private final Method cleanup; + private final Method getPushState; + private final Method setPushMetricsCallback; + private final Class pushMetricsCallbackClass; + private final Field inFlightRequestTracker; + private final Field totalInFlightRequests; + private final Field pushStateException; + private final ExecutorShufflePushAdmission admission; private final int shuffleId; private final int mapId; private final int encodedAttemptId; private final int numMappers; private final int numPartitions; + private final int maxFrameBytes; + private final int maxReservationBytes; + private final AtomicLongArray partitionLengths; + private final Object lifecycleLock = new Object(); + private final Object cleanupLock = new Object(); + private final ThreadLocal encodingReservation = new ThreadLocal<>(); + private final ArrayDeque pendingPushes = new ArrayDeque<>(); + private final IdentityHashMap observedPushStates = + new IdentityHashMap<>(); + + private State state = State.OPEN; + private int activePushes; + private int activeClientPushes; + private int activeEncoders; + private ScheduledFuture completionReconciliation; + private Thread mapperEndThread; + private boolean cleanupStarted; + private boolean cleanupAfterActivePushes; + private boolean finalCleanupStarted; + + private enum State { + OPEN, + FINISHING, + FINISHED, + ABORTED + } + + private static final class PushReservation { + private int bytes; + private ObservedPushState pushState; + private boolean submitted; + + private PushReservation(int bytes) { + this.bytes = bytes; + } + } + + private static final class ObservedPushState { + private final LongAdder inFlightRequests; + private final AtomicReference exception; + private long submittedPushes; + private long metricsCompletions; + private long metricsOnlyFailures; + private long releasedPushes; + private boolean failedBeforeObservation; + private boolean recoveredMissedFailure; + + private ObservedPushState(LongAdder inFlightRequests, AtomicReference exception) { + this.inFlightRequests = inFlightRequests; + this.exception = exception; + } + } + + private static ScheduledThreadPoolExecutor createCompletionReconciler() { + ScheduledThreadPoolExecutor executor = + new ScheduledThreadPoolExecutor( + 1, + runnable -> { + Thread thread = + new Thread(runnable, "comet-celeborn-shuffle-push-admission-reconciler"); + thread.setDaemon(true); + return thread; + }); + executor.setRemoveOnCancelPolicy(true); + return executor; + } /** * Binds the existing Celeborn client and all shuffle identity to one map task. @@ -59,6 +152,25 @@ public CelebornShufflePartitionPusher( int encodedAttemptId, int numMappers, int numPartitions) { + this( + shuffleClient, + shuffleId, + mapId, + encodedAttemptId, + numMappers, + numPartitions, + DEFAULT_MAX_IN_FLIGHT_BYTES); + } + + /** Binds a task-owned adapter to byte admission shared by its executor-side Celeborn client. */ + public CelebornShufflePartitionPusher( + Object shuffleClient, + int shuffleId, + int mapId, + int encodedAttemptId, + int numMappers, + int numPartitions, + int maxInFlightBytes) { if (shuffleClient == null) { throw new IllegalArgumentException("Celeborn shuffle client must not be null"); } @@ -80,10 +192,13 @@ public CelebornShufflePartitionPusher( if (numPartitions <= 0) { throw new IllegalArgumentException("Celeborn partition count must be positive"); } + if (maxInFlightBytes <= CELEBORN_BATCH_HEADER_BYTES) { + throw new IllegalArgumentException("Celeborn in-flight byte limit must fit a request"); + } - Method method; + Method pushOrMergeData; try { - method = + pushOrMergeData = shuffleClient .getClass() .getMethod( @@ -103,18 +218,134 @@ public CelebornShufflePartitionPusher( throw new IllegalArgumentException( "Celeborn shuffle client does not provide the required public raw-push API", e); } - if (method.getReturnType() != int.class || Modifier.isStatic(method.getModifiers())) { + if (pushOrMergeData.getReturnType() != int.class + || Modifier.isStatic(pushOrMergeData.getModifiers())) { throw new IllegalArgumentException( "Celeborn raw-push API must be an instance method returning an int"); } + Method mapperEndMethod = + resolveLifecycleMethod( + shuffleClient, "mapperEnd", int.class, int.class, int.class, int.class); + Method cleanupMethod = + resolveLifecycleMethod(shuffleClient, "cleanup", int.class, int.class, int.class); + final Method getPushStateMethod; + final Method setPushMetricsCallbackMethod; + final Class pushMetricsCallbackClass; + final Field inFlightRequestTrackerField; + final Field totalInFlightRequestsField; + final Field pushStateExceptionField; + try { + getPushStateMethod = shuffleClient.getClass().getMethod("getPushState", String.class); + setPushMetricsCallbackMethod = resolvePushMetricsCallback(getPushStateMethod.getReturnType()); + pushMetricsCallbackClass = setPushMetricsCallbackMethod.getParameterTypes()[0]; + pushMetricsCallbackClass.getMethod("incPushDataCount", long.class); + inFlightRequestTrackerField = + getPushStateMethod.getReturnType().getDeclaredField("inFlightRequestTracker"); + totalInFlightRequestsField = + inFlightRequestTrackerField.getType().getDeclaredField("totalInflightReqs"); + pushStateExceptionField = getPushStateMethod.getReturnType().getDeclaredField("exception"); + if (totalInFlightRequestsField.getType() != LongAdder.class) { + throw new NoSuchFieldException("InFlightRequestTracker.totalInflightReqs: LongAdder"); + } + if (pushStateExceptionField.getType() != AtomicReference.class) { + throw new NoSuchFieldException("PushState.exception: AtomicReference"); + } + inFlightRequestTrackerField.setAccessible(true); + totalInFlightRequestsField.setAccessible(true); + pushStateExceptionField.setAccessible(true); + } catch (ReflectiveOperationException | RuntimeException e) { + throw new IllegalArgumentException( + "Celeborn shuffle client does not provide observable completion-backed push admission", + e); + } + if (Modifier.isStatic(getPushStateMethod.getModifiers()) + || Modifier.isStatic(setPushMetricsCallbackMethod.getModifiers()) + || setPushMetricsCallbackMethod.getReturnType() != void.class + || !pushMetricsCallbackClass.isInterface()) { + throw new IllegalArgumentException("Celeborn push-state completion API is incompatible"); + } + this.shuffleClient = shuffleClient; - this.pushOrMergeData = method; + this.pushOrMergeData = pushOrMergeData; + this.mapperEnd = mapperEndMethod; + this.cleanup = cleanupMethod; + this.getPushState = getPushStateMethod; + this.setPushMetricsCallback = setPushMetricsCallbackMethod; + this.pushMetricsCallbackClass = pushMetricsCallbackClass; + this.inFlightRequestTracker = inFlightRequestTrackerField; + this.totalInFlightRequests = totalInFlightRequestsField; + this.pushStateException = pushStateExceptionField; + this.admission = ExecutorShufflePushAdmission.forClient(shuffleClient, maxInFlightBytes); this.shuffleId = shuffleId; this.mapId = mapId; this.encodedAttemptId = encodedAttemptId; this.numMappers = numMappers; this.numPartitions = numPartitions; + this.maxReservationBytes = maxInFlightBytes - CELEBORN_BATCH_HEADER_BYTES; + this.maxFrameBytes = maxReservationBytes / 3; + this.partitionLengths = new AtomicLongArray(numPartitions); + } + + private static Method resolvePushMetricsCallback(Class pushStateClass) + throws NoSuchMethodException { + for (Method method : pushStateClass.getMethods()) { + if (method.getName().equals("setMetricsCallback") && method.getParameterCount() == 1) { + return method; + } + } + throw new NoSuchMethodException("PushState.setMetricsCallback"); + } + + private static Method resolveLifecycleMethod( + Object shuffleClient, String name, Class... parameterTypes) { + final Method method; + try { + method = shuffleClient.getClass().getMethod(name, parameterTypes); + } catch (NoSuchMethodException | SecurityException e) { + throw new IllegalArgumentException( + "Celeborn shuffle client does not provide the required public " + name + " API", e); + } + if (method.getReturnType() != void.class || Modifier.isStatic(method.getModifiers())) { + throw new IllegalArgumentException( + "Celeborn " + name + " API must be an instance method returning void"); + } + return method; + } + + @Override + public void reservePartitionData(int maxLength) throws IOException { + if (maxLength <= 0 || maxLength > maxReservationBytes) { + throw new IOException("Celeborn native frame reservation exceeds its byte limit"); + } + if (encodingReservation.get() != null) { + throw new IOException("Celeborn native frame already has a reservation on this thread"); + } + + final int requestBytes = maxLength + CELEBORN_BATCH_HEADER_BYTES; + admission.acquire(requestBytes, this::isAborted); + synchronized (lifecycleLock) { + if (state != State.OPEN) { + admission.release(requestBytes); + throw new IOException("Celeborn shuffle map attempt no longer accepts frame encoding"); + } + activeEncoders++; + encodingReservation.set(new PushReservation(requestBytes)); + } + } + + @Override + public void releasePartitionDataReservation() { + final PushReservation reservation = encodingReservation.get(); + if (reservation == null) { + return; + } + encodingReservation.remove(); + synchronized (lifecycleLock) { + activeEncoders--; + lifecycleLock.notifyAll(); + } + admission.release(reservation.bytes); } @Override @@ -132,52 +363,537 @@ public int pushPartitionData(int partitionId, byte[] bytes, int length) throws I throw new IOException("Celeborn shuffle frame and transport header exceed the byte limit"); } - final int accepted; + beginPush(); + final int requestBytes = length + CELEBORN_BATCH_HEADER_BYTES; + PushReservation reservation = null; + boolean registered = false; + boolean submitted = false; + Throwable pushFailure = null; try { + reservation = claimEncodingReservation(requestBytes, length); + final Object pushState = + getPushState.invoke(shuffleClient, shuffleId + "-" + mapId + "-" + encodedAttemptId); + ObservedPushState observedPushState = observePushState(pushState); + // doPush=true sends this complete frame immediately. skipCompress=true preserves the // existing Comet frame, which already owns its compression and framing format. - accepted = - (int) - pushOrMergeData.invoke( - shuffleClient, - shuffleId, - mapId, - encodedAttemptId, - partitionId, - bytes, - 0, - length, - numMappers, - numPartitions, - true, - true); + beginClientPush(reservation, observedPushState); + registered = true; + final int accepted; + try { + accepted = + (int) + pushOrMergeData.invoke( + shuffleClient, + shuffleId, + mapId, + encodedAttemptId, + partitionId, + bytes, + 0, + length, + numMappers, + numPartitions, + true, + true); + submitted = accepted > 0; + + if (submitted) { + // The native frame and JNI payload overlap Celeborn's copied request only until its + // raw-push call returns. Completion callbacks can run inline inside that call, so keep + // all three copies charged while the reservation remains unsubmitted. + shrinkSubmittedReservation(reservation, requestBytes); + } + + if (submitted && isAborted()) { + // Cleanup can remove the initial state while Celeborn resolves locations. Capture the + // recreated state's actual in-flight count before this submission becomes observable. + final Object resumedPushState = + getPushState.invoke(shuffleClient, shuffleId + "-" + mapId + "-" + encodedAttemptId); + observedPushState = observePushState(resumedPushState); + if (resumedPushState != pushState) { + recoverUnobservedFailure(observedPushState); + } + } + if (submitted) { + acceptClientPush(reservation, observedPushState); + } + } finally { + endClientPush(); + } + + if (isAborted()) { + throw new IOException("Celeborn shuffle map attempt was aborted during its push"); + } + + int expected = requestBytes; + if (accepted != expected) { + throw new IOException( + "Celeborn raw shuffle push accepted " + + accepted + + " bytes; expected " + + expected + + " including its transport header"); + } + + partitionLengths.addAndGet(partitionId, length); + + // ShufflePartitionPusher reports Comet payload bytes, never Celeborn transport bytes. + return length; } catch (IllegalAccessException e) { - throw new IOException("Cannot invoke the public Celeborn raw-push API", e); + IOException failure = new IOException("Cannot invoke the public Celeborn raw-push API", e); + pushFailure = failure; + abortAndSuppress(failure); + throw failure; } catch (InvocationTargetException e) { - Throwable cause = e.getCause(); - if (cause instanceof IOException) { - throw (IOException) cause; + Throwable failure = unwrapFailure("Celeborn raw shuffle push failed", e); + pushFailure = failure; + abortAndSuppress(failure); + throwFailure(failure); + throw new AssertionError("unreachable"); + } catch (IOException | RuntimeException | Error failure) { + pushFailure = failure; + abortAndSuppress(failure); + throw failure; + } finally { + if (reservation != null && !submitted) { + if (registered) { + releaseUnsubmittedPush(reservation); + } else { + admission.release(reservation.bytes); + } + } + try { + endPush(); + } catch (IOException cleanupFailure) { + if (pushFailure == null) { + throw cleanupFailure; + } + if (cleanupFailure != pushFailure) { + pushFailure.addSuppressed(cleanupFailure); + } + } + } + } + + /** + * Drain all asynchronous Celeborn pushes, commit this map attempt, and return its payload sizes. + * + *

All native frames use {@code doPush=true}, so there is no merge buffer to flush. The + * deployed Celeborn client's {@code mapperEnd} waits for every in-flight push, propagates any + * asynchronous failure, and reports map completion to its lifecycle manager. + */ + public long[] finish() throws IOException { + try { + synchronized (lifecycleLock) { + if (state == State.FINISHED) { + return snapshotPartitionLengths(); + } + if (state != State.OPEN) { + throw new IOException("Celeborn shuffle map attempt is not available for completion"); + } + state = State.FINISHING; + while ((activePushes != 0 || activeEncoders != 0) && state == State.FINISHING) { + lifecycleLock.wait(); + } + if (state != State.FINISHING) { + throw new IOException("Celeborn shuffle map attempt was aborted before completion"); + } + } + + synchronized (lifecycleLock) { + if (state != State.FINISHING) { + throw new IOException("Celeborn shuffle map attempt was aborted before mapperEnd"); + } + mapperEndThread = Thread.currentThread(); } - if (cause instanceof RuntimeException) { - throw (RuntimeException) cause; + try { + mapperEnd.invoke(shuffleClient, shuffleId, mapId, encodedAttemptId, numMappers); + } finally { + synchronized (lifecycleLock) { + mapperEndThread = null; + } } - if (cause instanceof Error) { - throw (Error) cause; + + synchronized (lifecycleLock) { + if (state != State.FINISHING) { + throw new IOException("Celeborn shuffle map attempt was aborted during completion"); + } + state = State.FINISHED; + lifecycleLock.notifyAll(); } - throw new IOException("Celeborn raw shuffle push failed", cause); + return snapshotPartitionLengths(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + IOException failure = + new IOException("Interrupted while draining Celeborn shuffle pushes", e); + abortAndSuppress(failure); + throw failure; + } catch (IllegalAccessException e) { + IOException failure = new IOException("Cannot invoke the public Celeborn mapperEnd API", e); + abortAndSuppress(failure); + throw failure; + } catch (InvocationTargetException e) { + Throwable failure = unwrapFailure("Celeborn shuffle map completion failed", e); + abortAndSuppress(failure); + throwFailure(failure); + throw new AssertionError("unreachable"); + } catch (IOException | RuntimeException | Error failure) { + abortAndSuppress(failure); + throw failure; } + } - int expected = length + CELEBORN_BATCH_HEADER_BYTES; - if (accepted != expected) { - throw new IOException( - "Celeborn raw shuffle push accepted " - + accepted - + " bytes; expected " - + expected - + " including its transport header"); + /** Cancel this attempt and wake any Celeborn push blocked on transport backpressure. */ + public void abort() throws IOException { + final boolean shouldCleanup; + final Thread completionThread; + synchronized (lifecycleLock) { + if (state != State.FINISHED) { + state = State.ABORTED; + } + shouldCleanup = !cleanupStarted; + if (shouldCleanup) { + cleanupAfterActivePushes = activeClientPushes > 0; + } + cleanupStarted = true; + completionThread = mapperEndThread; + lifecycleLock.notifyAll(); } + if (completionThread != null && completionThread != Thread.currentThread()) { + completionThread.interrupt(); + } + if (!shouldCleanup) { + return; + } + + cleanupAttempt(); + } - // ShufflePartitionPusher reports Comet payload bytes, never Celeborn transport bytes. - return length; + private void cleanupAttempt() throws IOException { + synchronized (cleanupLock) { + try { + cleanup.invoke(shuffleClient, shuffleId, mapId, encodedAttemptId); + } catch (IllegalAccessException e) { + throw new IOException("Cannot invoke the public Celeborn cleanup API", e); + } catch (InvocationTargetException e) { + throwFailure(unwrapFailure("Celeborn shuffle map cleanup failed", e)); + } + } + } + + private void beginPush() throws IOException { + synchronized (lifecycleLock) { + if (state != State.OPEN) { + throw new IOException("Celeborn shuffle map attempt no longer accepts partition data"); + } + activePushes++; + } + } + + private PushReservation claimEncodingReservation(int requestBytes, int frameBytes) + throws IOException { + PushReservation reservation = encodingReservation.get(); + if (reservation != null) { + final long overlappingBytes = 3L * frameBytes + (long) CELEBORN_BATCH_HEADER_BYTES; + if (overlappingBytes > reservation.bytes) { + throw new IOException("Celeborn shuffle push exceeds its native encoding reservation"); + } + encodingReservation.remove(); + synchronized (lifecycleLock) { + activeEncoders--; + lifecycleLock.notifyAll(); + } + if (overlappingBytes < reservation.bytes) { + admission.release(reservation.bytes - (int) overlappingBytes); + return new PushReservation((int) overlappingBytes); + } + return reservation; + } + + admission.acquire(requestBytes, this::isAborted); + return new PushReservation(requestBytes); + } + + private void shrinkSubmittedReservation(PushReservation reservation, int requestBytes) { + final int releasedBytes; + synchronized (lifecycleLock) { + releasedBytes = reservation.bytes - requestBytes; + reservation.bytes = requestBytes; + } + if (releasedBytes > 0) { + admission.release(releasedBytes); + } + } + + private ObservedPushState observePushState(Object pushState) + throws IllegalAccessException, InvocationTargetException { + synchronized (lifecycleLock) { + ObservedPushState observed = observedPushStates.get(pushState); + if (observed != null) { + return observed; + } + + final Object tracker = inFlightRequestTracker.get(pushState); + final ObservedPushState created = + new ObservedPushState( + (LongAdder) totalInFlightRequests.get(tracker), + (AtomicReference) pushStateException.get(pushState)); + Object callback = + Proxy.newProxyInstance( + pushMetricsCallbackClass.getClassLoader(), + new Class[] {pushMetricsCallbackClass}, + (proxy, method, arguments) -> { + if (method.getName().equals("incPushDataCount")) { + boolean metricsOnlyFailure = + created.inFlightRequests.sum() > 0 && isTerminalFailureCallback(); + completeAcceptedPushes(created, (long) arguments[0], metricsOnlyFailure); + } else if (method.getName().equals("hashCode")) { + return System.identityHashCode(proxy); + } else if (method.getName().equals("equals")) { + return proxy == arguments[0]; + } else if (method.getName().equals("toString")) { + return "Celeborn native shuffle push completion callback"; + } + return null; + }); + setPushMetricsCallback.invoke(pushState, callback); + created.failedBeforeObservation = created.exception.get() != null; + observedPushStates.put(pushState, created); + return created; + } + } + + private void recoverUnobservedFailure(ObservedPushState pushState) { + synchronized (lifecycleLock) { + if (pushState.failedBeforeObservation && !pushState.recoveredMissedFailure) { + // A recreated state's inline terminal failure used Celeborn's NOOP metrics callback. The + // pinned failure path leaves one tracker batch behind but records its public exception. + pushState.metricsOnlyFailures++; + pushState.recoveredMissedFailure = true; + } + } + } + + private void beginClientPush(PushReservation reservation, ObservedPushState pushState) + throws IOException { + synchronized (lifecycleLock) { + if (state == State.ABORTED) { + throw new IOException("Celeborn shuffle map attempt was aborted before its client push"); + } + activeClientPushes++; + reservation.pushState = pushState; + pendingPushes.addLast(reservation); + ensureCompletionReconciliation(); + } + } + + private void acceptClientPush(PushReservation reservation, ObservedPushState pushState) { + synchronized (lifecycleLock) { + reservation.pushState = pushState; + reservation.submitted = true; + pushState.submittedPushes++; + } + reconcileAcceptedPushes(); + } + + private static boolean isTerminalFailureCallback() { + // In the pinned raw-push client only onFailure emits metrics without removing its batch. + // Skip the stack walk for the common final successful request, whose tracker is already zero. + return COMPLETION_CALL_STACK.walk( + frames -> + frames + .filter( + frame -> + frame.getMethodName().equals("onFailure") + || frame.getMethodName().equals("onSuccess")) + .findFirst() + .map(frame -> frame.getMethodName().equals("onFailure")) + .orElse(false)); + } + + private void completeAcceptedPushes( + ObservedPushState pushState, long completed, boolean metricsOnlyFailure) { + synchronized (lifecycleLock) { + if (completed > 0) { + pushState.metricsCompletions += completed; + if (metricsOnlyFailure) { + pushState.metricsOnlyFailures += completed; + } + } + } + reconcileAcceptedPushes(); + } + + private void ensureCompletionReconciliation() { + if (completionReconciliation == null || completionReconciliation.isDone()) { + completionReconciliation = + COMPLETION_RECONCILER.scheduleWithFixedDelay( + this::reconcileAcceptedPushes, + COMPLETION_RECONCILIATION_INTERVAL_MILLIS, + COMPLETION_RECONCILIATION_INTERVAL_MILLIS, + TimeUnit.MILLISECONDS); + } + } + + /** + * Celeborn's MAP_ENDED retry paths remove their transport batch without invoking its metrics + * callback. The pinned tracker preserves its in-flight request count even after cleanup, so the + * actual count also reconciles silent completions without releasing cancelled live requests. + */ + private void reconcileAcceptedPushes() { + final ArrayList completed = new ArrayList<>(); + synchronized (lifecycleLock) { + for (ObservedPushState pushState : observedPushStates.values()) { + long transportRequests = + Math.max(0L, pushState.inFlightRequests.sum() - pushState.metricsOnlyFailures); + long trackerCompletions = Math.max(0L, pushState.submittedPushes - transportRequests); + long knownCompletions = + hasUnsubmittedReservation(pushState) + ? trackerCompletions + : Math.max(pushState.metricsCompletions, trackerCompletions); + while (pushState.releasedPushes < knownCompletions) { + PushReservation reservation = smallestPendingReservation(pushState); + if (reservation == null) { + break; + } + pendingPushes.remove(reservation); + pushState.releasedPushes++; + completed.add(reservation); + } + } + + if (pendingPushes.isEmpty() && activeClientPushes == 0) { + stopCompletionReconciliation(); + } + } + for (PushReservation reservation : completed) { + admission.release(reservation.bytes); + } + } + + private PushReservation smallestPendingReservation(ObservedPushState pushState) { + // Metrics report a completion count without its batch identity. Releasing the smallest + // eligible reservation never restores more bytes than the request that actually completed. + PushReservation smallest = null; + for (PushReservation reservation : pendingPushes) { + if (reservation.pushState == pushState + && reservation.submitted + && (smallest == null || reservation.bytes < smallest.bytes)) { + smallest = reservation; + } + } + return smallest; + } + + private boolean hasUnsubmittedReservation(ObservedPushState pushState) { + for (PushReservation reservation : pendingPushes) { + if (reservation.pushState == pushState && !reservation.submitted) { + return true; + } + } + return false; + } + + private void stopCompletionReconciliation() { + if (completionReconciliation != null) { + completionReconciliation.cancel(false); + completionReconciliation = null; + } + } + + private void releaseUnsubmittedPush(PushReservation reservation) { + final boolean pending; + synchronized (lifecycleLock) { + pending = pendingPushes.remove(reservation); + if (pendingPushes.isEmpty() && activeClientPushes == 0) { + stopCompletionReconciliation(); + } + } + if (pending) { + admission.release(reservation.bytes); + } + } + + private void endClientPush() { + synchronized (lifecycleLock) { + activeClientPushes--; + if (pendingPushes.isEmpty() && activeClientPushes == 0) { + stopCompletionReconciliation(); + } + } + } + + private void endPush() throws IOException { + final boolean cleanupAfterPush; + synchronized (lifecycleLock) { + activePushes--; + if (activePushes == 0) { + lifecycleLock.notifyAll(); + } + cleanupAfterPush = + activePushes == 0 + && state == State.ABORTED + && cleanupAfterActivePushes + && !finalCleanupStarted; + if (cleanupAfterPush) { + finalCleanupStarted = true; + } + } + if (cleanupAfterPush) { + cleanupAttempt(); + } + } + + private boolean isAborted() { + synchronized (lifecycleLock) { + return state == State.ABORTED; + } + } + + public int numPartitions() { + return numPartitions; + } + + /** Largest native frame whose overlapping native, JNI, and Celeborn copies fit admission. */ + public int maxFrameBytes() { + return maxFrameBytes; + } + + private long[] snapshotPartitionLengths() { + long[] sizes = new long[numPartitions]; + for (int partition = 0; partition < numPartitions; partition++) { + sizes[partition] = partitionLengths.get(partition); + } + return sizes; + } + + private void abortAndSuppress(Throwable original) { + try { + abort(); + } catch (Throwable cleanupFailure) { + if (cleanupFailure != original) { + original.addSuppressed(cleanupFailure); + } + } + } + + private static Throwable unwrapFailure(String message, InvocationTargetException failure) { + Throwable cause = failure.getCause(); + return cause instanceof IOException + || cause instanceof RuntimeException + || cause instanceof Error + ? cause + : new IOException(message, cause); + } + + private static void throwFailure(Throwable failure) throws IOException { + if (failure instanceof IOException) { + throw (IOException) failure; + } + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + throw (Error) failure; } } diff --git a/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java b/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java new file mode 100644 index 00000000000..4abd8cf95c0 --- /dev/null +++ b/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java @@ -0,0 +1,93 @@ +/* + * 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; +import java.util.IdentityHashMap; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.function.BooleanSupplier; + +/** Completion-backed byte admission shared by all map attempts using one Celeborn client. */ +final class ExecutorShufflePushAdmission { + private static final IdentityHashMap BY_CLIENT = + new IdentityHashMap<>(); + + private final int limit; + private final Semaphore available; + + private ExecutorShufflePushAdmission(int limit) { + this.limit = limit; + this.available = new Semaphore(limit, true); + } + + static ExecutorShufflePushAdmission forClient(Object client, int limit) { + synchronized (BY_CLIENT) { + ExecutorShufflePushAdmission admission = BY_CLIENT.get(client); + if (admission == null) { + admission = new ExecutorShufflePushAdmission(limit); + BY_CLIENT.put(client, admission); + } else if (admission.limit != limit) { + throw new IllegalArgumentException( + "All map attempts sharing a Celeborn client must use the same in-flight byte limit"); + } + return admission; + } + } + + static void releaseClient(Object client) { + synchronized (BY_CLIENT) { + BY_CLIENT.remove(client); + } + } + + void acquire(int bytes, BooleanSupplier cancelled) throws IOException { + if (bytes > limit) { + throw new IOException( + "Celeborn request requires " + + bytes + + " bytes, exceeding the executor in-flight byte limit of " + + limit); + } + + try { + while (true) { + if (cancelled.getAsBoolean()) { + throw new IOException("Celeborn shuffle map attempt was cancelled during admission"); + } + if (available.tryAcquire(bytes, 25, TimeUnit.MILLISECONDS)) { + if (cancelled.getAsBoolean()) { + available.release(bytes); + throw new IOException("Celeborn shuffle map attempt was cancelled during admission"); + } + return; + } + } + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new IOException( + "Interrupted while waiting for Celeborn shuffle push admission", failure); + } + } + + void release(int bytes) { + available.release(bytes); + } +} diff --git a/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java index 35a59d243e8..c3161b3fe99 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java +++ b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java @@ -25,6 +25,17 @@ @FunctionalInterface public interface ShufflePartitionPusher { + /** + * Reserve executor-wide transport capacity before native frame encoding or JNI array allocation. + * + *

The reservation belongs to the calling thread until its next successful push. Backends + * without asynchronous admission can retain the default no-op behavior. + */ + default void reservePartitionData(int maxLength) throws IOException {} + + /** Release an unsubmitted reservation after native encoding fails or must split its batch. */ + default void releasePartitionDataReservation() {} + /** * Accepts one complete, independently decodable Comet frame for an output partition. * diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 475b73ed692..31a103a3998 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -605,6 +605,32 @@ object CometConf extends ShimCometConf { .checkValue(v => v >= 0, "Must not be negative") .createWithDefault(0) + val COMET_SHUFFLE_RSS_MAX_FRAME_BYTES: ConfigEntry[Long] = + conf("spark.comet.shuffle.rss.maxFrameBytes") + .category(CATEGORY_SHUFFLE) + .doc("Maximum encoded size of one complete native shuffle frame sent to a remote " + + "shuffle service. Frames are never split across remote push requests.") + .bytesConf(ByteUnit.BYTE) + .checkValue( + value => value >= 20 && value <= Int.MaxValue - 16, + "Remote shuffle frame size must fit a complete Comet frame and a Celeborn request") + .createWithDefault(64L * 1024 * 1024) + + val COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES: ConfigEntry[Long] = + conf("spark.comet.shuffle.rss.maxInFlightBytes") + .category(CATEGORY_SHUFFLE) + .doc( + "Maximum shuffle bytes admitted concurrently by all native Comet map attempts " + + "sharing an executor-side Celeborn client. Admission covers native encoding " + + "scratch and overlapping native, JNI, and Celeborn frame copies during submission; " + + "accepted request bytes remain reserved until Celeborn reports completion.") + .bytesConf(ByteUnit.BYTE) + .checkValue( + value => value >= 76 && value <= Int.MaxValue, + "Remote shuffle in-flight byte limit must fit three complete frame copies and a " + + "Celeborn request header") + .createWithDefault(256L * 1024 * 1024) + val COMET_DEBUG_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.debug.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 57578235f72..f66dbc88b8a 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 /** @@ -72,7 +73,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, + rssPartitionPusher: Option[CometExecIterator.RssPartitionPusherRegistration] = None) extends Iterator[ColumnarBatch] with Logging { @@ -106,7 +108,7 @@ class CometExecIterator( val memoryConfig = CometExecIterator.getMemoryConfig(conf) - nativeLib.createPlan( + val nativePlan = nativeLib.createPlan( id, inputObjects, protobufQueryPlan, @@ -130,6 +132,25 @@ class CometExecIterator( // worker has neither. See CometUdfBridge.evaluate. TaskContext.get(), Thread.currentThread().getContextClassLoader) + + try { + rssPartitionPusher.foreach { registration => + nativeLib.registerRssPartitionPusher( + nativePlan, + registration.handle, + registration.pusher, + registration.numPartitions, + registration.maxFrameBytes) + } + nativePlan + } catch { + case failure: Throwable => + try nativeLib.releasePlan(nativePlan) + catch { + case releaseFailure: Throwable => failure.addSuppressed(releaseFailure) + } + throw failure + } } private var nextBatch: Option[ColumnarBatch] = None @@ -255,6 +276,18 @@ class CometExecIterator( object CometExecIterator extends Logging { + /** A callback is resolved only inside the native execution context that owns this task. */ + final case class RssPartitionPusherRegistration( + handle: Long, + pusher: ShufflePartitionPusher, + numPartitions: Int, + maxFrameBytes: Int) { + require(handle > 0, "The native RSS partition-pusher handle must be positive") + require(pusher != null, "The native RSS partition pusher must not be null") + require(numPartitions > 0, "The native RSS partition count must be positive") + require(maxFrameBytes >= 20, "The native RSS frame byte limit must fit a complete frame") + } + private def cometSqlConfs: Map[String, String] = SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX)) diff --git a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala index 25cb3383cce..9dc48b40e12 100644 --- a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala +++ b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala @@ -19,7 +19,15 @@ package org.apache.comet.shuffle -import org.apache.spark.{SparkConf, TaskContext} +import java.io.IOException +import java.lang.reflect.InvocationTargetException + +import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext} +import org.apache.spark.shuffle.ShuffleHandle +import org.apache.spark.storage.BlockManagerId + +import org.apache.comet.CometConf +import org.apache.comet.util.ClassLoaders /** Creates task-owned Celeborn pushers using the application's existing Spark configuration. */ object CelebornShufflePusherFactory { @@ -35,6 +43,12 @@ object CelebornShufflePusherFactory { "org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager" private val CELEBORN_SHUFFLE_DATA_IO = "org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO" + private val CELEBORN_SHUFFLE_HANDLE = + "org.apache.spark.shuffle.celeborn.CelebornShuffleHandle" + private val CELEBORN_SPARK_UTILS = "org.apache.spark.shuffle.celeborn.SparkUtils" + private val CELEBORN_SHUFFLE_CLIENT = "org.apache.celeborn.client.ShuffleClient" + private val CELEBORN_CONF = "org.apache.celeborn.common.CelebornConf" + private val CELEBORN_USER_IDENTIFIER = "org.apache.celeborn.common.identity.UserIdentifier" private val MAX_STAGE_ATTEMPTS = 1 << 15 private val MAX_TASK_ATTEMPTS = 1 << 16 @@ -77,7 +91,7 @@ object CelebornShufflePusherFactory { celebornShuffleId: Int, numMappers: Int, numPartitions: Int, - taskContext: TaskContext): Option[ShufflePartitionPusher] = { + taskContext: TaskContext): Option[CelebornShufflePartitionPusher] = { if (!isEnabled(conf)) { None } else { @@ -85,6 +99,12 @@ object CelebornShufflePusherFactory { val stageAttempt = taskContext.stageAttemptNumber() val taskAttempt = taskContext.attemptNumber() val encodedAttempt = encodeAttemptNumber(stageAttempt, taskAttempt) + val maxInFlightBytes = conf.getSizeAsBytes( + CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.key, + CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.defaultValue.get) + require( + maxInFlightBytes >= 76 && maxInFlightBytes <= Int.MaxValue, + "Celeborn executor in-flight byte limit must fit three frame copies and a request header") Some( new CelebornShufflePartitionPusher( client, @@ -92,7 +112,276 @@ object CelebornShufflePusherFactory { mapId, encodedAttempt, numMappers, - numPartitions)) + numPartitions, + maxInFlightBytes.toInt)) + } + } + + /** + * Resolve the application's existing client and stage-attempt shuffle generation from its + * handle. Ownership is recorded immediately after client acquisition so application shutdown + * can release it even when the subsequent shuffle-generation lookup fails. + */ + def createFromHandle( + conf: SparkConf, + handle: ShuffleHandle, + taskContext: TaskContext, + onClientAcquired: AnyRef => Unit, + onShuffleGenerationResolved: (Int, Int) => Unit, + onShuffleGenerationInvalidated: (Int, Int) => Unit = (_, _) => (), + onShuffleGenerationInvalidationUnsafe: (Int, Int) => Boolean = (_, _) => false) + : ResolvedCelebornShufflePusher = { + if (!isEnabled(conf)) { + throw new IllegalStateException("Celeborn shuffle is not enabled for this application") + } + + try { + val handleClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_HANDLE) + if (!handleClass.isInstance(handle)) { + throw new IllegalArgumentException( + "Native Comet shuffle requires an actual Celeborn shuffle handle; " + + s"the delegated manager returned ${handle.getClass.getName}") + } + + val sparkUtilsClass = ClassLoaders.loadClass(CELEBORN_SPARK_UTILS) + val shuffleClientClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_CLIENT) + val celebornConfClass = ClassLoaders.loadClass(CELEBORN_CONF) + val userIdentifierClass = ClassLoaders.loadClass(CELEBORN_USER_IDENTIFIER) + val celebornConf = + sparkUtilsClass.getMethod("fromSparkConf", classOf[SparkConf]).invoke(null, conf) + + def handleValue(name: String): AnyRef = handleClass.getMethod(name).invoke(handle) + + val client = acquireClient( + shuffleClientClass + .getMethod( + "get", + classOf[String], + classOf[String], + java.lang.Integer.TYPE, + celebornConfClass, + userIdentifierClass, + classOf[Array[Byte]]) + .invoke( + null, + handleValue("appUniqueId"), + handleValue("lifecycleManagerHost"), + handleValue("lifecycleManagerPort"), + celebornConf, + handleValue("userIdentifier"), + handleValue("extension")) + .asInstanceOf[AnyRef], + onClientAcquired) + + val stageRerunEnabled = handleValue("stageRerunEnabled").asInstanceOf[Boolean] + if (!stageRerunEnabled) { + throw new IllegalStateException( + "Native Celeborn shuffle requires stage reruns to recover ambiguous map attempts") + } + + // Celeborn installs the barrier-stage failure hook before allocating a generation. + registerBarrierFailureListener( + sparkUtilsClass, + shuffleClientClass, + handleClass, + client, + taskContext, + handle) + + val celebornShuffleId = sparkUtilsClass + .getMethod( + "celebornShuffleId", + shuffleClientClass, + handleClass, + classOf[TaskContext], + classOf[java.lang.Boolean]) + .invoke(null, client, handle, taskContext, java.lang.Boolean.TRUE) + .asInstanceOf[Int] + + val numMappers = handleValue("numMappers").asInstanceOf[Int] + onShuffleGenerationResolved(celebornShuffleId, numMappers) + + if (taskContext.attemptNumber() > 0) { + rejectRetriedAttempt( + client, + handle.shuffleId, + celebornShuffleId, + taskContext, + true, + () => onShuffleGenerationInvalidated(handle.shuffleId, celebornShuffleId), + () => onShuffleGenerationInvalidationUnsafe(handle.shuffleId, celebornShuffleId)) + } + + val numPartitions = handleValue("dependency") + .asInstanceOf[ShuffleDependency[_, _, _]] + .partitioner + .numPartitions + + ResolvedCelebornShufflePusher( + create(conf, client, celebornShuffleId, numMappers, numPartitions, taskContext).get, + client, + celebornShuffleId) + } catch { + case failure: InvocationTargetException => + throw new IllegalStateException( + "Could not resolve the application-owned Celeborn shuffle client", + Option(failure.getCause).getOrElse(failure)) + case failure: ReflectiveOperationException => + throw new IllegalStateException( + "The Celeborn client does not expose the required native-shuffle handle API", + failure) + } + } + + private[shuffle] def acquireClient( + resolveClient: => AnyRef, + onClientAcquired: AnyRef => Unit): AnyRef = { + val client = resolveClient + onClientAcquired(client) + client + } + + private[shuffle] def registerBarrierFailureListener( + sparkUtilsClass: Class[_], + shuffleClientClass: Class[_], + handleClass: Class[_], + client: AnyRef, + taskContext: TaskContext, + handle: ShuffleHandle): Unit = { + sparkUtilsClass + .getMethod( + "addFailureListenerIfBarrierTask", + shuffleClientClass, + classOf[TaskContext], + handleClass) + .invoke(null, client, taskContext, handle) + } + + /** Check task liveness on the driver before mutating Celeborn's shuffle generation. */ + def shouldReportShuffleFetchFailure(taskAttemptId: Long): Boolean = + ClassLoaders + .loadClass(CELEBORN_SPARK_UTILS) + .getMethod("shouldReportShuffleFetchFailure", java.lang.Long.TYPE) + .invoke(null, Long.box(taskAttemptId)) + .asInstanceOf[Boolean] + + /** A retry cannot safely reconstruct the lengths of an already committed Celeborn attempt. */ + private[shuffle] def rejectRetriedAttempt( + client: AnyRef, + sparkShuffleId: Int, + celebornShuffleId: Int, + taskContext: TaskContext, + authorizedToCommit: Boolean, + onGenerationInvalidated: () => Unit = () => (), + onGenerationInvalidationUnsafe: () => Boolean = () => false): Unit = { + if (!authorizedToCommit) { + throw commitDenied(taskContext) + } + if (onGenerationInvalidationUnsafe()) { + throw commitDenied(taskContext) + } + + val invalidated = client.getClass + .getMethod( + "reportShuffleFetchFailure", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Long.TYPE) + .invoke( + client, + Int.box(sparkShuffleId), + Int.box(celebornShuffleId), + Long.box(taskContext.taskAttemptId())) + .asInstanceOf[Boolean] + + if (!invalidated) { + throw new IOException( + "Could not invalidate the Celeborn shuffle generation for a retried map attempt") + } + onGenerationInvalidated() + + // Spark scopes this type private[spark] in Scala even though its JVM constructor is public. + val fetchFailureClass = + ClassLoaders.loadClass("org.apache.spark.shuffle.FetchFailedException") + val fetchFailure = fetchFailureClass + .getConstructor( + classOf[BlockManagerId], + java.lang.Integer.TYPE, + java.lang.Long.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + classOf[String], + classOf[Throwable]) + .newInstance( + null, + Int.box(sparkShuffleId), + Long.box(-1L), + Int.box(-1), + Int.box(-1), + s"Retried Celeborn map attempt requires a new shuffle generation: " + + s"$sparkShuffleId/$celebornShuffleId", + null) + .asInstanceOf[Throwable] + throw fetchFailure + } + + /** + * Preserve Spark's non-counted task-failure classification for a speculative losing attempt. + */ + def commitDenied(taskContext: TaskContext): Throwable = { + val deniedClass = ClassLoaders.loadClass("org.apache.spark.executor.CommitDeniedException") + deniedClass + .getConstructor( + classOf[String], + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE) + .newInstance( + "Another Spark map attempt already owns the Celeborn shuffle commit", + Int.box(taskContext.stageId()), + Int.box(taskContext.partitionId()), + Int.box(taskContext.attemptNumber())) + .asInstanceOf[Throwable] + } + + /** Drop task-independent state when Spark unregisters this particular shuffle generation. */ + def cleanupShuffle(client: AnyRef, celebornShuffleId: Int): Unit = { + try { + client.getClass + .getMethod("cleanupShuffle", java.lang.Integer.TYPE) + .invoke(client, Int.box(celebornShuffleId)) + } catch { + case failure: InvocationTargetException => + throw new IllegalStateException( + "Could not clean up the Celeborn shuffle generation", + Option(failure.getCause).getOrElse(failure)) + case failure: ReflectiveOperationException => + throw new IllegalStateException("Celeborn shuffle cleanup is unavailable", failure) + } + } + + /** Release only the application-scoped client that this manager actually acquired. */ + def releaseClient(client: AnyRef): Unit = { + try { + val shuffleClientClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_CLIENT) + shuffleClientClass.getMethod("removeInstance", shuffleClientClass).invoke(null, client) + } catch { + case failure: InvocationTargetException => + throw new IllegalStateException( + "Could not release the application-owned Celeborn shuffle client", + Option(failure.getCause).getOrElse(failure)) + case failure: ReflectiveOperationException => + throw new IllegalStateException("Celeborn shuffle client cleanup is unavailable", failure) + } finally { + ExecutorShufflePushAdmission.releaseClient(client) } } } + +/** + * A task-owned pusher plus the existing application client and its resolved shuffle generation. + */ +final case class ResolvedCelebornShufflePusher( + pusher: CelebornShufflePartitionPusher, + client: AnyRef, + celebornShuffleId: Int) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index d4d92cbe39f..89784afbea2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -20,19 +20,27 @@ package org.apache.spark.sql.comet.execution.shuffle import java.lang.reflect.InvocationTargetException +import java.util.concurrent.ConcurrentHashMap -import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext} -import org.apache.spark.shuffle.{ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import scala.collection.mutable +import scala.jdk.CollectionConverters._ +import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext, TaskEndReason, UnknownReason} +import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint} +import org.apache.spark.scheduler.OutputCommitCoordinator +import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.spark.util.RpcUtils + +import org.apache.comet.CometConf +import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory} import org.apache.comet.util.ClassLoaders /** * Lets Comet execution coexist with the application's existing Celeborn shuffle manager. * - * Ordinary Spark shuffle dependencies are owned entirely by Celeborn. Native and JVM Comet - * shuffle dependencies remain unsupported until their remote writer, reader, and task lifecycle - * are integrated; rejecting them prevents either local-disk fallback or handing a native Comet - * input iterator to Celeborn's row writer. + * Ordinary Spark shuffle dependencies are owned entirely by Celeborn. Native Comet map tasks can + * write directly to Celeborn, but query planning keeps native shuffle disabled until the matching + * remote reader is available. JVM Comet shuffle remains unsupported. * * Celeborn is loaded reflectively because its client is an optional, application-provided * dependency rather than part of Comet's compile-time or runtime distribution. @@ -50,11 +58,34 @@ class CometCelebornShuffleManager private[shuffle] ( private val celebornManager = Option(backendFactory(conf, isDriver)).getOrElse { throw new IllegalStateException("Celeborn Spark shuffle manager factory returned null") } + private val nativeShuffleClients = + new ConcurrentHashMap[Int, ConcurrentHashMap[Int, AnyRef]]() + private val ownedNativeClients = new ConcurrentHashMap[AnyRef, java.lang.Boolean]() + @volatile private var nativeGenerationCoordinator: CelebornShuffleGenerationCoordinator = _ + @volatile private var nativeGenerationEndpoint: RpcEndpointRef = _ override def registerShuffle[K, V, C]( shuffleId: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle = { dependency match { + case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle => + val handle = celebornManager.registerShuffle(shuffleId, dependency) + if (!CometCelebornShuffleManager.isCelebornHandle(handle)) { + try celebornManager.unregisterShuffle(shuffleId) + catch { + case cleanupFailure: Throwable => + val failure = new UnsupportedOperationException( + "Native Comet shuffle cannot use Celeborn's local fallback writer") + failure.addSuppressed(cleanupFailure) + throw failure + } + throw new UnsupportedOperationException( + "Native Comet shuffle cannot use Celeborn's local fallback writer") + } + if (isDriver) { + initializeNativeGenerationCoordinator() + } + handle case _: CometShuffleDependency[_, _, _] => rejectCometShuffle() case _ => celebornManager.registerShuffle(shuffleId, dependency) } @@ -65,8 +96,76 @@ class CometCelebornShuffleManager private[shuffle] ( mapId: Long, context: TaskContext, metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { - rejectCometHandle(handle) - celebornManager.getWriter(handle, mapId, context, metrics) + nativeDependency(handle) match { + case Some(dependency) => + val earlyClaim = claimNativeShuffleAttempt(handle.shuffleId, context) + if (!earlyClaim.authorized && context.attemptNumber() > 0 && + !earlyClaim.requiresGenerationResolution) { + throw CelebornShufflePusherFactory.commitDenied(context) + } + var preparedClaim = earlyClaim + val resolved = CelebornShufflePusherFactory.createFromHandle( + conf, + handle, + context, + client => ownedNativeClients.put(client, java.lang.Boolean.TRUE), + (celebornShuffleId, numMappers) => + preparedClaim = prepareNativeShuffleGeneration( + handle.shuffleId, + celebornShuffleId, + numMappers, + context, + earlyClaim), + (sparkShuffleId, celebornShuffleId) => + invalidateNativeShuffleGeneration( + sparkShuffleId, + celebornShuffleId, + context, + preparedClaim), + (sparkShuffleId, celebornShuffleId) => + abandonNativeShuffleAttempt( + sparkShuffleId, + celebornShuffleId, + context, + preparedClaim)) + nativeShuffleClients + .computeIfAbsent(handle.shuffleId, _ => new ConcurrentHashMap[Int, AnyRef]()) + .put(resolved.celebornShuffleId, resolved.client) + + new CometNativeShuffleWriter[K, V]( + dependency.nativeShuffleSpec.getOrElse { + throw new IllegalStateException("Native Comet shuffle has no execution plan") + }, + dependency.outputPartitioning.getOrElse { + throw new IllegalStateException("Native Comet shuffle has no output partitioning") + }, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + mapId, + context, + metrics, + dependency.rangePartitionBounds, + Some( + CelebornNativeShuffleDestination( + resolved.pusher, + CometCelebornShuffleManager.maxNativeFrameBytes( + CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.get().toInt, + resolved.pusher), + dependency.partitioner.numPartitions, + commitAuthorized = true, + commitValidator = () => + validateNativeShuffleAttempt( + handle.shuffleId, + resolved.celebornShuffleId, + context, + preparedClaim)))) + + case None => + rejectCometHandle(handle) + celebornManager.getWriter(handle, mapId, context, metrics) + } } override def getReader[K, C]( @@ -77,6 +176,10 @@ class CometCelebornShuffleManager private[shuffle] ( endPartition: Int, context: TaskContext, metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + if (nativeDependency(handle).nonEmpty) { + throw new UnsupportedOperationException( + "Celeborn-backed Comet shuffle cannot be enabled until its native reader is available") + } rejectCometHandle(handle) celebornManager.getReader( handle, @@ -91,10 +194,160 @@ class CometCelebornShuffleManager private[shuffle] ( override def shuffleBlockResolver: ShuffleBlockResolver = celebornManager.shuffleBlockResolver - override def unregisterShuffle(shuffleId: Int): Boolean = + override def unregisterShuffle(shuffleId: Int): Boolean = { + Option(nativeShuffleClients.remove(shuffleId)).foreach { generations => + generations.forEach { (celebornShuffleId, client) => + CelebornShufflePusherFactory.cleanupShuffle(client, celebornShuffleId) + } + } + if (isDriver) { + Option(nativeGenerationCoordinator).foreach(_.unregisterShuffle(shuffleId)) + } celebornManager.unregisterShuffle(shuffleId) + } + + override def stop(): Unit = { + try celebornManager.stop() + finally { + try { + if (isDriver) { + Option(nativeGenerationEndpoint).foreach { endpoint => + SparkEnv.get.rpcEnv.stop(endpoint) + } + } + } finally { + ownedNativeClients.keySet().asScala.foreach(CelebornShufflePusherFactory.releaseClient) + ownedNativeClients.clear() + nativeShuffleClients.clear() + } + } + } + + private def initializeNativeGenerationCoordinator(): Unit = synchronized { + if (nativeGenerationEndpoint == null) { + val env = Option(SparkEnv.get).getOrElse { + throw new IllegalStateException("Spark environment is unavailable for native shuffle") + } + val coordinator = new CelebornShuffleGenerationCoordinator( + env.outputCommitCoordinator, + CelebornShufflePusherFactory.shouldReportShuffleFetchFailure) + val endpoint = env.rpcEnv.setupEndpoint( + CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT, + new CelebornShuffleGenerationEndpoint(env.rpcEnv, coordinator)) + nativeGenerationCoordinator = coordinator + nativeGenerationEndpoint = endpoint + } + } + + private def generationEndpoint: RpcEndpointRef = + Option(nativeGenerationEndpoint).getOrElse { + synchronized { + if (nativeGenerationEndpoint == null) { + if (isDriver) { + initializeNativeGenerationCoordinator() + } else { + nativeGenerationEndpoint = RpcUtils.makeDriverRef( + CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT, + conf, + SparkEnv.get.rpcEnv) + } + } + nativeGenerationEndpoint + } + } + + private def claimNativeShuffleAttempt( + shuffleId: Int, + taskContext: TaskContext): CelebornMapAttemptClaim = + generationEndpoint.askSync[CelebornMapAttemptClaim]( + ClaimCelebornMapAttempt( + shuffleId, + taskContext.stageId(), + taskContext.stageAttemptNumber(), + taskContext.partitionId(), + taskContext.attemptNumber())) + + private def prepareNativeShuffleGeneration( + shuffleId: Int, + celebornShuffleId: Int, + numMappers: Int, + taskContext: TaskContext, + earlyClaim: CelebornMapAttemptClaim): CelebornMapAttemptClaim = { + val prepared = generationEndpoint.askSync[CelebornMapAttemptClaim]( + PrepareCelebornShuffleGeneration( + shuffleId, + celebornShuffleId, + taskContext.stageId(), + taskContext.stageAttemptNumber(), + numMappers, + taskContext.partitionId(), + taskContext.attemptNumber(), + earlyClaim.epoch, + earlyClaim.authorized)) + if (!prepared.authorized) { + throw CelebornShufflePusherFactory.commitDenied(taskContext) + } + prepared + } + + private def validateNativeShuffleAttempt( + shuffleId: Int, + celebornShuffleId: Int, + taskContext: TaskContext, + claim: CelebornMapAttemptClaim): Boolean = + generationEndpoint.askSync[Boolean]( + ValidateCelebornMapAttempt( + shuffleId, + celebornShuffleId, + taskContext.stageId(), + taskContext.stageAttemptNumber(), + taskContext.partitionId(), + taskContext.attemptNumber(), + claim.epoch)) + + private def invalidateNativeShuffleGeneration( + shuffleId: Int, + celebornShuffleId: Int, + taskContext: TaskContext, + claim: CelebornMapAttemptClaim): Unit = { + generationEndpoint.askSync[Boolean]( + InvalidateCelebornShuffleGeneration( + shuffleId, + celebornShuffleId, + taskContext.stageId(), + taskContext.stageAttemptNumber(), + claim.epoch)) + () + } + + private def abandonNativeShuffleAttempt( + shuffleId: Int, + celebornShuffleId: Int, + taskContext: TaskContext, + claim: CelebornMapAttemptClaim): Boolean = + generationEndpoint.askSync[Boolean]( + AbandonCelebornMapAttempt( + shuffleId, + celebornShuffleId, + taskContext.stageId(), + taskContext.stageAttemptNumber(), + taskContext.partitionId(), + taskContext.attemptNumber(), + claim.epoch, + taskContext.taskAttemptId())) - override def stop(): Unit = celebornManager.stop() + private def nativeDependency(handle: ShuffleHandle): Option[CometShuffleDependency[_, _, _]] = { + if (!CometCelebornShuffleManager.isCelebornHandle(handle)) { + None + } else { + handle.asInstanceOf[BaseShuffleHandle[_, _, _]].dependency match { + case dependency: CometShuffleDependency[_, _, _] + if dependency.shuffleType == CometNativeShuffle => + Some(dependency) + case _ => None + } + } + } private def rejectCometHandle(handle: ShuffleHandle): Unit = handle match { case _: CometNativeShuffleHandle[_, _] => rejectCometShuffle() @@ -112,7 +365,18 @@ class CometCelebornShuffleManager private[shuffle] ( private[shuffle] object CometCelebornShuffleManager { + private[shuffle] val GENERATION_COORDINATOR_ENDPOINT = + "CometCelebornShuffleGenerationCoordinator" private val CELEBORN_MANAGER_CLASS = "org.apache.spark.shuffle.celeborn.SparkShuffleManager" + private val CELEBORN_HANDLE_CLASS = "org.apache.spark.shuffle.celeborn.CelebornShuffleHandle" + + private[shuffle] def maxNativeFrameBytes( + configuredMaxFrameBytes: Int, + pusher: CelebornShufflePartitionPusher): Int = + math.min(configuredMaxFrameBytes, pusher.maxFrameBytes()) + + private[shuffle] def isCelebornHandle(handle: ShuffleHandle): Boolean = + handle != null && handle.getClass.getName == CELEBORN_HANDLE_CLASS private[shuffle] def createBackend(conf: SparkConf, isDriver: Boolean): ShuffleManager = { try { @@ -146,3 +410,364 @@ private[shuffle] object CometCelebornShuffleManager { } } } + +private[shuffle] final case class PrepareCelebornShuffleGeneration( + shuffleId: Int, + celebornShuffleId: Int, + stageId: Int, + stageAttempt: Int, + numMappers: Int, + mapId: Int = -1, + taskAttempt: Int = -1, + claimEpoch: Long = -1L, + claimAuthorized: Boolean = false) + extends Serializable + +private[shuffle] final case class ClaimCelebornMapAttempt( + shuffleId: Int, + stageId: Int, + stageAttempt: Int, + mapId: Int, + taskAttempt: Int) + extends Serializable + +private[shuffle] final case class CelebornMapAttemptClaim( + authorized: Boolean, + epoch: Long, + requiresGenerationResolution: Boolean = false) + extends Serializable + +private[shuffle] final case class ValidateCelebornMapAttempt( + shuffleId: Int, + celebornShuffleId: Int, + stageId: Int, + stageAttempt: Int, + mapId: Int, + taskAttempt: Int, + claimEpoch: Long) + extends Serializable + +private[shuffle] final case class InvalidateCelebornShuffleGeneration( + shuffleId: Int, + celebornShuffleId: Int, + stageId: Int, + stageAttempt: Int, + claimEpoch: Long) + extends Serializable + +private[shuffle] final case class AbandonCelebornMapAttempt( + shuffleId: Int, + celebornShuffleId: Int, + stageId: Int, + stageAttempt: Int, + mapId: Int, + taskAttempt: Int, + claimEpoch: Long, + taskAttemptId: Long) + extends Serializable + +/** Keeps Spark's driver-owned commit authorization aligned with Celeborn shuffle generations. */ +private[shuffle] final class CelebornShuffleGenerationCoordinator( + outputCommitCoordinator: OutputCommitCoordinator, + shouldReportShuffleFetchFailure: Long => Boolean = _ => true) { + + private val generations = mutable.HashMap.empty[Int, PrepareCelebornShuffleGeneration] + private val invalidatedGenerations = mutable.HashSet.empty[Int] + private val generationEpochs = mutable.HashMap.empty[Int, Long] + private val claimOwners = + mutable.HashMap.empty[(Int, Int, Int, Int), (Int, Long)] + private val deniedAttempts = + mutable.HashMap.empty[(Int, Int, Int, Int), mutable.HashSet[Int]] + private val authorizeCommit = outputCommitCoordinator.getClass.getMethod( + "handleAskPermissionToCommit", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE) + + private def currentEpoch(shuffleId: Int): Long = generationEpochs.getOrElse(shuffleId, 0L) + + private def ownerKey( + shuffleId: Int, + stageId: Int, + stageAttempt: Int, + mapId: Int): (Int, Int, Int, Int) = + (shuffleId, stageId, stageAttempt, mapId) + + private def authorize( + shuffleId: Int, + stageId: Int, + stageAttempt: Int, + mapId: Int, + taskAttempt: Int): CelebornMapAttemptClaim = { + val epoch = currentEpoch(shuffleId) + val authorized = authorizeCommit + .invoke( + outputCommitCoordinator, + Int.box(stageId), + Int.box(stageAttempt), + Int.box(mapId), + Int.box(taskAttempt)) + .asInstanceOf[Boolean] + if (authorized) { + claimOwners.update(ownerKey(shuffleId, stageId, stageAttempt, mapId), (taskAttempt, epoch)) + } + CelebornMapAttemptClaim(authorized, epoch) + } + + private def invalidateOwners(shuffleId: Int): Unit = { + generationEpochs.update(shuffleId, currentEpoch(shuffleId) + 1L) + claimOwners.filterInPlace { case ((ownerShuffleId, _, _, _), _) => + ownerShuffleId != shuffleId + } + deniedAttempts.filterInPlace { case ((ownerShuffleId, _, _, _), _) => + ownerShuffleId != shuffleId + } + } + + def claimMapAttempt(claim: ClaimCelebornMapAttempt): CelebornMapAttemptClaim = synchronized { + val previousGeneration = generations.get(claim.shuffleId) + val stale = previousGeneration.exists { generation => + generation.stageId == claim.stageId && + (generation.stageAttempt > claim.stageAttempt || + (generation.stageAttempt == claim.stageAttempt && + invalidatedGenerations.contains(claim.shuffleId))) + } + if (stale) { + CelebornMapAttemptClaim(false, currentEpoch(claim.shuffleId)) + } else { + val key = ownerKey(claim.shuffleId, claim.stageId, claim.stageAttempt, claim.mapId) + val epoch = currentEpoch(claim.shuffleId) + if (claimOwners.get(key).contains((claim.taskAttempt, epoch))) { + return CelebornMapAttemptClaim(true, epoch) + } + + // ShuffleMapTask resolves its input iterator before asking for a writer. Reserve an + // unfailed lower-numbered attempt even when a speculative copy reaches the manager first. + // Spark does not remember TaskCommitDenied attempts as failures, so exclude ones that this + // coordinator has already rejected; otherwise they become permanent phantom owners. + val alreadyDenied = deniedAttempts.getOrElse(key, mutable.HashSet.empty[Int]) + var candidate = 0 + while (candidate < claim.taskAttempt) { + if (!alreadyDenied.contains(candidate) && + !claimOwners.get(key).contains((candidate, epoch))) { + val earlier = + authorize(claim.shuffleId, claim.stageId, claim.stageAttempt, claim.mapId, candidate) + if (earlier.authorized) { + deniedAttempts.getOrElseUpdate(key, mutable.HashSet.empty[Int]).add(claim.taskAttempt) + return CelebornMapAttemptClaim(false, epoch) + } + } + candidate += 1 + } + + val result = authorize( + claim.shuffleId, + claim.stageId, + claim.stageAttempt, + claim.mapId, + claim.taskAttempt) + if (!result.authorized && claim.taskAttempt > 0) { + val requiresResolution = previousGeneration.exists { generation => + generation.stageId == claim.stageId && generation.stageAttempt < claim.stageAttempt + } && !claimOwners.get(key).exists { case (_, ownerEpoch) => ownerEpoch == epoch } + if (!requiresResolution) { + deniedAttempts.getOrElseUpdate(key, mutable.HashSet.empty[Int]).add(claim.taskAttempt) + } + result.copy(requiresGenerationResolution = requiresResolution) + } else { + result + } + } + } + + def prepareGeneration(generation: PrepareCelebornShuffleGeneration): Boolean = synchronized { + require(generation.numMappers > 0, "Celeborn shuffle mapper count must be positive") + + generations.get(generation.shuffleId) match { + case Some(previous) + if invalidatedGenerations.contains(generation.shuffleId) && + previous.celebornShuffleId == generation.celebornShuffleId => + false + case Some(previous) + if previous.stageId == generation.stageId && + previous.stageAttempt > generation.stageAttempt => + false + case Some(previous) + if previous.stageId == generation.stageId && + previous.stageAttempt == generation.stageAttempt => + previous.celebornShuffleId == generation.celebornShuffleId + case Some(previous) if previous.celebornShuffleId != generation.celebornShuffleId => + // Spark scopes these lifecycle methods private[scheduler] despite public JVM methods. + outputCommitCoordinator.synchronized { + outputCommitCoordinator.getClass + .getMethod("stageEnd", java.lang.Integer.TYPE) + .invoke(outputCommitCoordinator, Int.box(generation.stageId)) + outputCommitCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke( + outputCommitCoordinator, + Int.box(generation.stageId), + Int.box(generation.numMappers - 1)) + } + invalidateOwners(generation.shuffleId) + invalidatedGenerations.remove(generation.shuffleId) + generations.update(generation.shuffleId, generation) + true + case Some(_) => + generations.update(generation.shuffleId, generation) + true + case None => + generations.update(generation.shuffleId, generation) + true + } + } + + def prepareGenerationAndClaim( + generation: PrepareCelebornShuffleGeneration): CelebornMapAttemptClaim = synchronized { + if (!prepareGeneration(generation)) { + deniedAttempts + .getOrElseUpdate( + ownerKey( + generation.shuffleId, + generation.stageId, + generation.stageAttempt, + generation.mapId), + mutable.HashSet.empty[Int]) + .add(generation.taskAttempt) + return CelebornMapAttemptClaim(false, currentEpoch(generation.shuffleId)) + } + + val epoch = currentEpoch(generation.shuffleId) + val expectedOwner = claimOwners.get( + ownerKey( + generation.shuffleId, + generation.stageId, + generation.stageAttempt, + generation.mapId)) + if (generation.claimAuthorized && generation.claimEpoch == epoch && + expectedOwner.contains((generation.taskAttempt, epoch))) { + CelebornMapAttemptClaim(true, epoch) + } else { + val claim = authorize( + generation.shuffleId, + generation.stageId, + generation.stageAttempt, + generation.mapId, + generation.taskAttempt) + if (!claim.authorized) { + deniedAttempts + .getOrElseUpdate( + ownerKey( + generation.shuffleId, + generation.stageId, + generation.stageAttempt, + generation.mapId), + mutable.HashSet.empty[Int]) + .add(generation.taskAttempt) + } + claim + } + } + + def validateMapAttempt(validation: ValidateCelebornMapAttempt): Boolean = synchronized { + generations.get(validation.shuffleId).exists { generation => + !invalidatedGenerations.contains(validation.shuffleId) && + generation.celebornShuffleId == validation.celebornShuffleId && + generation.stageId == validation.stageId && + generation.stageAttempt == validation.stageAttempt && + currentEpoch(validation.shuffleId) == validation.claimEpoch && + claimOwners + .get( + ownerKey( + validation.shuffleId, + validation.stageId, + validation.stageAttempt, + validation.mapId)) + .contains((validation.taskAttempt, validation.claimEpoch)) + } + } + + def invalidateGeneration(invalidation: InvalidateCelebornShuffleGeneration): Boolean = + synchronized { + val current = generations.get(invalidation.shuffleId).exists { generation => + generation.celebornShuffleId == invalidation.celebornShuffleId && + generation.stageId == invalidation.stageId && + generation.stageAttempt == invalidation.stageAttempt && + currentEpoch(invalidation.shuffleId) == invalidation.claimEpoch + } + if (current) { + invalidatedGenerations.add(invalidation.shuffleId) + invalidateOwners(invalidation.shuffleId) + } + current + } + + def abandonMapAttempt(abandoned: AbandonCelebornMapAttempt): Boolean = synchronized { + val current = validateMapAttempt( + ValidateCelebornMapAttempt( + abandoned.shuffleId, + abandoned.celebornShuffleId, + abandoned.stageId, + abandoned.stageAttempt, + abandoned.mapId, + abandoned.taskAttempt, + abandoned.claimEpoch)) + val abandon = current && !shouldReportShuffleFetchFailure(abandoned.taskAttemptId) + if (abandon) { + outputCommitCoordinator.getClass + .getMethod( + "taskCompleted", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + classOf[TaskEndReason]) + .invoke( + outputCommitCoordinator, + Int.box(abandoned.stageId), + Int.box(abandoned.stageAttempt), + Int.box(abandoned.mapId), + Int.box(abandoned.taskAttempt), + UnknownReason) + val key = + ownerKey(abandoned.shuffleId, abandoned.stageId, abandoned.stageAttempt, abandoned.mapId) + claimOwners.remove(key) + deniedAttempts.getOrElseUpdate(key, mutable.HashSet.empty[Int]).add(abandoned.taskAttempt) + } + abandon + } + + def unregisterShuffle(shuffleId: Int): Unit = synchronized { + generations.remove(shuffleId) + invalidatedGenerations.remove(shuffleId) + generationEpochs.remove(shuffleId) + claimOwners.filterInPlace { case ((ownerShuffleId, _, _, _), _) => + ownerShuffleId != shuffleId + } + deniedAttempts.filterInPlace { case ((ownerShuffleId, _, _, _), _) => + ownerShuffleId != shuffleId + } + } +} + +private[shuffle] final class CelebornShuffleGenerationEndpoint( + override val rpcEnv: RpcEnv, + coordinator: CelebornShuffleGenerationCoordinator) + extends ThreadSafeRpcEndpoint { + + override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { + case claim: ClaimCelebornMapAttempt => + context.reply(coordinator.claimMapAttempt(claim)) + case generation: PrepareCelebornShuffleGeneration if generation.mapId < 0 => + context.reply(coordinator.prepareGeneration(generation)) + case generation: PrepareCelebornShuffleGeneration => + context.reply(coordinator.prepareGenerationAndClaim(generation)) + case validation: ValidateCelebornMapAttempt => + context.reply(coordinator.validateMapAttempt(validation)) + case invalidation: InvalidateCelebornShuffleGeneration => + context.reply(coordinator.invalidateGeneration(invalidation)) + case abandoned: AbandonCelebornMapAttempt => + context.reply(coordinator.abandonMapAttempt(abandoned)) + } +} 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 7c64e963aa9..d411b56fb68 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 @@ -21,6 +21,8 @@ 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 import scala.collection.mutable import scala.jdk.CollectionConverters._ @@ -35,11 +37,13 @@ import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partition import org.apache.spark.sql.comet.{CometExec, CometMetricNode, CometScalarSubquery, PlanDataInjector} import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.StructField +import org.apache.spark.util.ThreadUtils import org.apache.comet.{CometConf, CometExecIterator} import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde} import org.apache.comet.serde.OperatorOuterClass.{CompressionCodec, Operator} import org.apache.comet.serde.operator.schema2Proto +import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory} /** * Drives the native shuffle write in a single [[CometExecIterator]] per partition. The plan is @@ -60,7 +64,8 @@ class CometNativeShuffleWriter[K, V]( mapId: Long, context: TaskContext, metricsReporter: ShuffleWriteMetricsReporter, - rangePartitionBounds: Option[Seq[InternalRow]] = None) + rangePartitionBounds: Option[Seq[InternalRow]] = None, + remoteDestination: Option[CelebornNativeShuffleDestination] = None) extends ShuffleWriter[K, V] with Logging { @@ -68,16 +73,56 @@ class CometNativeShuffleWriter[K, V]( var partitionLengths: Array[Long] = _ var mapStatus: MapStatus = _ + private var stopped = false + private lazy val effectivePartitionCount = + remoteDestination.map(_.numPartitions).getOrElse(outputPartitioning.numPartitions) + + private val cancellationWatch: Option[ScheduledFuture[_]] = remoteDestination.flatMap { + destination => + Option(context).map { taskContext => + val watcher = CelebornNativeShuffleDestination.watchForCancellation( + taskContext, + destination.pusher, + failure => + logWarning("Could not abort an interrupted Celeborn shuffle map attempt", failure)) + + taskContext.addTaskCompletionListener[Unit] { _ => + watcher.cancel(false) + destination.pusher.abort() + } + watcher + } + } override def write(inputs: Iterator[Product2[K, V]]): Unit = { - val shuffleBlockResolver = - SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver] - val dataFile = shuffleBlockResolver.getDataFile(shuffleId, mapId) - val indexFile = shuffleBlockResolver.getIndexFile(shuffleId, mapId) - val tempDataFilename = dataFile.getPath.replace(".data", ".data.tmp") - val tempIndexFilename = indexFile.getPath.replace(".index", ".index.tmp") - val tempDataFilePath = Paths.get(tempDataFilename) - val tempIndexFilePath = Paths.get(tempIndexFilename) + try { + writeInternal(inputs) + } catch { + case failure: Throwable => + remoteDestination.foreach { destination => + try destination.pusher.abort() + catch { + case cleanupFailure: Throwable => failure.addSuppressed(cleanupFailure) + } + } + throw failure + } + } + + private def writeInternal(inputs: Iterator[Product2[K, V]]): Unit = { + val localOutput = if (remoteDestination.isEmpty) { + 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"))) + } else { + None + } // The dep's _rdd is always a CometNativeShuffleInputRDD on this path. Pattern-match instead // of asInstanceOf so a future RDD-layering change produces a clear error here rather than a @@ -94,7 +139,10 @@ class CometNativeShuffleWriter[K, V]( val inputObjects = shuffleInputIter.inputObjects val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators - val unifiedPlan = buildUnifiedPlan(tempDataFilename, tempIndexFilename) + val unifiedPlan = localOutput match { + case Some(output) => buildUnifiedPlan(output.dataFile, output.indexFile) + case None => buildUnifiedPlan("", "") + } val ctx = spec.execContext val finalNativePlan = if (ctx.commonByKey.nonEmpty) { // This partition's plan-data slice rides on the input iterator's Partition object (populated @@ -137,6 +185,14 @@ class CometNativeShuffleWriter[K, V]( Option(context).foreach(nativeMetrics.reportScanInputMetrics) } + val rssRegistration = remoteDestination.map { destination => + CometExecIterator.RssPartitionPusherRegistration( + CelebornNativeShuffleDestination.TASK_PUSHER_HANDLE, + destination.pusher, + destination.numPartitions, + destination.maxFrameBytes) + } + val cometIter = new CometExecIterator( CometExec.newIterId, inputObjects, @@ -147,7 +203,8 @@ class CometNativeShuffleWriter[K, V]( partitionIdx, ctx.broadcastedHadoopConfForEncryption, ctx.encryptedFilePaths, - shuffleBlockIters) + shuffleBlockIters, + rssPartitionPusher = rssRegistration) // Register subqueries against the iterator id so native callbacks resolve them to values. ctx.subqueries.foreach { sub => @@ -161,51 +218,89 @@ class CometNativeShuffleWriter[K, V]( } } - while (cometIter.hasNext) { - cometIter.next() + try { + while (cometIter.hasNext) { + cometIter.next() + } + } finally { + cometIter.close() + } + + remoteDestination match { + case Some(destination) => + val completionStart = System.nanoTime() + // Manager-created destinations already acquired ownership before Celeborn lookup. + // Spark's coordinator rejects even the same owner on a second request, so validate its + // generation-aware driver token instead of asking the coordinator again. + val authorized = if (destination.commitAuthorized) { + destination.commitValidator() + } else { + SparkEnv.get.outputCommitCoordinator.canCommit( + context.stageId(), + context.stageAttemptNumber(), + context.partitionId(), + context.attemptNumber()) + } + if (!authorized) { + throw CelebornShufflePusherFactory.commitDenied(context) + } + partitionLengths = destination.pusher.finish() + if (destination.commitAuthorized && !destination.commitValidator()) { + throw CelebornShufflePusherFactory.commitDenied(context) + } + metricsReporter.incBytesWritten(partitionLengths.sum) + metricsReporter.incWriteTime(System.nanoTime() - completionStart) + mapStatus = MapStatus.apply( + SparkEnv.get.blockManager.shuffleServerId, + partitionLengths, + context.taskAttemptId()) + + case None => + val output = localOutput.get + val tempDataFilePath = Paths.get(output.dataFile) + val tempIndexFilePath = Paths.get(output.indexFile) + + // get partition lengths from shuffle write output index file + var offset = 0L + partitionLengths = Files + .readAllBytes(tempIndexFilePath) + .grouped(OFFSET_LENGTH) + .drop(1) // first partition offset is always 0 + .map(indexBytes => { + val partitionOffset = + ByteBuffer.wrap(indexBytes).order(ByteOrder.LITTLE_ENDIAN).getLong + val partitionLength = partitionOffset - offset + offset = partitionOffset + partitionLength + }) + .toArray + Files.delete(tempIndexFilePath) + + metricsReporter.incBytesWritten(Files.size(tempDataFilePath)) + + // commit + output.resolver.writeMetadataFileAndCommit( + shuffleId, + mapId, + partitionLengths, + Array.empty, // TODO: add checksums + tempDataFilePath.toFile) + mapStatus = + MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId) } - cometIter.close() - - // get partition lengths from shuffle write output index file - var offset = 0L - partitionLengths = Files - .readAllBytes(tempIndexFilePath) - .grouped(OFFSET_LENGTH) - .drop(1) // first partition offset is always 0 - .map(indexBytes => { - val partitionOffset = - ByteBuffer.wrap(indexBytes).order(ByteOrder.LITTLE_ENDIAN).getLong - val partitionLength = partitionOffset - offset - offset = partitionOffset - partitionLength - }) - .toArray - Files.delete(tempIndexFilePath) - - // Total written bytes at native - metricsReporter.incBytesWritten(Files.size(tempDataFilePath)) + metricsReporter.incRecordsWritten(metricsOutputRows.value) metricsReporter.incWriteTime(metricsWriteTime.value) - - // commit - shuffleBlockResolver.writeMetadataFileAndCommit( - shuffleId, - mapId, - partitionLengths, - Array.empty, // TODO: add checksums - tempDataFilePath.toFile) - mapStatus = - MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId) } private def isSinglePartitioning(p: Partitioning): Boolean = p match { case SinglePartition => true - case rp: RangePartitioning => + case _: RangePartitioning => // Spark sometimes generates RangePartitioning schemes with numPartitions == 1, // or the computed bounds results in a single target partition. // In this case Comet just serializes a SinglePartition scheme to native. - rp.numPartitions == 1 || rangePartitionBounds.forall(_.isEmpty) - case hp: HashPartitioning => hp.numPartitions == 1 + effectivePartitionCount == 1 || rangePartitionBounds.forall(_.isEmpty) + case _: HashPartitioning => effectivePartitionCount == 1 case _ => false } @@ -213,10 +308,21 @@ class CometNativeShuffleWriter[K, V]( * Build the unified `ShuffleWriter(child = childNativeOp)` plan with the partitioning serde, * compression settings, and output file paths. */ - private def buildUnifiedPlan(dataFile: String, indexFile: String): Operator = { + private[shuffle] def buildUnifiedPlan(dataFile: String, indexFile: String): Operator = { val shuffleWriterBuilder = OperatorOuterClass.ShuffleWriter.newBuilder() - shuffleWriterBuilder.setOutputDataFile(dataFile) - shuffleWriterBuilder.setOutputIndexFile(indexFile) + remoteDestination match { + case Some(_) => + val remoteWriter = PartitioningOuterClass.RssPartitionWriter + .newBuilder() + .setRssPartitionPusher(CelebornNativeShuffleDestination.TASK_PUSHER_HANDLE) + shuffleWriterBuilder.setPartitionWriter( + PartitioningOuterClass.PartitionWriter + .newBuilder() + .setRssPartitionWriter(remoteWriter)) + case None => + shuffleWriterBuilder.setOutputDataFile(dataFile) + shuffleWriterBuilder.setOutputIndexFile(indexFile) + } if (SparkEnv.get.conf.getBoolean("spark.shuffle.compress", true)) { val codec = CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.get() match { @@ -243,7 +349,7 @@ class CometNativeShuffleWriter[K, V]( case _: HashPartitioning => val hashPartitioning = outputPartitioning.asInstanceOf[HashPartitioning] val partitioning = PartitioningOuterClass.HashPartition.newBuilder() - partitioning.setNumPartitions(outputPartitioning.numPartitions) + partitioning.setNumPartitions(effectivePartitionCount) val partitionExprs = hashPartitioning.expressions .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) @@ -261,7 +367,7 @@ class CometNativeShuffleWriter[K, V]( case _: RangePartitioning => val rangePartitioning = outputPartitioning.asInstanceOf[RangePartitioning] val partitioning = PartitioningOuterClass.RangePartition.newBuilder() - partitioning.setNumPartitions(outputPartitioning.numPartitions) + partitioning.setNumPartitions(effectivePartitionCount) // Detect duplicates by tracking expressions directly, similar to DataFusion's LexOrdering // DataFusion will deduplicate identical sort expressions in LexOrdering, @@ -325,7 +431,7 @@ class CometNativeShuffleWriter[K, V]( case _: RoundRobinPartitioning => val partitioning = PartitioningOuterClass.RoundRobinPartition.newBuilder() - partitioning.setNumPartitions(outputPartitioning.numPartitions) + partitioning.setNumPartitions(effectivePartitionCount) partitioning.setMaxHashColumns( CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get()) @@ -354,12 +460,83 @@ class CometNativeShuffleWriter[K, V]( } override def stop(success: Boolean): Option[MapStatus] = { - if (success) { - Some(mapStatus) - } else { - None + remoteDestination match { + case None => + if (success) Some(mapStatus) else None + + case Some(destination) => + synchronized { + if (stopped) { + None + } else { + stopped = true + try { + if (success) { + if (mapStatus == null) { + throw new IllegalStateException( + "Cannot complete a Celeborn shuffle map task before writing its data") + } + Some(mapStatus) + } else { + None + } + } finally { + cancellationWatch.foreach(_.cancel(false)) + destination.pusher.abort() + } + } + } } } override def getPartitionLengths(): Array[Long] = partitionLengths + + private final case class LocalShuffleOutput( + resolver: IndexShuffleBlockResolver, + dataFile: String, + indexFile: String) +} + +/** A task-scoped remote destination; its callback is never shared between Spark map attempts. */ +private[shuffle] final case class CelebornNativeShuffleDestination( + pusher: CelebornShufflePartitionPusher, + maxFrameBytes: Int, + numPartitions: Int, + commitAuthorized: Boolean = false, + commitValidator: () => Boolean = () => true) { + require(pusher != null, "The Celeborn shuffle partition pusher must not be null") + require(maxFrameBytes >= 20, "The Celeborn shuffle frame byte limit must fit a complete frame") + require( + numPartitions == pusher.numPartitions(), + "The Celeborn shuffle destination must use its actual reducer partition count") +} + +private[shuffle] object CelebornNativeShuffleDestination { + // Handles are resolved only within their owning native execution context, so every map attempt + // can use the same nonzero identifier without an executor-wide callback registry. + val TASK_PUSHER_HANDLE: Long = 1L + private val cancellationWatcher = + ThreadUtils.newDaemonSingleThreadScheduledExecutor("comet-celeborn-task-cancellation") + + private[shuffle] def watchForCancellation( + taskContext: TaskContext, + pusher: CelebornShufflePartitionPusher, + reportFailure: Throwable => Unit): ScheduledFuture[_] = { + val handled = new AtomicBoolean(false) + cancellationWatcher.scheduleWithFixedDelay( + new Runnable { + override def run(): Unit = { + if ((taskContext.isInterrupted() || taskContext.isFailed()) && + handled.compareAndSet(false, true)) { + try pusher.abort() + catch { + case failure: Throwable => reportFailure(failure) + } + } + } + }, + 0L, + 25L, + TimeUnit.MILLISECONDS) + } } diff --git a/spark/src/test/java/org/apache/comet/shuffle/RecordingCelebornSparkUtils.java b/spark/src/test/java/org/apache/comet/shuffle/RecordingCelebornSparkUtils.java new file mode 100644 index 00000000000..44508900ced --- /dev/null +++ b/spark/src/test/java/org/apache/comet/shuffle/RecordingCelebornSparkUtils.java @@ -0,0 +1,39 @@ +/* + * 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 org.apache.spark.TaskContext; +import org.apache.spark.shuffle.ShuffleHandle; + +/** Verifies that Comet preserves Celeborn's static barrier-listener hook and argument order. */ +public final class RecordingCelebornSparkUtils { + public static Object client; + public static TaskContext taskContext; + public static ShuffleHandle shuffleHandle; + + private RecordingCelebornSparkUtils() {} + + public static void addFailureListenerIfBarrierTask( + Object currentClient, TaskContext currentTask, ShuffleHandle currentHandle) { + client = currentClient; + taskContext = currentTask; + shuffleHandle = currentHandle; + } +} diff --git a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala index abc89b1e2c8..ac306ee9f8a 100644 --- a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala +++ b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala @@ -20,19 +20,101 @@ package org.apache.comet.shuffle import java.io.IOException -import java.util.concurrent.atomic.AtomicReference +import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedDeque, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference, LongAdder} import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.{SparkConf, TaskContext} +import org.apache.spark.shuffle.ShuffleHandle + +import org.apache.comet.CometExecIterator + +/** Matches the pinned Celeborn completion hook without adding its optional client dependency. */ +trait RecordingCelebornPushMetricsCallback { + def incPushDataCount(count: Long): Unit + def incPushDataRetryCount(count: Long): Unit + def incPushDataTime(time: Long): Unit + def incInFlightWaitTime(time: Long): Unit +} /** Public so the production adapter can invoke its method through ordinary Java reflection. */ final class RecordingCelebornShuffleClient { @volatile var acceptedBytes: Option[Int] = None @volatile var failure: Throwable = _ + @volatile var mapperEndFailure: Throwable = _ + @volatile var cleanupFailure: Throwable = _ @volatile var observedTaskContext: TaskContext = _ @volatile var lastPush: RecordedCelebornPush = _ + @volatile var lastMapperEnd: (Int, Int, Int, Int) = _ + @volatile var lastCleanup: (Int, Int, Int) = _ + @volatile var pushStarted: CountDownLatch = _ + @volatile var allowPush: CountDownLatch = _ + @volatile var pushCompletedBeforeReturn: CountDownLatch = _ + @volatile var allowPushReturn: CountDownLatch = _ + @volatile var mapperEndStarted: CountDownLatch = _ + @volatile var allowMapperEnd: CountDownLatch = _ + @volatile var drainStarted: CountDownLatch = _ + @volatile var allowDrain: CountDownLatch = _ + @volatile var drainFailure: Throwable = _ + @volatile var drainTimedOut = false + @volatile var automaticallyCompletePushes = true + @volatile var automaticallyCompleteLatestPush = false + @volatile var silentlyCompletePushes = false + @volatile var metricsFailureWithoutBatchRemoval = false + @volatile var completionRemoved: CountDownLatch = _ + @volatile var allowCompletionMetrics: CountDownLatch = _ + @volatile var cleanupUnblocksPush = false + @volatile var recreatePushStateAfterCleanup = false + @volatile var cleanupUnblocksMapperEnd = false + @volatile var cleanupUnblocksDrain = false + @volatile var retainedPushState = false + @volatile var observedMapKey: String = _ + @volatile var generationInvalidated = true + @volatile var lastInvalidatedShuffle: (Int, Int, Long) = _ + val pushCalls = new AtomicInteger() + val pushCompletionCalls = new AtomicInteger() + val mapperEndCalls = new AtomicInteger() + val cleanupCalls = new AtomicInteger() + private val pushStates = new ConcurrentHashMap[String, RecordingCelebornPushState]() + private val pendingPushCompletions = new ConcurrentLinkedDeque[RecordingCelebornPushState]() + + def getPushState(mapKey: String): RecordingCelebornPushState = { + observedMapKey = mapKey + retainedPushState = true + pushStates.computeIfAbsent(mapKey, _ => new RecordingCelebornPushState(this)) + } + + def completeNextPush(): Boolean = { + completePush(pendingPushCompletions.pollFirst()) + } + + def completeLatestPush(): Boolean = { + completePush(pendingPushCompletions.pollLast()) + } + + private def completePush(pushState: RecordingCelebornPushState): Boolean = { + if (pushState == null) { + false + } else { + pushCompletionCalls.incrementAndGet() + if (metricsFailureWithoutBatchRemoval) { + pushState.failPushWithoutRemovingBatch() + } else { + pushState.completePush() + } + true + } + } + + def reportShuffleFetchFailure( + sparkShuffleId: Int, + celebornShuffleId: Int, + taskAttemptId: Long): Boolean = { + lastInvalidatedShuffle = (sparkShuffleId, celebornShuffleId, taskAttemptId) + generationInvalidated + } @throws[IOException] def pushOrMergeData( @@ -47,6 +129,7 @@ final class RecordingCelebornShuffleClient { numPartitions: Int, doPush: Boolean, skipCompress: Boolean): Int = { + pushCalls.incrementAndGet() observedTaskContext = TaskContext.get() lastPush = RecordedCelebornPush( shuffleId, @@ -61,12 +144,141 @@ final class RecordingCelebornShuffleClient { doPush, skipCompress) + if (pushStarted != null) { + pushStarted.countDown() + if (!allowPush.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting for the test push") + } + } + if (failure != null) { throw failure } + if (recreatePushStateAfterCleanup && cleanupCalls.get() > 0) { + retainedPushState = true + } + + val pushState = getPushState(s"$shuffleId-$mapId-$attemptId") + pushState.addPush() + pendingPushCompletions.add(pushState) + if (automaticallyCompleteLatestPush) { + completeLatestPush() + } else if (automaticallyCompletePushes) { + completeNextPush() + } + + if (pushCompletedBeforeReturn != null) { + pushCompletedBeforeReturn.countDown() + if (!allowPushReturn.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting for the completed test push to return") + } + } + acceptedBytes.getOrElse(length + 16) } + + @throws[IOException] + def mapperEnd(shuffleId: Int, mapId: Int, attemptId: Int, numMappers: Int): Unit = { + lastMapperEnd = (shuffleId, mapId, attemptId, numMappers) + mapperEndCalls.incrementAndGet() + if (getPushState(s"$shuffleId-$mapId-$attemptId").limitZeroInFlight()) { + throw new IOException("Timed out waiting for the accepted Celeborn shuffle push") + } + if (mapperEndStarted != null) { + mapperEndStarted.countDown() + if (!allowMapperEnd.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting for test map completion") + } + } + if (mapperEndFailure != null) { + throw mapperEndFailure + } + } + + def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = { + lastCleanup = (shuffleId, mapId, attemptId) + cleanupCalls.incrementAndGet() + retainedPushState = false + pushStates.remove(s"$shuffleId-$mapId-$attemptId") + if (cleanupUnblocksPush) { + failure = new IOException("Celeborn map attempt was cleaned up") + allowPush.countDown() + } + if (recreatePushStateAfterCleanup) { + allowPush.countDown() + } + if (cleanupUnblocksMapperEnd) { + mapperEndFailure = new IOException("Celeborn map completion was cleaned up") + allowMapperEnd.countDown() + } + if (cleanupUnblocksDrain) { + drainFailure = new IOException("Celeborn pending push was cleaned up") + allowDrain.countDown() + } + if (cleanupFailure != null) { + throw cleanupFailure + } + } +} + +/** Public so production reflection can wait for the fake client's accepted network request. */ +final class RecordingCelebornInFlightRequestTracker { + private val totalInflightReqs = new LongAdder() + + def addBatch(): Unit = totalInflightReqs.increment() + + def removeBatch(): Unit = totalInflightReqs.decrement() +} + +final class RecordingCelebornPushState(client: RecordingCelebornShuffleClient) { + private val inFlightRequestTracker = new RecordingCelebornInFlightRequestTracker() + private val exception = new AtomicReference[IOException]() + @volatile private var metricsCallback: RecordingCelebornPushMetricsCallback = _ + + def setMetricsCallback(callback: RecordingCelebornPushMetricsCallback): Unit = { + metricsCallback = callback + } + + def addPush(): Unit = inFlightRequestTracker.addBatch() + + def completePush(): Unit = { + inFlightRequestTracker.removeBatch() + if (client.completionRemoved != null) { + client.completionRemoved.countDown() + if (!client.allowCompletionMetrics.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting to report the completed test push") + } + } + if (!client.silentlyCompletePushes) { + onSuccess() + } + } + + def failPushWithoutRemovingBatch(): Unit = onFailure() + + private def onSuccess(): Unit = { + Option(metricsCallback).foreach(_.incPushDataCount(1)) + } + + private def onFailure(): Unit = { + Option(metricsCallback).foreach(_.incPushDataCount(1)) + exception.compareAndSet(null, new IOException("the test Celeborn push failed")) + } + + @throws[IOException] + def limitZeroInFlight(): Boolean = { + if (client.drainStarted != null) { + client.drainStarted.countDown() + if (!client.allowDrain.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting for the test request to complete") + } + } + if (client.drainFailure != null) { + throw client.drainFailure + } + client.drainTimedOut + } } final case class RecordedCelebornPush( @@ -137,6 +349,7 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { assert(push.numPartitions == 9) assert(push.doPush) assert(push.skipCompress) + assert(client.observedMapKey == s"19-3-${(4 << 16) | 7}") } test("raw Celeborn push requires exactly the payload plus its transport header") { @@ -165,6 +378,8 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { } assert(actual eq expected) + assert(client.cleanupCalls.get() == 1) + assert(client.lastCleanup == ((19, 3, (4 << 16) | 7))) } test("raw Celeborn push preserves unchecked client failures") { @@ -177,6 +392,893 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { } assert(actual eq expected) + assert(client.cleanupCalls.get() == 1) + } + + test("executor-wide byte admission stays reserved until the accepted push completes") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val first = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val second = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val secondFailure = new AtomicReference[Throwable]() + val bytes = Array.fill[Byte](32)(1) + + assert(first.pushPartitionData(0, bytes, bytes.length) == bytes.length) + assert(client.pushCompletionCalls.get() == 0) + + val secondThread = new Thread(() => { + try second.pushPartitionData(0, bytes, bytes.length) + catch { case failure: Throwable => secondFailure.set(failure) } + }) + secondThread.start() + + Thread.sleep(100) + assert(client.pushCalls.get() == 1) + + assert(client.completeNextPush()) + secondThread.join(5000) + + assert(!secondThread.isAlive) + assert(secondFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + + test("accepted frames pipeline without synchronously draining Celeborn requests") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val adapter = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 96) + + (0 until 3).foreach { partition => + assert(adapter.pushPartitionData(partition, Array.fill[Byte](8)(1), 8) == 8) + } + + assert(client.pushCalls.get() == 3) + assert(client.pushCompletionCalls.get() == 0) + assert(client.completeNextPush()) + assert(client.completeNextPush()) + assert(client.completeNextPush()) + } + + test("native encoding reservations shrink to actual frame bytes before transport completes") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val adapter = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 80) + val failure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](8)(1) + + adapter.reservePartitionData(32) + assert(adapter.pushPartitionData(0, frame, frame.length) == frame.length) + + val worker = new Thread(() => { + try { + adapter.reservePartitionData(32) + adapter.pushPartitionData(1, frame, frame.length) + } catch { case error: Throwable => failure.set(error) } + }) + worker.start() + worker.join(5000) + + if (worker.isAlive) { + assert(client.completeNextPush()) + worker.join(5000) + fail("the second small native frame waited for the first transport to complete") + } + + assert(failure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.pushCompletionCalls.get() == 0) + assert(client.completeNextPush()) + assert(client.completeNextPush()) + } + + test("native push admission covers every overlapping copy until the raw push returns") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + val original = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 112) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 112) + val originalFailure = new AtomicReference[Throwable]() + val replacementFailure = new AtomicReference[Throwable]() + val originalFrame = Array.fill[Byte](32)(1) + val replacementFrame = Array.fill[Byte](8)(1) + + val originalWorker = new Thread(() => { + try { + original.reservePartitionData(96) + original.pushPartitionData(0, originalFrame, originalFrame.length) + } catch { case failure: Throwable => originalFailure.set(failure) } + }) + originalWorker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + val replacementWorker = new Thread(() => { + try { + replacement.reservePartitionData(24) + replacement.pushPartitionData(0, replacementFrame, replacementFrame.length) + } catch { case failure: Throwable => replacementFailure.set(failure) } + }) + replacementWorker.start() + + Thread.sleep(100) + assert(replacementWorker.isAlive) + assert(client.pushCalls.get() == 1) + + client.allowPush.countDown() + originalWorker.join(5000) + replacementWorker.join(5000) + + if (replacementWorker.isAlive) { + assert(client.completeNextPush()) + replacementWorker.join(5000) + fail("native-copy admission was retained after the raw push returned") + } + + assert(!originalWorker.isAlive) + assert(originalFailure.get() == null) + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.pushCompletionCalls.get() == 0) + assert(client.completeNextPush()) + assert(client.completeNextPush()) + } + + test("inline transport completion cannot release native copies before the raw push returns") { + val client = new RecordingCelebornShuffleClient + client.pushCompletedBeforeReturn = new CountDownLatch(1) + client.allowPushReturn = new CountDownLatch(1) + val original = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 112) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 112) + val originalFailure = new AtomicReference[Throwable]() + val replacementFailure = new AtomicReference[Throwable]() + val originalFrame = Array.fill[Byte](32)(1) + val replacementFrame = Array.fill[Byte](8)(1) + + val originalWorker = new Thread(() => { + try { + original.reservePartitionData(96) + original.pushPartitionData(0, originalFrame, originalFrame.length) + } catch { case failure: Throwable => originalFailure.set(failure) } + }) + originalWorker.start() + assert(client.pushCompletedBeforeReturn.await(5, TimeUnit.SECONDS)) + assert(client.pushCompletionCalls.get() == 1) + + val replacementWorker = new Thread(() => { + try { + replacement.reservePartitionData(24) + replacement.pushPartitionData(0, replacementFrame, replacementFrame.length) + } catch { case failure: Throwable => replacementFailure.set(failure) } + }) + replacementWorker.start() + + Thread.sleep(100) + assert(replacementWorker.isAlive) + assert(client.pushCalls.get() == 1) + + client.allowPushReturn.countDown() + originalWorker.join(5000) + replacementWorker.join(5000) + + assert(!originalWorker.isAlive) + assert(!replacementWorker.isAlive) + assert(originalFailure.get() == null) + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.pushCompletionCalls.get() == 2) + } + + test("undersized native-copy reservations fail before the Celeborn client is invoked") { + val client = new RecordingCelebornShuffleClient + val rejected = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 112) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 112) + val frame = Array.fill[Byte](32)(1) + + rejected.reservePartitionData(32) + val failure = intercept[IOException] { + rejected.pushPartitionData(0, frame, frame.length) + } + rejected.releasePartitionDataReservation() + + assert(failure.getMessage.contains("native encoding reservation")) + assert(client.pushCalls.get() == 0) + + replacement.reservePartitionData(96) + assert(replacement.pushPartitionData(0, frame, frame.length) == frame.length) + assert(client.pushCalls.get() == 1) + } + + test("native scratch reservations can exceed the frame cap but not executor admission") { + val client = new RecordingCelebornShuffleClient + val adapter = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 112) + + assert(adapter.maxFrameBytes() == 32) + adapter.reservePartitionData(96) + adapter.releasePartitionDataReservation() + + val failure = intercept[IOException] { + adapter.reservePartitionData(97) + } + assert(failure.getMessage.contains("reservation exceeds its byte limit")) + assert(client.pushCalls.get() == 0) + } + + test("silent mapper-ended transport completion restores executor-wide byte admission") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + client.silentlyCompletePushes = true + val ended = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val failure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + assert(ended.pushPartitionData(0, frame, frame.length) == frame.length) + val worker = new Thread(() => { + try replacement.pushPartitionData(0, frame, frame.length) + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + + Thread.sleep(100) + assert(client.pushCalls.get() == 1) + assert(client.completeNextPush()) + client.silentlyCompletePushes = false + + worker.join(5000) + if (worker.isAlive) { + replacement.abort() + worker.join(5000) + fail("silent transport completion did not restore executor-wide admission") + } + + assert(failure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + + test("metrics-only failures and silent completions each restore their own admission") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val original = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 96) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 96) + val failure = new AtomicReference[Throwable]() + val originalFrame = Array.fill[Byte](32)(1) + val replacementFrame = Array.fill[Byte](80)(1) + + assert( + original.pushPartitionData(0, originalFrame, originalFrame.length) == originalFrame.length) + assert( + original.pushPartitionData(1, originalFrame, originalFrame.length) == originalFrame.length) + + client.metricsFailureWithoutBatchRemoval = true + assert(client.completeNextPush()) + client.metricsFailureWithoutBatchRemoval = false + client.silentlyCompletePushes = true + assert(client.completeNextPush()) + client.silentlyCompletePushes = false + + val worker = new Thread(() => { + try replacement.pushPartitionData(0, replacementFrame, replacementFrame.length) + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + worker.join(5000) + + if (worker.isAlive) { + replacement.abort() + worker.join(5000) + fail("a metrics-only failure masked a separate silently completed request") + } + + assert(failure.get() == null) + assert(client.pushCalls.get() == 3) + assert(client.completeNextPush()) + } + + test("tracker reconciliation and delayed completion metrics do not release a live request") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val original = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 72) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 72) + val later = new CelebornShufflePartitionPusher(client, 19, 5, 1, 12, 9, 72) + val replacementFailure = new AtomicReference[Throwable]() + val laterFailure = new AtomicReference[Throwable]() + val smallFrame = Array.fill[Byte](8)(1) + val largeFrame = Array.fill[Byte](32)(1) + + assert(original.pushPartitionData(0, smallFrame, smallFrame.length) == smallFrame.length) + assert(original.pushPartitionData(1, smallFrame, smallFrame.length) == smallFrame.length) + + client.completionRemoved = new CountDownLatch(1) + client.allowCompletionMetrics = new CountDownLatch(1) + val completion = new Thread(() => client.completeNextPush()) + completion.start() + assert(client.completionRemoved.await(5, TimeUnit.SECONDS)) + + val replacementWorker = new Thread(() => { + try replacement.pushPartitionData(0, largeFrame, largeFrame.length) + catch { case error: Throwable => replacementFailure.set(error) } + }) + replacementWorker.start() + replacementWorker.join(5000) + + if (replacementWorker.isAlive) { + client.allowCompletionMetrics.countDown() + replacement.abort() + replacementWorker.join(5000) + fail("tracker reconciliation did not release the first completed request") + } + + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 3) + client.completionRemoved = null + client.allowCompletionMetrics.countDown() + completion.join(5000) + assert(!completion.isAlive) + + val laterWorker = new Thread(() => { + try later.pushPartitionData(0, smallFrame, smallFrame.length) + catch { case error: Throwable => laterFailure.set(error) } + }) + laterWorker.start() + + Thread.sleep(100) + assert(laterWorker.isAlive) + assert(client.pushCalls.get() == 3) + + assert(client.completeNextPush()) + laterWorker.join(5000) + assert(!laterWorker.isAlive) + assert(laterFailure.get() == null) + assert(client.pushCalls.get() == 4) + assert(client.completeNextPush()) + assert(client.completeNextPush()) + } + + test("completed requests do not claim an overlapping push's unsubmitted reservation") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val original = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 72) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 72) + val overlappingFailure = new AtomicReference[Throwable]() + val replacementFailure = new AtomicReference[Throwable]() + val smallFrame = Array.fill[Byte](8)(1) + val largeFrame = Array.fill[Byte](32)(1) + + assert(original.pushPartitionData(0, largeFrame, largeFrame.length) == largeFrame.length) + + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + val overlapping = new Thread(() => { + try original.pushPartitionData(1, smallFrame, smallFrame.length) + catch { case error: Throwable => overlappingFailure.set(error) } + }) + overlapping.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + client.pushStarted = null + assert(client.completeNextPush()) + + val replacementWorker = new Thread(() => { + try replacement.pushPartitionData(0, largeFrame, largeFrame.length) + catch { case error: Throwable => replacementFailure.set(error) } + }) + replacementWorker.start() + replacementWorker.join(5000) + + if (replacementWorker.isAlive) { + client.allowPush.countDown() + overlapping.join(5000) + assert(client.completeNextPush()) + replacementWorker.join(5000) + fail("a completed request consumed the smaller unsubmitted reservation's credit") + } + + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 3) + client.allowPush.countDown() + overlapping.join(5000) + assert(!overlapping.isAlive) + assert(overlappingFailure.get() == null) + assert(client.completeNextPush()) + assert(client.completeNextPush()) + } + + test("inline completion cannot release an older live request before submission is accepted") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val original = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 72) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 72) + val replacementFailure = new AtomicReference[Throwable]() + val smallFrame = Array.fill[Byte](8)(1) + val largeFrame = Array.fill[Byte](32)(1) + + assert(original.pushPartitionData(0, largeFrame, largeFrame.length) == largeFrame.length) + client.automaticallyCompleteLatestPush = true + assert(original.pushPartitionData(1, smallFrame, smallFrame.length) == smallFrame.length) + client.automaticallyCompleteLatestPush = false + assert(client.pushCompletionCalls.get() == 1) + + val replacementWorker = new Thread(() => { + try replacement.pushPartitionData(0, largeFrame, largeFrame.length) + catch { case error: Throwable => replacementFailure.set(error) } + }) + replacementWorker.start() + + Thread.sleep(100) + assert(replacementWorker.isAlive) + assert(client.pushCalls.get() == 2) + + assert(client.completeNextPush()) + replacementWorker.join(5000) + assert(!replacementWorker.isAlive) + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 3) + assert(client.completeNextPush()) + } + + test( + "native frame admission is reserved before encoding and released when encoding is skipped") { + val client = new RecordingCelebornShuffleClient + val first = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val second = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val secondFailure = new AtomicReference[Throwable]() + + first.reservePartitionData(32) + val blocked = new Thread(() => { + try { + second.reservePartitionData(32) + second.releasePartitionDataReservation() + } catch { case failure: Throwable => secondFailure.set(failure) } + }) + blocked.start() + + Thread.sleep(100) + assert(blocked.isAlive) + assert(client.pushCalls.get() == 0) + + first.releasePartitionDataReservation() + blocked.join(5000) + assert(!blocked.isAlive) + assert(secondFailure.get() == null) + } + + test("cancelled accepted requests retain executor admission until their transport callback") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val cancelled = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val replacementFailure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + assert(cancelled.pushPartitionData(0, frame, frame.length) == frame.length) + cancelled.abort() + + val worker = new Thread(() => { + try replacement.pushPartitionData(0, frame, frame.length) + catch { case failure: Throwable => replacementFailure.set(failure) } + }) + worker.start() + + Thread.sleep(100) + assert(worker.isAlive) + assert(client.pushCalls.get() == 1) + + assert(client.completeNextPush()) + worker.join(5000) + assert(!worker.isAlive) + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + + test("map completion fails and cleans up when accepted Celeborn pushes time out") { + val client = new RecordingCelebornShuffleClient + client.drainTimedOut = true + val adapter = pusher(client) + + assert(adapter.pushPartitionData(0, Array[Byte](1), 1) == 1) + + val failure = intercept[IOException] { + adapter.finish() + } + + assert(failure.getMessage.contains("Timed out")) + assert(client.cleanupCalls.get() == 1) + assert(client.mapperEndCalls.get() == 1) + } + + test("executor-wide admission rejects a frame larger than its configured budget") { + val client = new RecordingCelebornShuffleClient + val adapter = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 20) + + val failure = intercept[IOException] { + adapter.pushPartitionData(0, Array.fill[Byte](8)(1), 8) + } + + assert(failure.getMessage.contains("in-flight byte limit")) + assert(client.pushCalls.get() == 0) + } + + test("successful map completion drains pending pushes and reports per-reducer frame bytes") { + val client = new RecordingCelebornShuffleClient + val adapter = pusher(client) + + adapter.pushPartitionData(2, Array[Byte](1, 2, 3), 3) + adapter.pushPartitionData(5, Array[Byte](4, 5), 2) + adapter.pushPartitionData(2, Array[Byte](6), 1) + + val partitionLengths = adapter.finish() + assert(partitionLengths.sameElements(Array[Long](0, 0, 4, 0, 0, 2, 0, 0, 0))) + assert(client.mapperEndCalls.get() == 1) + assert(client.lastMapperEnd == ((19, 3, (4 << 16) | 7, 12))) + assert(client.cleanupCalls.get() == 0) + + assert(adapter.finish().sameElements(partitionLengths)) + assert(client.mapperEndCalls.get() == 1) + + adapter.abort() + adapter.abort() + assert(client.cleanupCalls.get() == 1) + } + + test("an empty map still commits all reducer partitions to Celeborn") { + val client = new RecordingCelebornShuffleClient + + assert(pusher(client).finish().sameElements(Array.fill[Long](9)(0L))) + assert(client.mapperEndCalls.get() == 1) + } + + test("asynchronous push failures observed by mapperEnd fail the task and clean up") { + val client = new RecordingCelebornShuffleClient + val expected = new IOException("a pending Celeborn push failed") + client.mapperEndFailure = expected + val adapter = pusher(client) + + adapter.pushPartitionData(0, Array[Byte](1), 1) + val actual = intercept[IOException] { + adapter.finish() + } + + assert(actual eq expected) + assert(client.mapperEndCalls.get() == 1) + assert(client.cleanupCalls.get() == 1) + } + + test("cleanup errors are suppressed without replacing the original push failure") { + val client = new RecordingCelebornShuffleClient + val expected = new IOException("the shuffle worker rejected the frame") + val cleanupFailure = new IllegalStateException("attempt cleanup also failed") + client.failure = expected + client.cleanupFailure = cleanupFailure + + val actual = intercept[IOException] { + pusher(client).pushPartitionData(0, Array[Byte](1), 1) + } + + assert(actual eq expected) + assert(actual.getSuppressed.sameElements(Array[Throwable](cleanupFailure))) + assert(client.cleanupCalls.get() == 1) + } + + test("completed and aborted map attempts reject subsequent native pushes") { + val completedClient = new RecordingCelebornShuffleClient + val completed = pusher(completedClient) + completed.finish() + + intercept[IOException] { + completed.pushPartitionData(0, Array[Byte](1), 1) + } + + val abortedClient = new RecordingCelebornShuffleClient + val aborted = pusher(abortedClient) + aborted.abort() + + intercept[IOException] { + aborted.pushPartitionData(0, Array[Byte](1), 1) + } + + assert(completedClient.lastPush == null) + assert(abortedClient.lastPush == null) + assert(abortedClient.cleanupCalls.get() == 1) + } + + test("map completion waits for an active native callback before invoking mapperEnd") { + val client = new RecordingCelebornShuffleClient + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + val adapter = pusher(client) + val pushFailure = new AtomicReference[Throwable]() + val finishFailure = new AtomicReference[Throwable]() + val finished = new CountDownLatch(1) + + val pushThread = new Thread(() => { + try adapter.pushPartitionData(4, Array[Byte](1, 2), 2) + catch { case error: Throwable => pushFailure.set(error) } + }) + pushThread.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + val finishThread = new Thread(() => { + try adapter.finish() + catch { case error: Throwable => finishFailure.set(error) } + finally finished.countDown() + }) + finishThread.start() + + assert(!finished.await(100, TimeUnit.MILLISECONDS)) + assert(client.mapperEndCalls.get() == 0) + + client.allowPush.countDown() + pushThread.join(5000) + finishThread.join(5000) + + assert(!pushThread.isAlive) + assert(!finishThread.isAlive) + assert(pushFailure.get() == null) + assert(finishFailure.get() == null) + assert(client.mapperEndCalls.get() == 1) + } + + test("failed task cleanup wakes a native push blocked in Celeborn backpressure") { + val client = new RecordingCelebornShuffleClient + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + client.cleanupUnblocksPush = true + val adapter = pusher(client) + val failure = new AtomicReference[Throwable]() + + val worker = new Thread(() => { + try adapter.pushPartitionData(0, Array[Byte](1), 1) + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + adapter.abort() + worker.join(5000) + + assert(!worker.isAlive) + assert(failure.get().isInstanceOf[IOException]) + assert(client.cleanupCalls.get() == 2) + assert(client.mapperEndCalls.get() == 0) + } + + test("an aborted client push cleans up Celeborn state recreated after initial cleanup") { + val client = new RecordingCelebornShuffleClient + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + client.recreatePushStateAfterCleanup = true + val adapter = pusher(client) + val failure = new AtomicReference[Throwable]() + + val worker = new Thread(() => { + try adapter.pushPartitionData(0, Array[Byte](1), 1) + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + assert(client.retainedPushState) + + adapter.abort() + worker.join(5000) + + assert(!worker.isAlive) + assert(failure.get().isInstanceOf[IOException]) + assert(failure.get().getMessage.contains("aborted during its push")) + assert(client.cleanupCalls.get() == 2) + assert(!client.retainedPushState) + + adapter.abort() + assert(client.cleanupCalls.get() == 2) + } + + test("a recreated cancelled push retains byte admission until its own transport completes") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + client.recreatePushStateAfterCleanup = true + val cancelled = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val cancelledFailure = new AtomicReference[Throwable]() + val replacementFailure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + val cancelledWorker = new Thread(() => { + try cancelled.pushPartitionData(0, frame, frame.length) + catch { case failure: Throwable => cancelledFailure.set(failure) } + }) + cancelledWorker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + cancelled.abort() + cancelledWorker.join(5000) + assert(!cancelledWorker.isAlive) + assert(cancelledFailure.get().isInstanceOf[IOException]) + assert(client.cleanupCalls.get() == 2) + + client.pushStarted = null + client.allowPush = null + val replacementWorker = new Thread(() => { + try replacement.pushPartitionData(0, frame, frame.length) + catch { case failure: Throwable => replacementFailure.set(failure) } + }) + replacementWorker.start() + + Thread.sleep(100) + assert(replacementWorker.isAlive) + assert(client.pushCalls.get() == 1) + + assert(client.completeNextPush()) + replacementWorker.join(5000) + assert(!replacementWorker.isAlive) + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + + test("inline completion on a recreated cancelled state does not strand byte admission") { + val client = new RecordingCelebornShuffleClient + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + client.recreatePushStateAfterCleanup = true + val cancelled = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val cancelledFailure = new AtomicReference[Throwable]() + val replacementFailure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + val cancelledWorker = new Thread(() => { + try cancelled.pushPartitionData(0, frame, frame.length) + catch { case error: Throwable => cancelledFailure.set(error) } + }) + cancelledWorker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + cancelled.abort() + cancelledWorker.join(5000) + assert(!cancelledWorker.isAlive) + assert(cancelledFailure.get().isInstanceOf[IOException]) + assert(client.pushCompletionCalls.get() == 1) + + client.pushStarted = null + client.allowPush = null + val replacementWorker = new Thread(() => { + try replacement.pushPartitionData(0, frame, frame.length) + catch { case error: Throwable => replacementFailure.set(error) } + }) + replacementWorker.start() + replacementWorker.join(5000) + + if (replacementWorker.isAlive) { + replacement.abort() + replacementWorker.join(5000) + fail("inline completion on the recreated state permanently retained byte admission") + } + + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.pushCompletionCalls.get() == 2) + } + + test("inline failure on a recreated cancelled state does not strand byte admission") { + val client = new RecordingCelebornShuffleClient + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + client.recreatePushStateAfterCleanup = true + client.metricsFailureWithoutBatchRemoval = true + val cancelled = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val cancelledFailure = new AtomicReference[Throwable]() + val replacementFailure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + val cancelledWorker = new Thread(() => { + try cancelled.pushPartitionData(0, frame, frame.length) + catch { case error: Throwable => cancelledFailure.set(error) } + }) + cancelledWorker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + cancelled.abort() + cancelledWorker.join(5000) + assert(!cancelledWorker.isAlive) + assert(cancelledFailure.get().isInstanceOf[IOException]) + assert(client.pushCompletionCalls.get() == 1) + + client.pushStarted = null + client.allowPush = null + client.metricsFailureWithoutBatchRemoval = false + val replacementWorker = new Thread(() => { + try replacement.pushPartitionData(0, frame, frame.length) + catch { case error: Throwable => replacementFailure.set(error) } + }) + replacementWorker.start() + replacementWorker.join(5000) + + if (replacementWorker.isAlive) { + replacement.abort() + replacementWorker.join(5000) + fail("inline failure on the recreated state permanently retained byte admission") + } + + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.pushCompletionCalls.get() == 2) + } + + test("failed task cleanup wakes map completion while it drains pending Celeborn pushes") { + val client = new RecordingCelebornShuffleClient + client.mapperEndStarted = new CountDownLatch(1) + client.allowMapperEnd = new CountDownLatch(1) + client.cleanupUnblocksMapperEnd = true + val adapter = pusher(client) + val failure = new AtomicReference[Throwable]() + + adapter.pushPartitionData(0, Array[Byte](1), 1) + val worker = new Thread(() => { + try adapter.finish() + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + assert(client.mapperEndStarted.await(5, TimeUnit.SECONDS)) + + adapter.abort() + worker.join(5000) + + assert(!worker.isAlive) + assert(failure.get().isInstanceOf[IOException]) + assert(client.mapperEndCalls.get() == 1) + assert(client.cleanupCalls.get() == 1) + } + + test("failed task cleanup wakes mapper-end completion drain for an accepted push") { + val client = new RecordingCelebornShuffleClient + client.drainStarted = new CountDownLatch(1) + client.allowDrain = new CountDownLatch(1) + client.cleanupUnblocksDrain = true + val adapter = pusher(client) + val failure = new AtomicReference[Throwable]() + + assert(adapter.pushPartitionData(0, Array[Byte](1), 1) == 1) + val worker = new Thread(() => { + try adapter.finish() + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + assert(client.drainStarted.await(5, TimeUnit.SECONDS)) + + adapter.abort() + worker.join(5000) + + assert(!worker.isAlive) + assert(failure.get().isInstanceOf[IOException]) + assert(client.cleanupCalls.get() == 1) + } + + test("cancellation interrupts mapperEnd RPC even when client cleanup cannot wake it") { + val client = new RecordingCelebornShuffleClient + client.mapperEndStarted = new CountDownLatch(1) + client.allowMapperEnd = new CountDownLatch(1) + val adapter = pusher(client) + val failure = new AtomicReference[Throwable]() + + val worker = new Thread(() => { + try adapter.finish() + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + assert(client.mapperEndStarted.await(5, TimeUnit.SECONDS)) + + adapter.abort() + worker.join(5000) + + assert(!worker.isAlive) + assert(failure.get().isInstanceOf[IOException]) + assert(client.cleanupCalls.get() == 1) + assert(client.allowMapperEnd.getCount == 1) } test("adapter rejects a missing client or incompatible Celeborn raw-push API") { @@ -239,6 +1341,27 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { assert(client.lastPush == null) } + test("task-scoped native registration validates its handle, callback, reducers, and frame") { + val adapter = pusher(new RecordingCelebornShuffleClient) + + val registration = CometExecIterator.RssPartitionPusherRegistration(1L, adapter, 9, 20) + assert(registration.handle == 1L) + assert(registration.pusher eq adapter) + assert(registration.numPartitions == 9) + assert(registration.maxFrameBytes == 20) + + Seq( + (0L, adapter, 9, 20), + (-1L, adapter, 9, 20), + (1L, null, 9, 20), + (1L, adapter, 0, 20), + (1L, adapter, 9, 19)).foreach { case (handle, callback, partitions, maxFrame) => + intercept[IllegalArgumentException] { + CometExecIterator.RssPartitionPusherRegistration(handle, callback, partitions, maxFrame) + } + } + } + test("factory recognizes the existing Celeborn shuffle manager") { val conf = new SparkConf(false).set(managerKey, managerClass) @@ -248,6 +1371,134 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { assert(!CelebornShufflePusherFactory.isEnabled(conf)) } + test("an acquired client remains owned when Celeborn shuffle-generation resolution fails") { + val client = new RecordingCelebornShuffleClient + val expected = new IOException("Celeborn shuffle-generation lookup failed") + var ownedClient: AnyRef = null + + val actual = intercept[IOException] { + val acquired = CelebornShufflePusherFactory.acquireClient(client, ownedClient = _) + assert(ownedClient eq acquired) + throw expected + } + + assert(actual eq expected) + assert(ownedClient eq client) + assert(client.cleanupCalls.get() == 0) + } + + test("failed Celeborn client acquisition does not record application ownership") { + val expected = new IOException("Celeborn client acquisition failed") + var ownedClient: AnyRef = null + + val actual = intercept[IOException] { + CelebornShufflePusherFactory.acquireClient(throw expected, ownedClient = _) + } + + assert(actual eq expected) + assert(ownedClient == null) + } + + test("an authorized retried map invalidates its ambiguous Celeborn shuffle generation") { + val client = new RecordingCelebornShuffleClient + val context = emptyTaskContext() + var driverInvalidated = false + + val failure = intercept[Exception] { + CelebornShufflePusherFactory.rejectRetriedAttempt( + client, + 7, + 91, + context, + true, + () => driverInvalidated = true) + } + + assert(failure.getClass.getName == "org.apache.spark.shuffle.FetchFailedException") + assert(failure.getMessage.contains("new shuffle generation")) + assert(client.lastInvalidatedShuffle == ((7, 91, context.taskAttemptId()))) + assert(driverInvalidated) + } + + test("factory invokes Celeborn's barrier failure-listener hook with the task handle") { + val client = new Object + val context = emptyTaskContext() + val handle = new ShuffleHandle(7) {} + + CelebornShufflePusherFactory.registerBarrierFailureListener( + classOf[RecordingCelebornSparkUtils], + classOf[Object], + classOf[ShuffleHandle], + client, + context, + handle) + + assert(RecordingCelebornSparkUtils.client eq client) + assert(RecordingCelebornSparkUtils.taskContext eq context) + assert(RecordingCelebornSparkUtils.shuffleHandle eq handle) + } + + test("a losing speculative map attempt never invalidates the winning shuffle generation") { + val client = new RecordingCelebornShuffleClient + val taskContext = emptyTaskContext() + + val failure = intercept[Exception] { + CelebornShufflePusherFactory.rejectRetriedAttempt(client, 7, 91, taskContext, false) + } + + assert(failure.getClass.getName == "org.apache.spark.executor.CommitDeniedException") + assert(failure.getMessage.contains("already owns")) + val reason = failure.getClass.getMethod("toTaskCommitDeniedReason").invoke(failure) + assert(reason.getClass.getName == "org.apache.spark.TaskCommitDenied") + assert( + !reason.getClass.getMethod("countTowardsTaskFailures").invoke(reason).asInstanceOf[Boolean]) + assert(client.lastInvalidatedShuffle == null) + } + + test("an authorized retry cannot proceed when Celeborn rejects generation invalidation") { + val client = new RecordingCelebornShuffleClient + client.generationInvalidated = false + var driverInvalidated = false + + val failure = intercept[IOException] { + CelebornShufflePusherFactory.rejectRetriedAttempt( + client, + 7, + 91, + emptyTaskContext(), + true, + () => driverInvalidated = true) + } + + assert(failure.getMessage.contains("Could not invalidate")) + assert(client.lastInvalidatedShuffle._1 == 7) + assert(!driverInvalidated) + } + + test("a speculative owner abandons its commit when Celeborn preserves the live original") { + val client = new RecordingCelebornShuffleClient + client.generationInvalidated = false + var driverAbandoned = false + + val failure = intercept[Exception] { + CelebornShufflePusherFactory.rejectRetriedAttempt( + client, + 7, + 91, + emptyTaskContext(), + true, + () => fail("a rejected generation must not be marked invalid"), + () => { + driverAbandoned = true + true + }) + } + + assert(failure.getClass.getName == "org.apache.spark.executor.CommitDeniedException") + assert(driverAbandoned) + assert(client.lastInvalidatedShuffle == null) + } + test("factory recognizes the composite Comet and Celeborn shuffle manager") { val conf = new SparkConf(false).set(managerKey, compositeManagerClass) @@ -327,6 +1578,33 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { assert(client.lastPush.numPartitions == 4) } + test("enabled factory requires admission for all three copies of a minimum native frame") { + val admissionKey = "spark.comet.shuffle.rss.maxInFlightBytes" + val taskContext = emptyTaskContext() + + val rejected = intercept[IllegalArgumentException] { + CelebornShufflePusherFactory.create( + enabledConf.set(admissionKey, "75"), + new RecordingCelebornShuffleClient, + celebornShuffleId = 27, + numMappers = 8, + numPartitions = 4, + taskContext = taskContext) + } + assert(rejected.getMessage.contains("in-flight byte limit")) + + val minimum = CelebornShufflePusherFactory + .create( + enabledConf.set(admissionKey, "76"), + new RecordingCelebornShuffleClient, + celebornShuffleId = 27, + numMappers = 8, + numPartitions = 4, + taskContext = taskContext) + .get + assert(minimum.maxFrameBytes() == 20) + } + test("factory encodes the stage and task attempt without loading the Celeborn client jar") { assert(CelebornShufflePusherFactory.encodeAttemptNumber(4, 7) == ((4 << 16) | 7)) assert(CelebornShufflePusherFactory.encodeAttemptNumber(0, 0) == 0) 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 new file mode 100644 index 00000000000..79f7c07b2e8 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala @@ -0,0 +1,620 @@ +/* + * 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.spark.sql.comet.execution.shuffle + +import java.io.IOException +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference + +import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext} +import org.apache.spark.executor.CommitDeniedException +import org.apache.spark.scheduler.OutputCommitCoordinator +import org.apache.spark.shuffle.{BaseShuffleHandle, IndexShuffleBlockResolver, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.RpcUtils + +import org.apache.comet.CometConf +import org.apache.comet.shuffle.{CelebornShufflePusherFactory, RecordingCelebornShuffleClient} + +/** + * Exercises a real Spark map task, native RSS planning, JNI callback, and Celeborn map lifecycle. + */ +class CometCelebornNativeShuffleWriterSuite extends CometTestBase { + + import testImplicits._ + + private final class LocalFallbackShuffleManager extends ShuffleManager { + var unregisteredShuffle: Option[Int] = None + + override def registerShuffle[K, V, C]( + shuffleId: Int, + dependency: ShuffleDependency[K, V, C]): ShuffleHandle = + new BaseShuffleHandle(shuffleId, dependency) + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = + throw new AssertionError("Native Comet data must never reach a local fallback writer") + + override def getReader[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = + throw new AssertionError("Native Comet data must never reach a local fallback reader") + + override def shuffleBlockResolver: ShuffleBlockResolver = null + + override def unregisterShuffle(shuffleId: Int): Boolean = { + unregisteredShuffle = Some(shuffleId) + true + } + + override def stop(): Unit = () + } + + private def withNativeShuffleDependency( + run: CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch] => Unit): Unit = + withNativeShuffleDependency(false, run) + + private def withNativeShuffleDependency( + useRangePartitioning: Boolean, + run: CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch] => Unit): Unit = { + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + "spark.sql.adaptive.enabled" -> "false") { + val rows = (0 until 24).map { value => + val key = if (useRangePartitioning) value % 2 else value + (key, s"row-$value") + } + withParquetTable(rows, "celeborn_rows") { + val input = sql("SELECT * FROM celeborn_rows") + val shuffled = + if (useRangePartitioning) input.repartitionByRange(10, $"_1") + else input.repartition(3, $"_1") + val exchange = shuffled.queryExecution.executedPlan + .collectFirst { case value: CometShuffleExchangeExec => + value + } + .getOrElse { + fail("Expected a native Comet shuffle exchange") + } + val dependency = exchange.shuffleDependency + .asInstanceOf[CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]] + run(dependency) + } + } + } + + test("preauthorized native Spark maps reject speculation and commit remotely without files") { + withNativeShuffleDependency { dependency => + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val client = new RecordingCelebornShuffleClient + val taskConf = SparkEnv.get.conf + .clone() + .set("spark.celeborn.master.endpoints", "recording-celeborn:9097") + val pusher = CelebornShufflePusherFactory + .create( + taskConf, + client, + celebornShuffleId = 91, + numMappers = numMappers, + numPartitions = dependency.partitioner.numPartitions, + taskContext = context) + .get + val coordinator = SparkEnv.get.outputCommitCoordinator + assert( + coordinator.canCommit( + context.stageId(), + context.stageAttemptNumber(), + context.partitionId(), + context.attemptNumber())) + assert( + !coordinator.canCommit( + context.stageId(), + context.stageAttemptNumber(), + context.partitionId(), + context.attemptNumber() + 1)) + assert(client.lastInvalidatedShuffle == null) + var commitValidations = 0 + val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds, + Some( + CelebornNativeShuffleDestination( + pusher, + 1024 * 1024, + dependency.partitioner.numPartitions, + commitAuthorized = true, + commitValidator = () => { + commitValidations += 1 + client.lastInvalidatedShuffle == null + }))) + + writer.write(inputs) + val status = writer.stop(success = true).get + val localFile = SparkEnv.get.shuffleManager.shuffleBlockResolver + .asInstanceOf[IndexShuffleBlockResolver] + .getDataFile(dependency.shuffleId, context.taskAttemptId()) + assert(commitValidations == 2) + assert(client.lastInvalidatedShuffle == null) + + ( + writer.getPartitionLengths(), + status.mapId, + context.taskAttemptId(), + client.mapperEndCalls.get(), + client.cleanupCalls.get(), + Option(client.lastPush).map(_.length).getOrElse(0), + localFile.exists(), + writer.stop(success = false).isEmpty) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1.length == 3)) + assert(results.exists(_._1.sum > 0)) + assert(results.forall(result => result._2 == result._3)) + assert(results.forall(_._4 == 1)) + assert(results.forall(_._5 == 1)) + assert(results.exists(_._6 >= 20)) + assert(results.forall(!_._7)) + assert(results.forall(_._8)) + } + } + + test("asynchronous remote push failures abort the Spark map attempt without a MapStatus") { + withNativeShuffleDependency { dependency => + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val client = new RecordingCelebornShuffleClient + val expected = new IOException("a pending Celeborn worker push failed") + client.mapperEndFailure = expected + val taskConf = SparkEnv.get.conf + .clone() + .set("spark.celeborn.master.endpoints", "recording-celeborn:9097") + val pusher = CelebornShufflePusherFactory + .create( + taskConf, + client, + celebornShuffleId = 93, + numMappers = numMappers, + numPartitions = dependency.partitioner.numPartitions, + taskContext = context) + .get + val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds, + Some( + CelebornNativeShuffleDestination( + pusher, + 1024 * 1024, + dependency.partitioner.numPartitions))) + + val failure = + try { + writer.write(inputs) + null + } catch { + case error: IOException => error + } + + ( + failure eq expected, + client.mapperEndCalls.get(), + client.cleanupCalls.get(), + writer.stop(success = false).isEmpty, + writer.mapStatus == null) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1)) + assert(results.forall(_._2 == 1)) + assert(results.forall(_._3 == 1)) + assert(results.forall(_._4)) + assert(results.forall(_._5)) + } + } + + test("low-cardinality range shuffles use Spark's actual reducer partition count") { + withNativeShuffleDependency( + useRangePartitioning = true, + dependency => { + val actualPartitions = dependency.partitioner.numPartitions + val requestedPartitions = dependency.outputPartitioning.get.numPartitions + assert(actualPartitions < requestedPartitions) + + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val client = new RecordingCelebornShuffleClient + val taskConf = SparkEnv.get.conf + .clone() + .set("spark.celeborn.master.endpoints", "recording-celeborn:9097") + val pusher = CelebornShufflePusherFactory + .create(taskConf, client, 97, numMappers, actualPartitions, context) + .get + val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds, + Some(CelebornNativeShuffleDestination(pusher, 1024 * 1024, actualPartitions))) + + writer.write(inputs) + writer.stop(success = true).get + ( + writer.getPartitionLengths().length, + Option(client.lastPush).map(_.numPartitions).getOrElse(actualPartitions)) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1 == actualPartitions)) + assert(results.forall(_._2 == actualPartitions)) + }) + } + + test("a losing map attempt cannot commit Celeborn data or publish Spark MapStatus") { + withNativeShuffleDependency { dependency => + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val client = new RecordingCelebornShuffleClient + val taskConf = SparkEnv.get.conf + .clone() + .set("spark.celeborn.master.endpoints", "recording-celeborn:9097") + val pusher = CelebornShufflePusherFactory + .create( + taskConf, + client, + 98, + numMappers, + dependency.partitioner.numPartitions, + context) + .get + val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds, + Some( + CelebornNativeShuffleDestination( + pusher, + 1024 * 1024, + dependency.partitioner.numPartitions))) + + assert( + SparkEnv.get.outputCommitCoordinator.canCommit( + context.stageId(), + context.stageAttemptNumber(), + context.partitionId(), + context.attemptNumber() + 1)) + val failure = + try { + writer.write(inputs) + null + } catch { + case error: CommitDeniedException => error + } + + ( + failure != null && failure.getMessage.contains("already owns") && + !failure.toTaskCommitDeniedReason.countTowardsTaskFailures, + client.mapperEndCalls.get(), + client.cleanupCalls.get(), + writer.mapStatus == null, + writer.stop(success = false).isEmpty) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1)) + assert(results.forall(_._2 == 0)) + assert(results.forall(_._3 == 1)) + assert(results.forall(_._4)) + assert(results.forall(_._5)) + } + } + + test("a generation invalidated during mapperEnd cannot publish a stale Spark MapStatus") { + withNativeShuffleDependency { dependency => + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val client = new RecordingCelebornShuffleClient + val taskConf = SparkEnv.get.conf + .clone() + .set("spark.celeborn.master.endpoints", "recording-celeborn:9097") + val pusher = CelebornShufflePusherFactory + .create( + taskConf, + client, + 99, + numMappers, + dependency.partitioner.numPartitions, + context) + .get + assert( + SparkEnv.get.outputCommitCoordinator.canCommit( + context.stageId(), + context.stageAttemptNumber(), + context.partitionId(), + context.attemptNumber())) + var validations = 0 + val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds, + Some( + CelebornNativeShuffleDestination( + pusher, + 1024 * 1024, + dependency.partitioner.numPartitions, + commitAuthorized = true, + commitValidator = () => { + validations += 1 + validations == 1 + }))) + + val failure = + try { + writer.write(inputs) + null + } catch { + case error: CommitDeniedException => error + } + + ( + failure != null && !failure.toTaskCommitDeniedReason.countTowardsTaskFailures, + validations, + client.mapperEndCalls.get(), + client.cleanupCalls.get(), + writer.mapStatus == null, + writer.stop(success = false).isEmpty) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1)) + assert(results.forall(_._2 == 2)) + assert(results.forall(_._3 == 1)) + assert(results.forall(_._4 == 1)) + assert(results.forall(_._5)) + assert(results.forall(_._6)) + } + } + + test("executor RPC coordinates Celeborn generations with the Spark driver's commit owners") { + val sparkEnvironment = SparkEnv.get + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val stageId = 812345 + val endpointName = "CometCelebornShuffleGenerationCoordinatorSuite" + val endpoint = sparkEnvironment.rpcEnv.setupEndpoint( + endpointName, + new CelebornShuffleGenerationEndpoint(sparkEnvironment.rpcEnv, coordinator)) + + try { + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(0)) + val canCommit = sparkCoordinator.getClass + .getMethod( + "handleAskPermissionToCommit", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE) + val driver = + RpcUtils.makeDriverRef(endpointName, sparkEnvironment.conf, sparkEnvironment.rpcEnv) + val original = PrepareCelebornShuffleGeneration(7, 91, stageId, 0, 1) + val replacement = original.copy(celebornShuffleId = 92, stageAttempt = 1) + + assert(driver.askSync[Boolean](original)) + assert( + canCommit + .invoke(sparkCoordinator, Int.box(stageId), Int.box(0), Int.box(0), Int.box(0)) + .asInstanceOf[Boolean]) + assert(driver.askSync[Boolean](replacement)) + assert( + canCommit + .invoke(sparkCoordinator, Int.box(stageId), Int.box(1), Int.box(0), Int.box(0)) + .asInstanceOf[Boolean]) + assert(!driver.askSync[Boolean](original)) + } finally { + sparkEnvironment.rpcEnv.stop(endpoint) + } + } + + test("task interruption proactively cleans up a blocked Celeborn native push") { + val client = new RecordingCelebornShuffleClient + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + client.cleanupUnblocksPush = true + val taskContext = TaskContext.empty() + val pusher = + new org.apache.comet.shuffle.CelebornShufflePartitionPusher(client, 91, 0, 0, 1, 1) + val callbackFailure = new AtomicReference[Throwable]() + val cleanupFailure = new AtomicReference[Throwable]() + val watcher = CelebornNativeShuffleDestination.watchForCancellation( + taskContext, + pusher, + failure => cleanupFailure.set(failure)) + + try { + val worker = new Thread(() => { + try pusher.pushPartitionData(0, Array[Byte](1), 1) + catch { case failure: Throwable => callbackFailure.set(failure) } + }) + worker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + taskContext.markInterrupted("the Spark task was cancelled") + worker.join(5000) + + assert(!worker.isAlive) + assert(callbackFailure.get().isInstanceOf[IOException]) + assert(cleanupFailure.get() == null) + assert(client.cleanupCalls.get() == 2) + } finally { + watcher.cancel(false) + pusher.abort() + } + } + + test("task interruption proactively wakes blocked Celeborn map completion") { + val client = new RecordingCelebornShuffleClient + client.mapperEndStarted = new CountDownLatch(1) + client.allowMapperEnd = new CountDownLatch(1) + client.cleanupUnblocksMapperEnd = true + val taskContext = TaskContext.empty() + val pusher = + new org.apache.comet.shuffle.CelebornShufflePartitionPusher(client, 92, 0, 0, 1, 1) + val callbackFailure = new AtomicReference[Throwable]() + val cleanupFailure = new AtomicReference[Throwable]() + val watcher = CelebornNativeShuffleDestination.watchForCancellation( + taskContext, + pusher, + failure => cleanupFailure.set(failure)) + + try { + val worker = new Thread(() => { + try pusher.finish() + catch { case failure: Throwable => callbackFailure.set(failure) } + }) + worker.start() + assert(client.mapperEndStarted.await(5, TimeUnit.SECONDS)) + + taskContext.markInterrupted("the Spark task was cancelled during map completion") + worker.join(5000) + + assert(!worker.isAlive) + assert(callbackFailure.get().isInstanceOf[IOException]) + assert(cleanupFailure.get() == null) + assert(client.mapperEndCalls.get() == 1) + assert(client.cleanupCalls.get() == 1) + } finally { + watcher.cancel(false) + pusher.abort() + } + } + + test("the existing local native map writer still commits local files and MapStatus") { + withNativeShuffleDependency { dependency => + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds) + + writer.write(inputs) + val status = writer.stop(success = true).get + val resolver = SparkEnv.get.shuffleManager.shuffleBlockResolver + .asInstanceOf[IndexShuffleBlockResolver] + + ( + writer.getPartitionLengths(), + status.mapId, + context.taskAttemptId(), + resolver.getDataFile(dependency.shuffleId, context.taskAttemptId()).exists(), + resolver.getIndexFile(dependency.shuffleId, context.taskAttemptId()).exists()) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1.length == 3)) + assert(results.exists(_._1.sum > 0)) + assert(results.forall(result => result._2 == result._3)) + assert(results.forall(_._4)) + assert(results.forall(_._5)) + } + } + + test("Celeborn native registration rejects and cleans up a delegated local fallback") { + withNativeShuffleDependency { dependency => + val backend = new LocalFallbackShuffleManager + val manager = + new CometCelebornShuffleManager(spark.sparkContext.getConf, true, (_, _) => backend) + + val failure = intercept[UnsupportedOperationException] { + manager.registerShuffle(dependency.shuffleId, dependency) + } + + assert(failure.getMessage.contains("local fallback")) + assert(backend.unregisteredShuffle.contains(dependency.shuffleId)) + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala index 4393a0070c3..c37f575527e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala @@ -21,9 +21,12 @@ package org.apache.spark.sql.comet.execution.shuffle import org.scalatest.funsuite.AnyFunSuite -import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext} +import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext, TaskEndReason, UnknownReason} +import org.apache.spark.scheduler.OutputCommitCoordinator import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, RecordingCelebornShuffleClient} + class CometCelebornShuffleManagerSuite extends AnyFunSuite { private class RecordingShuffleManager extends ShuffleManager { @@ -112,6 +115,352 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(conf.get("spark.celeborn.master.endpoints") == "existing-master:9097") } + test("native frame limits reserve all submission copies within executor admission") { + val client = new RecordingCelebornShuffleClient + val pusher = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 112) + + assert(CometCelebornShuffleManager.maxNativeFrameBytes(64, pusher) == 32) + assert(CometCelebornShuffleManager.maxNativeFrameBytes(32, pusher) == 32) + assert(CometCelebornShuffleManager.maxNativeFrameBytes(24, pusher) == 24) + pusher.reservePartitionData(96) + pusher.releasePartitionDataReservation() + assert(pusher.pushPartitionData(0, Array.fill[Byte](32)(1), 32) == 32) + assert(client.lastPush.length == 32) + + val minimum = + new CelebornShufflePartitionPusher(new RecordingCelebornShuffleClient, 19, 3, 1, 12, 9, 76) + assert(CometCelebornShuffleManager.maxNativeFrameBytes(64, minimum) == 20) + } + + test("a new Celeborn generation resets successful map commit owners on the Spark driver") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val firstGeneration = PrepareCelebornShuffleGeneration(7, 91, 12, 0, 2) + val replacementGeneration = PrepareCelebornShuffleGeneration(7, 92, 12, 1, 2) + val stageStart = sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + val canCommit = sparkCoordinator.getClass + .getMethod( + "handleAskPermissionToCommit", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE) + + def authorize(stageAttempt: Int, partition: Int, taskAttempt: Int): Boolean = + canCommit + .invoke( + sparkCoordinator, + Int.box(12), + Int.box(stageAttempt), + Int.box(partition), + Int.box(taskAttempt)) + .asInstanceOf[Boolean] + + stageStart.invoke(sparkCoordinator, Int.box(12), Int.box(1)) + assert(coordinator.prepareGeneration(firstGeneration)) + assert(authorize(0, 0, 0)) + assert(authorize(0, 1, 1)) + assert(!coordinator.prepareGeneration(firstGeneration.copy(celebornShuffleId = 92))) + assert(!authorize(1, 0, 0)) + + assert(coordinator.prepareGeneration(replacementGeneration)) + assert(authorize(1, 0, 0)) + assert(coordinator.prepareGeneration(replacementGeneration)) + assert(!authorize(1, 0, 1)) + assert(authorize(1, 1, 0)) + assert(!coordinator.prepareGeneration(firstGeneration)) + } + + test("early map ownership denies speculation and releases only after genuine task failure") { + val coordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val stageId = 74 + coordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(coordinator, Int.box(stageId), Int.box(0)) + val canCommit = coordinator.getClass + .getMethod( + "handleAskPermissionToCommit", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE) + + def authorize(attempt: Int): Boolean = + canCommit + .invoke(coordinator, Int.box(stageId), Int.box(0), Int.box(0), Int.box(attempt)) + .asInstanceOf[Boolean] + + assert(authorize(0)) + assert(!authorize(1)) + // Spark's coordinator also rejects the exact same owner when asked twice. + assert(!authorize(0)) + + coordinator.getClass + .getMethod( + "taskCompleted", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + classOf[TaskEndReason]) + .invoke(coordinator, Int.box(stageId), Int.box(0), Int.box(0), Int.box(0), UnknownReason) + + assert(authorize(1)) + } + + test("driver claims survive concurrent generation resets and reject stale map commits") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val stageId = 75 + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(1)) + + val firstGeneration = PrepareCelebornShuffleGeneration(7, 91, stageId, 0, 2) + assert(coordinator.prepareGeneration(firstGeneration)) + val originalClaim = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 0, 0, 0)) + assert(originalClaim.authorized) + assert(!coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 0, 0, 1)).authorized) + val originalPrepared = coordinator.prepareGenerationAndClaim(firstGeneration + .copy(mapId = 0, taskAttempt = 0, claimEpoch = originalClaim.epoch, claimAuthorized = true)) + assert(originalPrepared == originalClaim) + val originalValidation = + ValidateCelebornMapAttempt(7, 91, stageId, 0, 0, 0, originalPrepared.epoch) + assert(coordinator.validateMapAttempt(originalValidation)) + + // Another partition can claim before its peer prepares and resets the next generation. + val concurrentClaim = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 1, 1, 0)) + assert(concurrentClaim.authorized) + val blockedClaim = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 1, 0, 0)) + assert(!blockedClaim.authorized) + val replacementGeneration = PrepareCelebornShuffleGeneration(7, 92, stageId, 1, 2) + val replacementClaim = coordinator.prepareGenerationAndClaim( + replacementGeneration.copy( + mapId = 0, + taskAttempt = 0, + claimEpoch = blockedClaim.epoch, + claimAuthorized = blockedClaim.authorized)) + assert(replacementClaim.authorized) + assert(replacementClaim.epoch > originalPrepared.epoch) + + // The reset erased the other partition's otherwise successful early authorization. + val recoveredConcurrentClaim = coordinator.prepareGenerationAndClaim( + replacementGeneration.copy( + mapId = 1, + taskAttempt = 0, + claimEpoch = concurrentClaim.epoch, + claimAuthorized = true)) + assert(recoveredConcurrentClaim.authorized) + assert(recoveredConcurrentClaim.epoch == replacementClaim.epoch) + assert(!coordinator.validateMapAttempt(originalValidation)) + assert(!coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 0, 1, 2)).authorized) + + val replacementValidation = + ValidateCelebornMapAttempt(7, 92, stageId, 1, 0, 0, replacementClaim.epoch) + assert(coordinator.validateMapAttempt(replacementValidation)) + assert( + coordinator.invalidateGeneration( + InvalidateCelebornShuffleGeneration(7, 92, stageId, 1, replacementClaim.epoch))) + assert(!coordinator.validateMapAttempt(replacementValidation)) + val rejectedOldClaim = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 1, 1, 1)) + assert(!rejectedOldClaim.authorized) + assert( + !coordinator + .prepareGenerationAndClaim( + replacementGeneration.copy( + mapId = 1, + taskAttempt = 1, + claimEpoch = rejectedOldClaim.epoch, + claimAuthorized = false)) + .authorized) + + val nextGeneration = PrepareCelebornShuffleGeneration(7, 93, stageId, 2, 2) + val nextEarlyClaim = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 2, 0, 0)) + val nextPrepared = coordinator.prepareGenerationAndClaim( + nextGeneration.copy( + mapId = 0, + taskAttempt = 0, + claimEpoch = nextEarlyClaim.epoch, + claimAuthorized = nextEarlyClaim.authorized)) + assert(nextPrepared.authorized) + assert( + coordinator.validateMapAttempt( + ValidateCelebornMapAttempt(7, 93, stageId, 2, 0, 0, nextPrepared.epoch))) + } + + test("speculation arriving first reserves its original without creating phantom owners") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val stageId = 76 + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(0)) + + val speculative = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(8, stageId, 0, 0, 1)) + assert(!speculative.authorized) + assert(!speculative.requiresGenerationResolution) + val original = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(8, stageId, 0, 0, 0)) + assert(original.authorized) + assert(original.epoch == speculative.epoch) + + sparkCoordinator.getClass + .getMethod( + "taskCompleted", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + classOf[TaskEndReason]) + .invoke( + sparkCoordinator, + Int.box(stageId), + Int.box(0), + Int.box(0), + Int.box(0), + UnknownReason) + + // Spark deliberately does not mark CommitDenied attempt 1 as failed; Comet must skip it. + val genuineRetry = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(8, stageId, 0, 0, 2)) + assert(genuineRetry.authorized) + } + + test("a first retry can resolve a replacement generation after its original failed in input") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val stageId = 77 + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(0)) + + val oldGeneration = PrepareCelebornShuffleGeneration(9, 91, stageId, 0, 1) + assert(coordinator.prepareGeneration(oldGeneration)) + assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(9, stageId, 0, 0, 0)).authorized) + sparkCoordinator.getClass + .getMethod( + "taskCompleted", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + classOf[TaskEndReason]) + .invoke( + sparkCoordinator, + Int.box(stageId), + Int.box(1), + Int.box(0), + Int.box(0), + UnknownReason) + + val retry = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(9, stageId, 1, 0, 1)) + assert(!retry.authorized) + assert(retry.requiresGenerationResolution) + + val prepared = coordinator.prepareGenerationAndClaim( + PrepareCelebornShuffleGeneration( + 9, + 92, + stageId, + 1, + 1, + mapId = 0, + taskAttempt = 1, + claimEpoch = retry.epoch, + claimAuthorized = retry.authorized)) + assert(prepared.authorized) + assert(!prepared.requiresGenerationResolution) + } + + test("a replacement-stage original blocks speculation before its generation is resolved") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val stageId = 78 + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(0)) + + assert(coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(10, 91, stageId, 0, 1))) + val original = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(10, stageId, 1, 0, 0)) + assert(original.authorized) + + val speculative = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(10, stageId, 1, 0, 1)) + assert(!speculative.authorized) + assert(!speculative.requiresGenerationResolution) + } + + test("a speculative replacement owner yields its commit to an original blocked in input") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator, _ => false) + val stageId = 79 + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(0)) + + val oldGeneration = PrepareCelebornShuffleGeneration(11, 91, stageId, 0, 1) + assert(coordinator.prepareGeneration(oldGeneration)) + assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 0, 0, 0)).authorized) + + val speculative = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 1, 0, 1)) + assert(!speculative.authorized) + assert(speculative.requiresGenerationResolution) + val replacementGeneration = PrepareCelebornShuffleGeneration(11, 92, stageId, 1, 1) + val speculativeOwner = coordinator.prepareGenerationAndClaim( + replacementGeneration.copy( + mapId = 0, + taskAttempt = 1, + claimEpoch = speculative.epoch, + claimAuthorized = speculative.authorized)) + assert(speculativeOwner.authorized) + assert( + coordinator.abandonMapAttempt( + AbandonCelebornMapAttempt(11, 92, stageId, 1, 0, 1, speculativeOwner.epoch, 101L))) + + val original = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 1, 0, 0)) + assert(original.authorized) + val preparedOriginal = coordinator.prepareGenerationAndClaim( + replacementGeneration.copy( + mapId = 0, + taskAttempt = 0, + claimEpoch = original.epoch, + claimAuthorized = original.authorized)) + assert(preparedOriginal.authorized) + assert( + coordinator.validateMapAttempt( + ValidateCelebornMapAttempt(11, 92, stageId, 1, 0, 0, preparedOriginal.epoch))) + } + + test( + "a later stage attempt preserves commit owners when the Celeborn generation is unchanged") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val generation = PrepareCelebornShuffleGeneration(7, 91, 12, 0, 1) + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(12), Int.box(0)) + val canCommit = sparkCoordinator.getClass + .getMethod( + "handleAskPermissionToCommit", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE) + + assert(coordinator.prepareGeneration(generation)) + assert( + canCommit + .invoke(sparkCoordinator, Int.box(12), Int.box(0), Int.box(0), Int.box(0)) + .asInstanceOf[Boolean]) + + assert(coordinator.prepareGeneration(generation.copy(stageAttempt = 1))) + assert( + !canCommit + .invoke(sparkCoordinator, Int.box(12), Int.box(1), Int.box(0), Int.box(1)) + .asInstanceOf[Boolean]) + } + test("ordinary shuffle registration preserves the existing Celeborn handle and fallback") { val backend = new RecordingShuffleManager val composite = manager(backend) From 413edb63d2d32bae101e6bbf2271f8cad326d769 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:19:20 +0000 Subject: [PATCH 07/12] fix: handle Celeborn push cancellation, retries, and nested batches --- .../src/writers/rss/rss_partition_writer.rs | 295 +++++++++++++++++- .../CelebornShufflePartitionPusher.java | 36 ++- .../shuffle/CometCelebornShuffleManager.scala | 37 ++- .../CelebornShufflePartitionPusherSuite.scala | 86 ++++- .../CometCelebornShuffleManagerSuite.scala | 77 +++++ 5 files changed, 500 insertions(+), 31 deletions(-) diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs index 406003601dc..15ecd2ccac2 100644 --- a/native/shuffle/src/writers/rss/rss_partition_writer.rs +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -168,10 +168,12 @@ impl RssPartitionWriter

{ // compacting first; ordinary oversized batches can still be split before admission. let original_ipc_data_size = Self::estimated_pre_compaction_ipc_data_size(batch)?; let compaction_scratch = Self::estimated_compaction_scratch(batch)?; - if compaction_scratch == 0 - && original_ipc_data_size > self.max_frame_bytes - && batch.num_rows() > 1 - { + // Dictionary compaction can shrink values but cannot remove keys, list/map offsets, or + // other ordinary columns. Split before admission when those surviving buffers alone make + // one frame impossible, even if a different or nested column contains a dictionary. + let minimum_compacted_ipc_data_size = + Self::estimated_minimum_compacted_ipc_data_size(batch)?; + if minimum_compacted_ipc_data_size > self.max_frame_bytes && batch.num_rows() > 1 { return self.write_split_batch(partition_id, batch, encode_time, write_time); } @@ -300,11 +302,31 @@ impl RssPartitionWriter

{ /// dictionary values remain fully charged because dense garbage collection traverses them. fn estimated_pre_compaction_ipc_data_size(batch: &RecordBatch) -> Result { batch.columns().iter().try_fold(0usize, |total, column| { - Ok(total.saturating_add(Self::estimated_live_array_ipc_data_size(column.as_ref())?)) + Ok( + total.saturating_add(Self::estimated_live_array_ipc_data_size( + column.as_ref(), + true, + )?), + ) + }) + } + + /// Lower-bound the IPC buffers that must survive any dictionary-value compaction. + fn estimated_minimum_compacted_ipc_data_size(batch: &RecordBatch) -> Result { + batch.columns().iter().try_fold(0usize, |total, column| { + Ok( + total.saturating_add(Self::estimated_live_array_ipc_data_size( + column.as_ref(), + false, + )?), + ) }) } - fn estimated_live_array_ipc_data_size(array: &dyn Array) -> Result { + fn estimated_live_array_ipc_data_size( + array: &dyn Array, + include_dictionary_values: bool, + ) -> Result { let data = array.to_data(); let child_logical = data.child_data().iter().try_fold(0usize, |total, child| { Ok::<_, DataFusionError>(total.saturating_add(child.get_slice_memory_size()?)) @@ -319,7 +341,14 @@ impl RssPartitionWriter

{ let children = if let Some(dictionary) = array.as_any_dictionary_opt() { - Self::estimated_live_array_ipc_data_size(dictionary.values().as_ref())? + if include_dictionary_values { + Self::estimated_live_array_ipc_data_size( + dictionary.values().as_ref(), + include_dictionary_values, + )? + } else { + 0 + } } else { match array.data_type() { DataType::List(_) => { @@ -328,7 +357,10 @@ impl RssPartitionWriter

{ let start = offsets[0] as usize; let end = offsets[offsets.len() - 1] as usize; let live_values = list.values().slice(start, end - start); - Self::estimated_live_array_ipc_data_size(live_values.as_ref())? + Self::estimated_live_array_ipc_data_size( + live_values.as_ref(), + include_dictionary_values, + )? } DataType::LargeList(_) => { let list = array.as_any().downcast_ref::().unwrap(); @@ -336,11 +368,17 @@ impl RssPartitionWriter

{ let start = offsets[0] as usize; let end = offsets[offsets.len() - 1] as usize; let live_values = list.values().slice(start, end - start); - Self::estimated_live_array_ipc_data_size(live_values.as_ref())? + Self::estimated_live_array_ipc_data_size( + live_values.as_ref(), + include_dictionary_values, + )? } DataType::FixedSizeList(_, _) => { let list = array.as_any().downcast_ref::().unwrap(); - Self::estimated_live_array_ipc_data_size(list.values().as_ref())? + Self::estimated_live_array_ipc_data_size( + list.values().as_ref(), + include_dictionary_values, + )? } DataType::Struct(_) => { let structure = array.as_any().downcast_ref::().unwrap(); @@ -349,7 +387,10 @@ impl RssPartitionWriter

{ .iter() .try_fold(0usize, |total, column| { Ok::<_, DataFusionError>(total.saturating_add( - Self::estimated_live_array_ipc_data_size(column.as_ref())?, + Self::estimated_live_array_ipc_data_size( + column.as_ref(), + include_dictionary_values, + )?, )) })? } @@ -359,7 +400,10 @@ impl RssPartitionWriter

{ let start = offsets[0] as usize; let end = offsets[offsets.len() - 1] as usize; let live_entries = map.entries().slice(start, end - start); - Self::estimated_live_array_ipc_data_size(&live_entries)? + Self::estimated_live_array_ipc_data_size( + &live_entries, + include_dictionary_values, + )? } _ => data.child_data().iter().try_fold(0usize, |total, child| { let logical = child.get_slice_memory_size()?; @@ -1604,6 +1648,233 @@ mod tests { } } + #[test] + fn rss_partition_writer_splits_oversized_lists_and_maps_before_admission() { + struct BoundedAdmissionPusher { + recorder: ReservationRecordingPusher, + capacity: usize, + } + + impl PartitionPusher for BoundedAdmissionPusher { + fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> { + if reservation_bytes > self.capacity { + return Err(DataFusionError::External(Box::new(io::Error::other( + "nested batch exceeded admission before it was split", + )))); + } + self.recorder.reserve_partition_data(reservation_bytes) + } + + fn release_partition_data_reservation(&self) -> Result<()> { + self.recorder.release_partition_data_reservation() + } + + fn push_partition_data(&self, partition_id: usize, frame: &[u8]) -> Result<()> { + self.recorder.push_partition_data(partition_id, frame) + } + } + + let row_count = 4_096; + let values: ArrayRef = Arc::new(Int32Array::from_iter_values(0..row_count)); + let offsets = OffsetBuffer::new((0..=row_count).collect::>().into()); + let item = Arc::new(Field::new("item", DataType::Int32, false)); + let list = ListArray::try_new( + Arc::clone(&item), + offsets.clone(), + Arc::clone(&values), + None, + ) + .unwrap(); + let list_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + DataType::List(item), + false, + )])), + vec![Arc::new(list)], + ) + .unwrap(); + + let fields = vec![ + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::new(Field::new("value", DataType::Int32, false)), + ]; + let entries = StructArray::try_new( + fields.clone().into(), + vec![Arc::clone(&values), Arc::clone(&values)], + None, + ) + .unwrap(); + let entry = Arc::new(Field::new( + "entries", + DataType::Struct(fields.into()), + false, + )); + let map = + MapArray::try_new(Arc::clone(&entry), offsets.clone(), entries, None, false).unwrap(); + let map_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "entries", + DataType::Map(entry, false), + false, + )])), + vec![Arc::new(map)], + ) + .unwrap(); + + let dictionary_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)); + let tiny_dictionary: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![0; row_count as usize]), + Arc::new(Int32Array::from(vec![42])), + ) + .unwrap(), + ); + let dictionary_field = Field::new("dictionary", dictionary_type.clone(), false); + + let mixed_list_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + list_batch.schema().field(0).clone(), + dictionary_field.clone(), + ])), + vec![ + Arc::clone(list_batch.column(0)), + Arc::clone(&tiny_dictionary), + ], + ) + .unwrap(); + let mixed_map_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + map_batch.schema().field(0).clone(), + dictionary_field, + ])), + vec![ + Arc::clone(map_batch.column(0)), + Arc::clone(&tiny_dictionary), + ], + ) + .unwrap(); + + let dictionary_item = Arc::new(Field::new("item", dictionary_type.clone(), false)); + let nested_list = ListArray::try_new( + Arc::clone(&dictionary_item), + offsets.clone(), + Arc::clone(&tiny_dictionary), + None, + ) + .unwrap(); + let nested_dictionary_list_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + DataType::List(dictionary_item), + false, + )])), + vec![Arc::new(nested_list)], + ) + .unwrap(); + + let dictionary_entries_fields = vec![ + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::new(Field::new("value", dictionary_type, false)), + ]; + let dictionary_entries = StructArray::try_new( + dictionary_entries_fields.clone().into(), + vec![values, tiny_dictionary], + None, + ) + .unwrap(); + let dictionary_entry = Arc::new(Field::new( + "entries", + DataType::Struct(dictionary_entries_fields.into()), + false, + )); + let nested_map = MapArray::try_new( + Arc::clone(&dictionary_entry), + offsets, + dictionary_entries, + None, + false, + ) + .unwrap(); + let nested_dictionary_map_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "entries", + DataType::Map(dictionary_entry, false), + false, + )])), + vec![Arc::new(nested_map)], + ) + .unwrap(); + + let frame_limit = 2_048; + let admission_capacity = 8_192 - 16; + for input in [ + list_batch, + map_batch, + mixed_list_batch, + mixed_map_batch, + nested_dictionary_list_batch, + nested_dictionary_map_batch, + ] { + let first_split = input.slice(0, row_count as usize / 2); + let split_ipc_data_size = + RssPartitionWriter::::estimated_pre_compaction_ipc_data_size( + &first_split, + ) + .unwrap(); + let split_compaction_scratch = + RssPartitionWriter::::estimated_compaction_scratch( + &first_split, + ) + .unwrap(); + assert!(split_compaction_scratch > 0); + assert!(split_ipc_data_size + split_compaction_scratch > admission_capacity); + + let recorder = ReservationRecordingPusher::default(); + let recorded = Arc::clone(&recorder.0); + let pusher = BoundedAdmissionPusher { + recorder, + capacity: admission_capacity, + }; + let encoder = + ShuffleBlockWriter::try_new(input.schema().as_ref(), CompressionCodec::None) + .unwrap(); + let mut writer = RssPartitionWriter::try_new(pusher, encoder, 1, frame_limit).unwrap(); + let metrics = metrics(); + + writer + .finish_partition(0, &mut [Ok(input.clone())].into_iter(), &metrics) + .unwrap(); + writer.finish_all(&metrics).unwrap(); + + let recorded = recorded.lock().unwrap(); + assert!(recorded.active.is_none()); + assert!(recorded.frames.len() > 1); + assert!(recorded + .reservations + .iter() + .all(|reservation| *reservation <= admission_capacity)); + assert_eq!( + recorded.reservations.len(), + recorded.releases + recorded.frames.len() + ); + + let decoded = recorded + .frames + .iter() + .map(|frame| { + assert!(frame.len() <= frame_limit); + read_ipc_compressed(&frame[16..]).unwrap() + }) + .collect::>(); + assert_eq!( + arrow_select::concat::concat_batches(&input.schema(), &decoded).unwrap(), + input + ); + } + } + #[test] fn rss_partition_writer_counts_validity_offsets_and_view_backing_buffers() { let booleans: ArrayRef = Arc::new(BooleanArray::from(vec![true; 8_192])); diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java index d5dcdd79edc..19f87577a11 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java +++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java @@ -28,6 +28,7 @@ import java.util.ArrayDeque; import java.util.ArrayList; import java.util.IdentityHashMap; +import java.util.Map; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; @@ -55,6 +56,7 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus private final Method getPushState; private final Method setPushMetricsCallback; private final Class pushMetricsCallbackClass; + private final Field clientPushStates; private final Field inFlightRequestTracker; private final Field totalInFlightRequests; private final Field pushStateException; @@ -232,6 +234,7 @@ public CelebornShufflePartitionPusher( final Method getPushStateMethod; final Method setPushMetricsCallbackMethod; final Class pushMetricsCallbackClass; + final Field clientPushStatesField; final Field inFlightRequestTrackerField; final Field totalInFlightRequestsField; final Field pushStateExceptionField; @@ -240,6 +243,7 @@ public CelebornShufflePartitionPusher( setPushMetricsCallbackMethod = resolvePushMetricsCallback(getPushStateMethod.getReturnType()); pushMetricsCallbackClass = setPushMetricsCallbackMethod.getParameterTypes()[0]; pushMetricsCallbackClass.getMethod("incPushDataCount", long.class); + clientPushStatesField = resolveClientPushStates(shuffleClient.getClass()); inFlightRequestTrackerField = getPushStateMethod.getReturnType().getDeclaredField("inFlightRequestTracker"); totalInFlightRequestsField = @@ -251,6 +255,7 @@ public CelebornShufflePartitionPusher( if (pushStateExceptionField.getType() != AtomicReference.class) { throw new NoSuchFieldException("PushState.exception: AtomicReference"); } + clientPushStatesField.setAccessible(true); inFlightRequestTrackerField.setAccessible(true); totalInFlightRequestsField.setAccessible(true); pushStateExceptionField.setAccessible(true); @@ -273,6 +278,7 @@ public CelebornShufflePartitionPusher( this.getPushState = getPushStateMethod; this.setPushMetricsCallback = setPushMetricsCallbackMethod; this.pushMetricsCallbackClass = pushMetricsCallbackClass; + this.clientPushStates = clientPushStatesField; this.inFlightRequestTracker = inFlightRequestTrackerField; this.totalInFlightRequests = totalInFlightRequestsField; this.pushStateException = pushStateExceptionField; @@ -297,6 +303,22 @@ private static Method resolvePushMetricsCallback(Class pushStateClass) throw new NoSuchMethodException("PushState.setMetricsCallback"); } + private static Field resolveClientPushStates(Class clientClass) throws NoSuchFieldException { + for (Class current = clientClass; current != null; current = current.getSuperclass()) { + try { + Field field = current.getDeclaredField("pushStates"); + if (!Map.class.isAssignableFrom(field.getType()) + || Modifier.isStatic(field.getModifiers())) { + throw new NoSuchFieldException("ShuffleClientImpl.pushStates: instance Map"); + } + return field; + } catch (NoSuchFieldException ignored) { + // Test clients and application wrappers can declare the pinned field on a superclass. + } + } + throw new NoSuchFieldException("ShuffleClientImpl.pushStates: instance Map"); + } + private static Method resolveLifecycleMethod( Object shuffleClient, String name, Class... parameterTypes) { final Method method; @@ -406,12 +428,14 @@ public int pushPartitionData(int partitionId, byte[] bytes, int length) throws I } if (submitted && isAborted()) { - // Cleanup can remove the initial state while Celeborn resolves locations. Capture the - // recreated state's actual in-flight count before this submission becomes observable. - final Object resumedPushState = - getPushState.invoke(shuffleClient, shuffleId + "-" + mapId + "-" + encodedAttemptId); - observedPushState = observePushState(resumedPushState); - if (resumedPushState != pushState) { + // Cleanup removes the map entry, but Celeborn can still submit using the PushState it + // already captured. A side-effect-free lookup distinguishes that state from a genuine + // post-cleanup replacement without creating an empty state and losing the live request. + Object resumedPushState = + ((Map) clientPushStates.get(shuffleClient)) + .get(shuffleId + "-" + mapId + "-" + encodedAttemptId); + if (resumedPushState != null && resumedPushState != pushState) { + observedPushState = observePushState(resumedPushState); recoverUnobservedFailure(observedPushState); } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index 89784afbea2..8fc8eb48846 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -525,6 +525,30 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( } } + private def resetSparkCommitOwners(generation: PrepareCelebornShuffleGeneration): Unit = + outputCommitCoordinator.synchronized { + val stageStatesField = classOf[OutputCommitCoordinator].getDeclaredField("stageStates") + stageStatesField.setAccessible(true) + val stageStates = stageStatesField + .get(outputCommitCoordinator) + .asInstanceOf[mutable.Map[Int, AnyRef]] + val stageState = stageStates.getOrElse( + generation.stageId, + throw new IllegalStateException( + s"Spark output commit state is unavailable for stage ${generation.stageId}")) + val authorizedCommitters = stageState.getClass + .getMethod("authorizedCommitters") + .invoke(stageState) + .asInstanceOf[Array[AnyRef]] + require( + authorizedCommitters.length == generation.numMappers, + s"Spark output commit mapper count ${authorizedCommitters.length} does not match " + + s"Celeborn mapper count ${generation.numMappers}") + + // Ending and restarting the stage would also erase failures that Spark has already recorded. + java.util.Arrays.fill(authorizedCommitters, null) + } + def claimMapAttempt(claim: ClaimCelebornMapAttempt): CelebornMapAttemptClaim = synchronized { val previousGeneration = generations.get(claim.shuffleId) val stale = previousGeneration.exists { generation => @@ -598,18 +622,7 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( previous.stageAttempt == generation.stageAttempt => previous.celebornShuffleId == generation.celebornShuffleId case Some(previous) if previous.celebornShuffleId != generation.celebornShuffleId => - // Spark scopes these lifecycle methods private[scheduler] despite public JVM methods. - outputCommitCoordinator.synchronized { - outputCommitCoordinator.getClass - .getMethod("stageEnd", java.lang.Integer.TYPE) - .invoke(outputCommitCoordinator, Int.box(generation.stageId)) - outputCommitCoordinator.getClass - .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) - .invoke( - outputCommitCoordinator, - Int.box(generation.stageId), - Int.box(generation.numMappers - 1)) - } + resetSparkCommitOwners(generation) invalidateOwners(generation.shuffleId) invalidatedGenerations.remove(generation.shuffleId) generations.update(generation.shuffleId, generation) diff --git a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala index ac306ee9f8a..420a895f6d4 100644 --- a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala +++ b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala @@ -51,6 +51,10 @@ final class RecordingCelebornShuffleClient { @volatile var lastCleanup: (Int, Int, Int) = _ @volatile var pushStarted: CountDownLatch = _ @volatile var allowPush: CountDownLatch = _ + @volatile var pushStateCaptured: CountDownLatch = _ + @volatile var allowPushStateCapture: CountDownLatch = _ + @volatile var pushRegistered: CountDownLatch = _ + @volatile var allowPushRegistration: CountDownLatch = _ @volatile var pushCompletedBeforeReturn: CountDownLatch = _ @volatile var allowPushReturn: CountDownLatch = _ @volatile var mapperEndStarted: CountDownLatch = _ @@ -74,6 +78,7 @@ final class RecordingCelebornShuffleClient { @volatile var generationInvalidated = true @volatile var lastInvalidatedShuffle: (Int, Int, Long) = _ val pushCalls = new AtomicInteger() + val pushStateCreations = new AtomicInteger() val pushCompletionCalls = new AtomicInteger() val mapperEndCalls = new AtomicInteger() val cleanupCalls = new AtomicInteger() @@ -83,7 +88,12 @@ final class RecordingCelebornShuffleClient { def getPushState(mapKey: String): RecordingCelebornPushState = { observedMapKey = mapKey retainedPushState = true - pushStates.computeIfAbsent(mapKey, _ => new RecordingCelebornPushState(this)) + pushStates.computeIfAbsent( + mapKey, + _ => { + pushStateCreations.incrementAndGet() + new RecordingCelebornPushState(this) + }) } def completeNextPush(): Boolean = { @@ -161,7 +171,19 @@ final class RecordingCelebornShuffleClient { val pushState = getPushState(s"$shuffleId-$mapId-$attemptId") pushState.addPush() + if (pushStateCaptured != null) { + pushStateCaptured.countDown() + if (!allowPushStateCapture.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting for the captured test push state") + } + } pendingPushCompletions.add(pushState) + if (pushRegistered != null) { + pushRegistered.countDown() + if (!allowPushRegistration.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting for the registered test push") + } + } if (automaticallyCompleteLatestPush) { completeLatestPush() } else if (automaticallyCompletePushes) { @@ -1074,6 +1096,68 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { assert(client.cleanupCalls.get() == 2) } + Seq("connection creation", "transport registration").foreach { cancellationWindow => + test( + s"cancellation during $cancellationWindow retains the captured state's transport admission") { + val client = new RecordingCelebornShuffleClient + client.automaticallyCompletePushes = false + val pausedPush = new CountDownLatch(1) + val allowPush = new CountDownLatch(1) + if (cancellationWindow == "connection creation") { + client.pushStateCaptured = pausedPush + client.allowPushStateCapture = allowPush + } else { + client.pushRegistered = pausedPush + client.allowPushRegistration = allowPush + } + + val cancelled = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 112) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 112) + val cancelledFailure = new AtomicReference[Throwable]() + val replacementFailure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + val cancelledWorker = new Thread(() => { + try { + cancelled.reservePartitionData(96) + cancelled.pushPartitionData(0, frame, frame.length) + } catch { case failure: Throwable => cancelledFailure.set(failure) } + }) + cancelledWorker.start() + assert(pausedPush.await(5, TimeUnit.SECONDS)) + + cancelled.abort() + allowPush.countDown() + cancelledWorker.join(5000) + assert(!cancelledWorker.isAlive) + assert(cancelledFailure.get().isInstanceOf[IOException]) + assert(client.cleanupCalls.get() == 2) + + client.pushStateCaptured = null + client.pushRegistered = null + val replacementWorker = new Thread(() => { + try { + replacement.reservePartitionData(96) + replacement.pushPartitionData(0, frame, frame.length) + } catch { case failure: Throwable => replacementFailure.set(failure) } + }) + replacementWorker.start() + + Thread.sleep(100) + assert(replacementWorker.isAlive) + assert(client.pushCalls.get() == 1) + assert(client.pushStateCreations.get() == 1) + assert(client.pushCompletionCalls.get() == 0) + + assert(client.completeNextPush()) + replacementWorker.join(5000) + assert(!replacementWorker.isAlive) + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + } + test("a recreated cancelled push retains byte admission until its own transport completes") { val client = new RecordingCelebornShuffleClient client.automaticallyCompletePushes = false diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala index c37f575527e..184e01809d8 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala @@ -374,6 +374,83 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(!prepared.requiresGenerationResolution) } + test("a replacement generation preserves another partition's previously reported failure") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val stageId = 80 + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(1)) + + val originalGeneration = PrepareCelebornShuffleGeneration(12, 91, stageId, 0, 2) + assert(coordinator.prepareGeneration(originalGeneration)) + assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(12, stageId, 0, 0, 0)).authorized) + + // Partition zero fails while resolving its input, before it can ask Comet for a writer. + sparkCoordinator.getClass + .getMethod( + "taskCompleted", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + classOf[TaskEndReason]) + .invoke( + sparkCoordinator, + Int.box(stageId), + Int.box(1), + Int.box(0), + Int.box(0), + UnknownReason) + + // Partition one is the first replacement-stage mapper to prepare the new generation. + val firstReplacementClaim = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(12, stageId, 1, 1, 0)) + assert(firstReplacementClaim.authorized) + val replacementGeneration = PrepareCelebornShuffleGeneration(12, 92, stageId, 1, 2) + val preparedReplacement = coordinator.prepareGenerationAndClaim( + replacementGeneration.copy( + mapId = 1, + taskAttempt = 0, + claimEpoch = firstReplacementClaim.epoch, + claimAuthorized = firstReplacementClaim.authorized)) + assert(preparedReplacement.authorized) + + // Its peer's already-failed attempt must not become a phantom commit owner after the reset. + val retry = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(12, stageId, 1, 0, 1)) + assert(retry.authorized) + val preparedRetry = coordinator.prepareGenerationAndClaim( + replacementGeneration.copy( + mapId = 0, + taskAttempt = 1, + claimEpoch = retry.epoch, + claimAuthorized = retry.authorized)) + assert(preparedRetry.authorized) + assert(preparedRetry.epoch == preparedReplacement.epoch) + } + + test("a replacement generation cannot recreate a completed Spark stage") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + val stageId = 81 + sparkCoordinator.getClass + .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId), Int.box(0)) + + val originalGeneration = PrepareCelebornShuffleGeneration(13, 91, stageId, 0, 1) + assert(coordinator.prepareGeneration(originalGeneration)) + sparkCoordinator.getClass + .getMethod("stageEnd", java.lang.Integer.TYPE) + .invoke(sparkCoordinator, Int.box(stageId)) + + val failure = intercept[IllegalStateException] { + coordinator.prepareGeneration( + originalGeneration.copy(celebornShuffleId = 92, stageAttempt = 1)) + } + assert(failure.getMessage.contains(s"stage $stageId")) + assert(sparkCoordinator.isEmpty) + } + test("a replacement-stage original blocks speculation before its generation is resolved") { val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) From e62f4ca5e3622cda3998910b4e9316915dd3c5c6 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:19:44 +0000 Subject: [PATCH 08/12] feat: add native Celeborn shuffle reader --- .../CometBlockStoreShuffleReader.scala | 5 +- .../shuffle/CometCelebornShuffleManager.scala | 69 +- .../shuffle/CometCelebornShuffleReader.scala | 618 +++++++++++++ .../shuffle/CometShuffleReader.scala | 29 + .../shuffle/CometShuffledRowRDD.scala | 12 +- .../shuffle/RecordingCelebornRawClient.java | 241 +++++ .../exception/CelebornRuntimeException.scala | 23 + .../celeborn/CelebornShuffleHandle.scala | 30 + .../celeborn/CelebornShuffleReader.scala | 27 + .../CometCelebornShuffleReaderSuite.scala | 827 ++++++++++++++++++ 10 files changed, 1855 insertions(+), 26 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleReader.scala create mode 100644 spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java create mode 100644 spark/src/test/scala/org/apache/celeborn/common/exception/CelebornRuntimeException.scala create mode 100644 spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala create mode 100644 spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala index 3bedc0e0f4c..3048456ea78 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala @@ -26,7 +26,6 @@ import org.apache.spark.internal.{config, Logging} import org.apache.spark.io.CompressionCodec import org.apache.spark.serializer.SerializerManager import org.apache.spark.shuffle.BaseShuffleHandle -import org.apache.spark.shuffle.ShuffleReader import org.apache.spark.shuffle.ShuffleReadMetricsReporter import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.storage.BlockId @@ -51,7 +50,7 @@ class CometBlockStoreShuffleReader[K, C]( blockManager: BlockManager = SparkEnv.get.blockManager, mapOutputTracker: MapOutputTracker = SparkEnv.get.mapOutputTracker, shouldBatchFetch: Boolean = false) - extends ShuffleReader[K, C] + extends CometShuffleReader[K, C] with Logging { private val dep = handle.dependency.asInstanceOf[CometShuffleDependency[_, _, _]] @@ -157,7 +156,7 @@ class CometBlockStoreShuffleReader[K, C]( * Returns the raw concatenated InputStream of all shuffle blocks, bypassing the decode step. * Used by ShuffleScan direct read path. */ - def readAsRawStream(): InputStream = { + override def readAsRawStream(): InputStream = { val streams = fetchIterator.map(_._2) new java.io.SequenceInputStream(new java.util.Enumeration[InputStream] { override def hasMoreElements: Boolean = streams.hasNext diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index 8fc8eb48846..533ce59a7cc 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -39,8 +39,8 @@ import org.apache.comet.util.ClassLoaders * Lets Comet execution coexist with the application's existing Celeborn shuffle manager. * * Ordinary Spark shuffle dependencies are owned entirely by Celeborn. Native Comet map tasks can - * write directly to Celeborn, but query planning keeps native shuffle disabled until the matching - * remote reader is available. JVM Comet shuffle remains unsupported. + * write to and read from Celeborn, while query planning keeps native shuffle disabled until its + * remaining production-readiness work is complete. JVM Comet shuffle remains unsupported. * * Celeborn is loaded reflectively because its client is an optional, application-provided * dependency rather than part of Comet's compile-time or runtime distribution. @@ -48,12 +48,17 @@ import org.apache.comet.util.ClassLoaders class CometCelebornShuffleManager private[shuffle] ( conf: SparkConf, isDriver: Boolean, - backendFactory: (SparkConf, Boolean) => ShuffleManager) + backendFactory: (SparkConf, Boolean) => ShuffleManager, + readerApi: CelebornRawPartitionReader.Api = CelebornRawPartitionReader.reflectedApi) extends ShuffleManager { /** Constructor selected by Spark for driver and executor shuffle managers. */ def this(conf: SparkConf, isDriver: Boolean) = - this(conf, isDriver, CometCelebornShuffleManager.createBackend) + this( + conf, + isDriver, + CometCelebornShuffleManager.createBackend, + CelebornRawPartitionReader.reflectedApi) private val celebornManager = Option(backendFactory(conf, isDriver)).getOrElse { throw new IllegalStateException("Celeborn Spark shuffle manager factory returned null") @@ -176,19 +181,51 @@ class CometCelebornShuffleManager private[shuffle] ( endPartition: Int, context: TaskContext, metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { - if (nativeDependency(handle).nonEmpty) { - throw new UnsupportedOperationException( - "Celeborn-backed Comet shuffle cannot be enabled until its native reader is available") + nativeDependency(handle) match { + case Some(dependency) => + if (startMapIndex > endMapIndex) { + throw new UnsupportedOperationException( + "Celeborn physical-skew chunk reads are not supported by native Comet shuffle") + } + val backendReader = celebornManager.getReader[K, C]( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) + val rawReader = CelebornRawPartitionReader.fromBackendReader( + conf, + handle, + backendReader, + CelebornRawPartitionReader.ReadRange( + startMapIndex, + endMapIndex, + startPartition, + endPartition), + context, + metrics, + client => ownedNativeClients.put(client, java.lang.Boolean.TRUE), + (client, celebornShuffleId) => { + nativeShuffleClients + .computeIfAbsent(handle.shuffleId, _ => new ConcurrentHashMap[Int, AnyRef]()) + .put(celebornShuffleId, client) + }, + readerApi) + new CometCelebornShuffleReader[K, C](dependency, context, metrics, rawReader) + + case None => + rejectCometHandle(handle) + celebornManager.getReader( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) } - rejectCometHandle(handle) - celebornManager.getReader( - handle, - startMapIndex, - endMapIndex, - startPartition, - endPartition, - context, - metrics) } override def shuffleBlockResolver: ShuffleBlockResolver = diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala new file mode 100644 index 00000000000..ac3dcddfa05 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala @@ -0,0 +1,618 @@ +/* + * 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.spark.sql.comet.execution.shuffle + +import java.io.{FilterInputStream, InputStream, IOException} +import java.lang.reflect.{InvocationHandler, InvocationTargetException, Method, Proxy} +import java.util.{ArrayList, Map => JMap, Set => JSet} +import java.util.concurrent.{ConcurrentHashMap, TimeoutException, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean + +import scala.util.control.NonFatal + +import org.apache.spark.{InterruptibleIterator, SparkConf, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.shuffle.{FetchFailedException, ShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter} +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.CompletionIterator + +import org.apache.comet.{CometConf, Native} +import org.apache.comet.util.ClassLoaders +import org.apache.comet.vector.NativeUtil + +/** Reads native shuffle frames from Celeborn without applying Celeborn's row decompressor. */ +private[shuffle] final class CometCelebornShuffleReader[K, C]( + dependency: CometShuffleDependency[_, _, _], + context: TaskContext, + readMetrics: ShuffleReadMetricsReporter, + partitionReader: CelebornRawPartitionReader) + extends CometShuffleReader[K, C] { + + private val consumed = new AtomicBoolean(false) + private val metricsMerged = new AtomicBoolean(false) + + private def mergeReadMetrics(): Unit = { + if (metricsMerged.compareAndSet(false, true)) { + context.taskMetrics().mergeShuffleReadMetrics() + } + } + + override def readAsRawStream(): InputStream = { + if (!consumed.compareAndSet(false, true)) { + throw new IllegalStateException("A Celeborn shuffle reader can only be consumed once") + } + val input = new FilterInputStream(partitionReader.openPartitions()) { + override def read(): Int = { + val value = in.read() + if (value < 0) mergeReadMetrics() + value + } + + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = { + val count = in.read(buffer, offset, length) + if (count < 0) mergeReadMetrics() + count + } + + override def close(): Unit = { + try super.close() + finally mergeReadMetrics() + } + } + context.addTaskCompletionListener[Unit](_ => input.close()) + input + } + + override def read(): Iterator[Product2[K, C]] = { + if (dependency.aggregator.isDefined) { + throw new UnsupportedOperationException("aggregate not allowed") + } + if (dependency.keyOrdering.isDefined) { + throw new UnsupportedOperationException("order not allowed") + } + + val input = readAsRawStream() + val nativeUtil = new NativeUtil() + var decoder: NativeBatchDecoderIterator = null + context.addTaskCompletionListener[Unit] { _ => + try { + if (decoder != null) decoder.close() + else input.close() + } finally { + nativeUtil.close() + } + } + + try { + decoder = NativeBatchDecoderIterator( + input, + dependency.decodeTime, + new Native(), + nativeUtil, + CometConf.COMET_TRACING_ENABLED.get()) + } catch { + case NonFatal(failure) => + try input.close() + finally nativeUtil.close() + throw failure + } + + val rows = decoder.map { batch: ColumnarBatch => + readMetrics.incRecordsRead(batch.numRows()) + (0, batch) + } + val completed = CompletionIterator[(Int, ColumnarBatch), Iterator[(Int, ColumnarBatch)]]( + rows, + mergeReadMetrics()) + new InterruptibleIterator[(Int, ColumnarBatch)](context, completed) + .asInstanceOf[Iterator[Product2[K, C]]] + } +} + +/** + * Reflects the optional application-owned Celeborn client and keeps one reducer-file-group + * snapshot for all reducers in an AQE partition. + */ +private[shuffle] final class CelebornRawPartitionReader( + client: AnyRef, + sparkShuffleId: Int, + celebornShuffleId: Int, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + readMetrics: ShuffleReadMetricsReporter, + rpcRetryLimit: Int, + stageRerunEnabled: Boolean = true) + extends Logging { + + import CelebornRawPartitionReader._ + + require(startMapIndex >= 0, s"Invalid Celeborn start map index: $startMapIndex") + require( + startMapIndex <= endMapIndex, + "Celeborn physical-skew chunk reads are not supported by the native shuffle reader") + require( + startPartition >= 0 && startPartition <= endPartition, + s"Invalid Celeborn reducer range: [$startPartition, $endPartition)") + require(rpcRetryLimit >= 0, s"Invalid Celeborn RPC retry limit: $rpcRetryLimit") + + private val updateFileGroup = + client.getClass.getMethod("updateFileGroup", java.lang.Integer.TYPE, java.lang.Integer.TYPE) + private val readPartitionMethod = client.getClass.getMethods + .find { method => + val parameters = method.getParameterTypes + method.getName == "readPartition" && parameters.length == 16 && + parameters(0) == java.lang.Integer.TYPE && parameters(4) == java.lang.Long.TYPE && + parameters(15) == java.lang.Boolean.TYPE && parameters(14).isInterface && + classOf[InputStream].isAssignableFrom(method.getReturnType) + } + .getOrElse { + throw new IllegalStateException( + "The Celeborn client does not expose its raw 16-argument partition reader") + } + private val metricsCallback = + createMetricsCallback(readPartitionMethod.getParameterTypes.apply(14), readMetrics) + + private lazy val fileGroups: AnyRef = loadFileGroups() + + private def loadFileGroups(): AnyRef = { + val started = System.nanoTime() + var retries = 0 + try { + while (true) { + val stageEnded = + try { + invoke( + client.getClass.getMethod("isShuffleStageEnd", java.lang.Integer.TYPE), + client, + Int.box(celebornShuffleId)).asInstanceOf[Boolean] + } catch { + case NonFatal(failure) => + logInfo( + s"Could not check whether Celeborn shuffle $celebornShuffleId ended", + failure) + true + } + + try { + return Option( + invoke(updateFileGroup, client, Int.box(celebornShuffleId), Int.box(startPartition))) + .getOrElse { + throw new IllegalStateException( + "Celeborn returned a null reducer-file-group snapshot") + } + } catch { + case failure if isFileGroupTimeout(failure) && !stageEnded && retries < rpcRetryLimit => + retries += 1 + logInfo( + s"Retrying Celeborn reducer-file-group snapshot $celebornShuffleId " + + s"($retries/$rpcRetryLimit)", + failure) + case failure if isUnreportableFileGroupFailure(failure) => + throw failure + case NonFatal(failure) => + reportFetchFailure(startPartition, failure) + } + } + throw new IllegalStateException("The Celeborn file-group retry loop ended unexpectedly") + } finally { + readMetrics.incFetchWaitTime(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started)) + } + } + + private def snapshotField(snapshot: AnyRef, name: String): AnyRef = + snapshot.getClass.getField(name).get(snapshot).asInstanceOf[AnyRef] + + private def reducers: Iterator[(Int, ArrayList[AnyRef], AnyRef, Array[Int])] = { + val snapshot = fileGroups + val partitionGroups = Option(snapshotField(snapshot, "partitionGroups")) + .getOrElse { + throw new IllegalStateException("Celeborn returned null reducer partition groups") + } + .asInstanceOf[JMap[Integer, JSet[AnyRef]]] + val mapAttempts = Option(snapshotField(snapshot, "mapAttempts")) + .getOrElse { + throw new IllegalStateException("Celeborn returned null mapper-attempt metadata") + } + .asInstanceOf[Array[Int]] + val pushFailedBatches = snapshotField(snapshot, "pushFailedBatches") + + (startPartition until endPartition).iterator.flatMap { partition => + Option(partitionGroups.get(Int.box(partition))) + .filter(!_.isEmpty) + .map(locations => + (partition, new ArrayList[AnyRef](locations), pushFailedBatches, mapAttempts)) + } + } + + private def openPartition( + partition: Int, + locations: ArrayList[AnyRef], + pushFailedBatches: AnyRef, + mapAttempts: Array[Int]): InputStream = { + val attempt = encodeAttemptNumber(context.stageAttemptNumber(), context.attemptNumber()) + val stream = + try { + Option( + invoke( + readPartitionMethod, + client, + Int.box(celebornShuffleId), + Int.box(sparkShuffleId), + Int.box(partition), + Int.box(attempt), + Long.box(context.taskAttemptId()), + Int.box(startMapIndex), + Int.box(endMapIndex), + null, + locations, + null, + pushFailedBatches, + null, + null, + mapAttempts, + metricsCallback, + java.lang.Boolean.FALSE)) + .getOrElse { + throw new IOException(s"Celeborn returned a null stream for reducer $partition") + } + .asInstanceOf[InputStream] + } catch { + case NonFatal(failure) => reportFetchFailure(partition, failure) + } + + new FilterInputStream(stream) { + override def read(): Int = + try in.read() + catch { case failure: IOException => reportFetchFailure(partition, failure) } + + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = + try in.read(buffer, offset, length) + catch { case failure: IOException => reportFetchFailure(partition, failure) } + + override def skip(length: Long): Long = + try in.skip(length) + catch { case failure: IOException => reportFetchFailure(partition, failure) } + } + } + + private def reportFetchFailure(partition: Int, failure: Throwable): Nothing = { + if (failure.isInstanceOf[FetchFailedException] || context.isInterrupted() || + context.isCompleted()) { + throw failure + } + + if (stageRerunEnabled && invoke( + client.getClass.getMethod( + "reportShuffleFetchFailure", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Long.TYPE), + client, + Int.box(sparkShuffleId), + Int.box(celebornShuffleId), + Long.box(context.taskAttemptId())).asInstanceOf[Boolean]) { + throw new FetchFailedException( + null, + sparkShuffleId, + -1L, + -1, + partition, + s"Celeborn FetchFailure appShuffleId/shuffleId: $sparkShuffleId/$celebornShuffleId", + failure) + } + throw failure + } + + def openPartitions(): InputStream = { + new InputStream { + private lazy val partitions = reducers + private val stateLock = new Object + @volatile private var current: InputStream = _ + @volatile private var closed = false + @volatile private var exhausted = false + + private def nextStream(): Boolean = { + if (closed || exhausted) { + false + } else if (!partitions.hasNext) { + exhausted = true + false + } else { + val (partition, locations, pushFailedBatches, mapAttempts) = partitions.next() + if (closed) { + false + } else { + val opened = openPartition(partition, locations, pushFailedBatches, mapAttempts) + val accepted = stateLock.synchronized { + if (closed) false + else { + current = opened + true + } + } + if (!accepted) opened.close() + accepted + } + } + } + + private def releaseCurrent(stream: InputStream): Unit = { + val shouldClose = stateLock.synchronized { + if (current eq stream) { + current = null + true + } else false + } + if (shouldClose) stream.close() + } + + private def currentStream(): InputStream = { + val active = current + if (active != null) active + else if (nextStream()) current + else null + } + + override def read(): Int = { + if (closed) throw new IOException("Celeborn shuffle input stream is closed") + var active = currentStream() + while (active != null && !closed) { + val value = active.read() + if (value >= 0) return value + releaseCurrent(active) + active = currentStream() + } + -1 + } + + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = { + java.util.Objects.checkFromIndexSize(offset, length, buffer.length) + if (closed) throw new IOException("Celeborn shuffle input stream is closed") + if (length == 0) return 0 + var active = currentStream() + while (active != null && !closed) { + val count = active.read(buffer, offset, length) + if (count >= 0) return count + releaseCurrent(active) + active = currentStream() + } + -1 + } + + override def close(): Unit = { + val previous = stateLock.synchronized { + if (closed) null + else { + closed = true + exhausted = true + val active = current + current = null + active + } + } + if (previous != null) previous.close() + } + } + } +} + +private[shuffle] object CelebornRawPartitionReader { + + private val SPARK_UTILS_CLASS = "org.apache.spark.shuffle.celeborn.SparkUtils" + private val SHUFFLE_CLIENT_CLASS = "org.apache.celeborn.client.ShuffleClient" + private val SHUFFLE_READER_COMPANION_CLASS = + "org.apache.spark.shuffle.celeborn.CelebornShuffleReader$" + private val CELEBORN_RUNTIME_EXCEPTION = + "org.apache.celeborn.common.exception.CelebornRuntimeException" + private val MAX_STAGE_ATTEMPTS = 1 << 15 + private val MAX_TASK_ATTEMPTS = 1 << 16 + + private[shuffle] def encodeAttemptNumber(stageAttempt: Int, taskAttempt: Int): Int = { + require( + stageAttempt >= 0 && stageAttempt < MAX_STAGE_ATTEMPTS, + s"Celeborn stage attempt must be between 0 and ${MAX_STAGE_ATTEMPTS - 1}: " + + stageAttempt) + require( + taskAttempt >= 0 && taskAttempt < MAX_TASK_ATTEMPTS, + s"Celeborn task attempt must be between 0 and ${MAX_TASK_ATTEMPTS - 1}: " + + taskAttempt) + (stageAttempt << 16) | taskAttempt + } + + private def isCelebornRuntimeException(failure: Throwable): Boolean = { + var current: Class[_] = failure.getClass + while (current != null) { + if (current.getName == CELEBORN_RUNTIME_EXCEPTION) return true + current = current.getSuperclass + } + false + } + + private[shuffle] final case class ReadRange( + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int) + + private[shuffle] trait Api { + def rpcRetryLimit(conf: SparkConf): Int + + def shuffleId( + client: AnyRef, + handle: ShuffleHandle, + context: TaskContext, + isWriter: Boolean): Int + } + + private[shuffle] val reflectedApi: Api = new Api { + override def rpcRetryLimit(conf: SparkConf): Int = { + val sparkUtilsClass = ClassLoaders.loadClass(SPARK_UTILS_CLASS) + val celebornConf = + invoke(sparkUtilsClass.getMethod("fromSparkConf", classOf[SparkConf]), null, conf) + invoke(celebornConf.getClass.getMethod("clientRpcMaxRetries"), celebornConf) + .asInstanceOf[Int] + } + + override def shuffleId( + client: AnyRef, + handle: ShuffleHandle, + context: TaskContext, + isWriter: Boolean): Int = { + val sparkUtilsClass = ClassLoaders.loadClass(SPARK_UTILS_CLASS) + val shuffleClientClass = ClassLoaders.loadClass(SHUFFLE_CLIENT_CLASS) + val method = sparkUtilsClass.getMethod( + "celebornShuffleId", + shuffleClientClass, + handle.getClass, + classOf[TaskContext], + classOf[java.lang.Boolean]) + invoke(method, null, client, handle, context, Boolean.box(isWriter)).asInstanceOf[Int] + } + } + + def fromBackendReader( + conf: SparkConf, + handle: ShuffleHandle, + backendReader: ShuffleReader[_, _], + range: ReadRange, + context: TaskContext, + readMetrics: ShuffleReadMetricsReporter, + onClientAcquired: AnyRef => Unit, + onGenerationResolved: (AnyRef, Int) => Unit, + api: Api = reflectedApi): CelebornRawPartitionReader = { + try { + val client = invoke(backendReader.getClass.getMethod("shuffleClient"), backendReader) + onClientAcquired(client) + + // The delegated reader normally initializes this companion from read(). Its static setup + // registers the application client's broadcast reducer-file-group deserializer. + ClassLoaders.loadClass(SHUFFLE_READER_COMPANION_CLASS) + + val stageRerunEnabled = + invoke(handle.getClass.getMethod("stageRerunEnabled"), handle).asInstanceOf[Boolean] + if (!stageRerunEnabled) { + throw new IllegalStateException("Native Celeborn shuffle requires stage reruns") + } + + val retryLimit = api.rpcRetryLimit(conf) + val celebornShuffleId = { + val started = System.nanoTime() + try { + api.shuffleId(client, handle, context, isWriter = false) + } catch { + case failure if isCelebornRuntimeException(failure) => + throw new FetchFailedException( + null, + handle.shuffleId, + -1L, + -1, + range.startPartition, + s"Celeborn could not resolve shuffle generation ${handle.shuffleId}", + failure) + } finally { + readMetrics.incFetchWaitTime(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started)) + } + } + + onGenerationResolved(client, celebornShuffleId) + new CelebornRawPartitionReader( + client, + handle.shuffleId, + celebornShuffleId, + range.startMapIndex, + range.endMapIndex, + range.startPartition, + range.endPartition, + context, + readMetrics, + retryLimit, + stageRerunEnabled) + } catch { + case failure: ReflectiveOperationException => + throw new IllegalStateException( + "The Celeborn client does not expose the required native-shuffle reader API", + failure) + } + } + + private def invoke(method: Method, instance: AnyRef, arguments: AnyRef*): AnyRef = + try method.invoke(instance, arguments: _*) + catch { + case failure: InvocationTargetException => + throw Option(failure.getCause).getOrElse(failure) + } + + private def isFileGroupTimeout(failure: Throwable): Boolean = + Option(failure.getCause).exists(_.isInstanceOf[TimeoutException]) + + private def isUnreportableFileGroupFailure(failure: Throwable): Boolean = + Option(failure.getCause).exists { cause => + cause.isInstanceOf[InterruptedException] || cause.isInstanceOf[TimeoutException] || + cause.getClass.getName == "org.apache.celeborn.common.exception.CelebornBroadcastException" + } + + private def createMetricsCallback( + callbackClass: Class[_], + reporter: ShuffleReadMetricsReporter): AnyRef = { + val remoteWorkers = ConcurrentHashMap.newKeySet[String]() + val unavailableMetrics = ConcurrentHashMap.newKeySet[String]() + + def optionalMetric(name: String, value: Long): Unit = { + if (!unavailableMetrics.contains(name)) { + try { + val method = reporter.getClass.getMethod(name, java.lang.Long.TYPE) + method.setAccessible(true) + invoke(method, reporter, Long.box(value)) + } catch { + case NonFatal(_) => unavailableMetrics.add(name) + } + } + () + } + + val callback = new InvocationHandler { + override def invoke(proxy: AnyRef, method: Method, arguments: Array[AnyRef]): AnyRef = { + method.getName match { + case "incBytesRead" => + reporter.incRemoteBytesRead(arguments(0).asInstanceOf[Long]) + reporter.incRemoteBlocksFetched(1) + case "incReadTime" => reporter.incFetchWaitTime(arguments(0).asInstanceOf[Long]) + case "recordRemoteReadWorker" => + val worker = arguments(0).asInstanceOf[String] + if (worker != null && remoteWorkers.add(worker)) { + optionalMetric("incCelebornDistinctRemoteWorkersRead", 1) + } + case "hashCode" => return Int.box(System.identityHashCode(proxy)) + case "equals" => return Boolean.box(proxy eq arguments(0)) + case "toString" => return "CometCelebornShuffleMetricsCallback" + case name if name.startsWith("inc") => + optionalMetric("incCeleborn" + name.substring(3), arguments(0).asInstanceOf[Long]) + case _ => + } + null + } + } + + Proxy.newProxyInstance(callbackClass.getClassLoader, Array(callbackClass), callback) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleReader.scala new file mode 100644 index 00000000000..325c47b6d79 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleReader.scala @@ -0,0 +1,29 @@ +/* + * 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.spark.sql.comet.execution.shuffle + +import java.io.InputStream + +import org.apache.spark.shuffle.ShuffleReader + +/** The local and remote shuffle readers support the same decoded and native consumption paths. */ +private[shuffle] trait CometShuffleReader[K, C] extends ShuffleReader[K, C] { + def readAsRawStream(): InputStream +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala index 7604910b06c..9b3bb9b735d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala @@ -92,9 +92,7 @@ class CometShuffledBatchRDD( } } - private def createReader( - split: Partition, - context: TaskContext): CometBlockStoreShuffleReader[_, _] = { + private def createReader(split: Partition, context: TaskContext): CometShuffleReader[_, _] = { val tempMetrics = context.taskMetrics().createTempShuffleReadMetrics() // `SQLShuffleReadMetricsReporter` will update its own metrics for SQL exchange operator, // as well as the `tempMetrics` for basic shuffle metrics. @@ -108,7 +106,7 @@ class CometShuffledBatchRDD( endReducerIndex, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] + .asInstanceOf[CometShuffleReader[_, _]] case PartialReducerPartitionSpec(reducerIndex, startMapIndex, endMapIndex, _) => SparkEnv.get.shuffleManager @@ -120,7 +118,7 @@ class CometShuffledBatchRDD( reducerIndex + 1, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] + .asInstanceOf[CometShuffleReader[_, _]] case PartialMapperPartitionSpec(mapIndex, startReducerIndex, endReducerIndex) => SparkEnv.get.shuffleManager @@ -132,7 +130,7 @@ class CometShuffledBatchRDD( endReducerIndex, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] + .asInstanceOf[CometShuffleReader[_, _]] case CoalescedMapperPartitionSpec(startMapIndex, endMapIndex, numReducers) => SparkEnv.get.shuffleManager @@ -144,7 +142,7 @@ class CometShuffledBatchRDD( numReducers, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] + .asInstanceOf[CometShuffleReader[_, _]] } } diff --git a/spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java b/spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java new file mode 100644 index 00000000000..9833c6efbe8 --- /dev/null +++ b/spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java @@ -0,0 +1,241 @@ +/* + * 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.spark.sql.comet.execution.shuffle; + +import java.io.IOException; +import java.io.InputStream; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; + +import scala.Product2; +import scala.collection.Iterator; + +import org.apache.spark.shuffle.ShuffleReader; + +/** Exposes the pinned Celeborn read contract without making Celeborn a Comet test dependency. */ +public final class RecordingCelebornRawClient { + + private static final AtomicBoolean BROADCAST_DECODER_REGISTERED = new AtomicBoolean(false); + + public static void registerBroadcastDecoder() { + BROADCAST_DECODER_REGISTERED.set(true); + } + + public static boolean broadcastDecoderRegistered() { + return BROADCAST_DECODER_REGISTERED.get(); + } + + public interface MetricsCallback { + void incBytesRead(long bytes); + + void incReadTime(long time); + + default void incRemoteReadRetryCount(long count) {} + + default void recordRemoteReadWorker(String worker) {} + } + + public interface OptionalMetricsReporter { + void incCelebornRemoteReadRetryCount(long count); + } + + public static final class BackendReader implements ShuffleReader { + private final RecordingCelebornRawClient client; + + public BackendReader(RecordingCelebornRawClient client) { + this.client = client; + } + + public RecordingCelebornRawClient shuffleClient() { + return client; + } + + @Override + public Iterator> read() { + throw new UnsupportedOperationException("The delegated row reader must never be consumed"); + } + } + + public static final class ReduceFileGroups { + public Map> partitionGroups = new HashMap<>(); + public Map pushFailedBatches = new HashMap<>(); + public int[] mapAttempts = new int[] {0}; + } + + public static final class ReadRequest { + public final int shuffleId; + public final int appShuffleId; + public final int partitionId; + public final int attemptNumber; + public final long taskId; + public final int startMapIndex; + public final int endMapIndex; + public final Object exceptionMaker; + public final ArrayList locations; + public final ArrayList streamHandlers; + public final Map pushFailedBatches; + public final Map chunksRange; + public final Map coalescedPartitionInfos; + public final int[] mapAttempts; + public final MetricsCallback metricsCallback; + public final boolean needDecompress; + + ReadRequest( + int shuffleId, + int appShuffleId, + int partitionId, + int attemptNumber, + long taskId, + int startMapIndex, + int endMapIndex, + Object exceptionMaker, + ArrayList locations, + ArrayList streamHandlers, + Map pushFailedBatches, + Map chunksRange, + Map coalescedPartitionInfos, + int[] mapAttempts, + MetricsCallback metricsCallback, + boolean needDecompress) { + this.shuffleId = shuffleId; + this.appShuffleId = appShuffleId; + this.partitionId = partitionId; + this.attemptNumber = attemptNumber; + this.taskId = taskId; + this.startMapIndex = startMapIndex; + this.endMapIndex = endMapIndex; + this.exceptionMaker = exceptionMaker; + this.locations = locations; + this.streamHandlers = streamHandlers; + this.pushFailedBatches = pushFailedBatches; + this.chunksRange = chunksRange; + this.coalescedPartitionInfos = coalescedPartitionInfos; + this.mapAttempts = mapAttempts; + this.metricsCallback = metricsCallback; + this.needDecompress = needDecompress; + } + } + + public final ReduceFileGroups fileGroups = new ReduceFileGroups(); + public final Map streams = new HashMap<>(); + public final List requests = new ArrayList<>(); + public int updateFileGroupCalls; + public int stageEndChecks; + public int failureReports; + public int cleanupCalls; + public int timeoutFailures; + public boolean stageEnded = true; + public boolean invalidateOnFetchFailure = true; + public boolean requiresBroadcastDecoder; + public CountDownLatch readPartitionStarted; + public CountDownLatch allowReadPartition; + public IOException updateFileGroupFailure; + public IOException readPartitionFailure; + + public boolean isShuffleStageEnd(int shuffleId) { + stageEndChecks += 1; + return stageEnded; + } + + public ReduceFileGroups updateFileGroup(int shuffleId, int partitionId) throws IOException { + updateFileGroupCalls += 1; + if (requiresBroadcastDecoder && !broadcastDecoderRegistered()) { + throw new IOException("Celeborn's broadcast reducer-file-group decoder was not registered"); + } + if (timeoutFailures > 0) { + timeoutFailures -= 1; + throw new IOException("reducer-file-group RPC timed out", new TimeoutException()); + } + if (updateFileGroupFailure != null) { + throw updateFileGroupFailure; + } + return fileGroups; + } + + public InputStream readPartition( + int shuffleId, + int appShuffleId, + int partitionId, + int attemptNumber, + long taskId, + int startMapIndex, + int endMapIndex, + Object exceptionMaker, + ArrayList locations, + ArrayList streamHandlers, + Map pushFailedBatches, + Map chunksRange, + Map coalescedPartitionInfos, + int[] mapAttempts, + MetricsCallback metricsCallback, + boolean needDecompress) + throws IOException { + requests.add( + new ReadRequest( + shuffleId, + appShuffleId, + partitionId, + attemptNumber, + taskId, + startMapIndex, + endMapIndex, + exceptionMaker, + locations, + streamHandlers, + pushFailedBatches, + chunksRange, + coalescedPartitionInfos, + mapAttempts, + metricsCallback, + needDecompress)); + if (readPartitionFailure != null) { + throw readPartitionFailure; + } + if (readPartitionStarted != null) { + readPartitionStarted.countDown(); + try { + if (!allowReadPartition.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting to open the test reducer stream"); + } + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new IOException("interrupted while opening the test reducer stream", failure); + } + } + return streams.get(partitionId); + } + + public boolean reportShuffleFetchFailure(int appShuffleId, int shuffleId, long taskId) { + failureReports += 1; + return invalidateOnFetchFailure; + } + + public boolean cleanupShuffle(int shuffleId) { + cleanupCalls += 1; + return true; + } +} diff --git a/spark/src/test/scala/org/apache/celeborn/common/exception/CelebornRuntimeException.scala b/spark/src/test/scala/org/apache/celeborn/common/exception/CelebornRuntimeException.scala new file mode 100644 index 00000000000..3c69efe3c72 --- /dev/null +++ b/spark/src/test/scala/org/apache/celeborn/common/exception/CelebornRuntimeException.scala @@ -0,0 +1,23 @@ +/* + * 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.celeborn.common.exception + +/** Exact-name fixture for the optional Celeborn generation-failure contract. */ +class CelebornRuntimeException(message: String) extends RuntimeException(message) diff --git a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala b/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala new file mode 100644 index 00000000000..6104c320b93 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala @@ -0,0 +1,30 @@ +/* + * 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.spark.shuffle.celeborn + +import org.apache.spark.ShuffleDependency +import org.apache.spark.shuffle.BaseShuffleHandle + +/** Exact-name test handle for the optional Celeborn manager's reflective reader integration. */ +class CelebornShuffleHandle[K, V, C]( + shuffleId: Int, + dependency: ShuffleDependency[K, V, C], + val stageRerunEnabled: Boolean = true) + extends BaseShuffleHandle[K, V, C](shuffleId, dependency) diff --git a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala b/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala new file mode 100644 index 00000000000..bea907ebae5 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala @@ -0,0 +1,27 @@ +/* + * 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.spark.shuffle.celeborn + +import org.apache.spark.sql.comet.execution.shuffle.RecordingCelebornRawClient + +/** Matches the pinned Celeborn companion's broadcast decoder-registration side effect. */ +object CelebornShuffleReader { + RecordingCelebornRawClient.registerBroadcastDecoder() +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala new file mode 100644 index 00000000000..fcd93e53e5a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -0,0 +1,827 @@ +/* + * 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.spark.sql.comet.execution.shuffle + +import java.io.{ByteArrayInputStream, InputStream, IOException} +import java.lang.reflect.{InvocationHandler, Method, Proxy} +import java.nio.{ByteBuffer, ByteOrder} +import java.nio.file.Files +import java.util.{LinkedHashSet, Set => JSet} +import java.util.concurrent.{CountDownLatch, TimeoutException, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{HashPartitioner, ShuffleDependency, SparkConf, SparkEnv, TaskContext} +import org.apache.spark.shuffle.{FetchFailedException, IndexShuffleBlockResolver, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.{CometConf, CometShuffleBlockIterator} + +class CometCelebornShuffleReaderSuite extends CometTestBase { + + import testImplicits._ + + private final class TrackingInputStream(bytes: Array[Byte]) + extends ByteArrayInputStream(bytes) { + var closeCalls = 0 + + override def close(): Unit = { + closeCalls += 1 + super.close() + } + } + + private final class RecordingReaderBackend(client: RecordingCelebornRawClient) + extends ShuffleManager { + var readRange: Option[(Int, Int, Int, Int)] = None + var unregistered: Option[Int] = None + + override def registerShuffle[K, V, C]( + shuffleId: Int, + dependency: ShuffleDependency[K, V, C]): ShuffleHandle = + throw new UnsupportedOperationException("This fixture only reads existing native shuffles") + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = + throw new UnsupportedOperationException("This fixture only reads existing native shuffles") + + override def getReader[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + readRange = Some((startMapIndex, endMapIndex, startPartition, endPartition)) + new RecordingCelebornRawClient.BackendReader(client).asInstanceOf[ShuffleReader[K, C]] + } + + override def shuffleBlockResolver: ShuffleBlockResolver = null + + override def unregisterShuffle(shuffleId: Int): Boolean = { + unregistered = Some(shuffleId) + true + } + + override def stop(): Unit = () + } + + private final class RecordingReaderApi extends CelebornRawPartitionReader.Api { + var resolvedClient: AnyRef = _ + var resolvedHandle: ShuffleHandle = _ + var resolvedContext: TaskContext = _ + var resolvedAsWriter: Option[Boolean] = None + var retryLimitFailure: Throwable = _ + var generationFailure: Throwable = _ + + override def rpcRetryLimit(conf: SparkConf): Int = { + if (retryLimitFailure != null) throw retryLimitFailure + 2 + } + + override def shuffleId( + client: AnyRef, + handle: ShuffleHandle, + context: TaskContext, + isWriter: Boolean): Int = { + resolvedClient = client + resolvedHandle = handle + resolvedContext = context + resolvedAsWriter = Some(isWriter) + if (generationFailure != null) throw generationFailure + 91 + } + } + + private def location(values: String*): JSet[Object] = { + val locations = new LinkedHashSet[Object]() + values.foreach(locations.add) + locations + } + + private def rawReader( + client: RecordingCelebornRawClient, + context: TaskContext, + startMap: Int = 0, + endMap: Int = Int.MaxValue, + startPartition: Int = 0, + endPartition: Int = 3, + retries: Int = 2): CelebornRawPartitionReader = + new CelebornRawPartitionReader( + client, + sparkShuffleId = 17, + celebornShuffleId = 91, + startMap, + endMap, + startPartition, + endPartition, + context, + context.taskMetrics().createTempShuffleReadMetrics(), + retries) + + private def simpleDependency( + partitions: Int = 3): CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch] = + new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( + spark.sparkContext.emptyRDD[(Int, ColumnarBatch)], + new HashPartitioner(partitions), + decodeTime = SQLMetrics.createMetric(spark.sparkContext, "Celeborn decode time")) + + private def ownedClients(manager: CometCelebornShuffleManager) + : java.util.concurrent.ConcurrentHashMap[AnyRef, java.lang.Boolean] = { + val ownership = classOf[CometCelebornShuffleManager].getDeclaredField("ownedNativeClients") + ownership.setAccessible(true) + ownership + .get(manager) + .asInstanceOf[java.util.concurrent.ConcurrentHashMap[AnyRef, java.lang.Boolean]] + } + + private def completionListenerCount(context: TaskContext): Int = { + val listeners = context.getClass.getDeclaredField("onCompleteCallbacks") + listeners.setAccessible(true) + listeners.get(context).asInstanceOf[java.util.Stack[_]].size() + } + + private def reader( + dependency: CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch], + client: RecordingCelebornRawClient, + context: TaskContext): CometCelebornShuffleReader[Int, ColumnarBatch] = { + val metrics = context.taskMetrics().createTempShuffleReadMetrics() + val partitions = new CelebornRawPartitionReader( + client, + dependency.shuffleId, + 91, + 0, + Int.MaxValue, + 0, + dependency.partitioner.numPartitions, + context, + metrics, + 2) + new CometCelebornShuffleReader[Int, ColumnarBatch](dependency, context, metrics, partitions) + } + + private def withNativeFrame( + run: (CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch], Array[Byte]) => Unit) + : Unit = { + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + "spark.sql.adaptive.enabled" -> "false") { + withParquetTable((0 until 12).map(i => (i, s"row-$i")), "celeborn_reader_rows") { + val dependency = sql("SELECT * FROM celeborn_reader_rows") + .repartition(3, $"_1") + .queryExecution + .executedPlan + .collectFirst { case exchange: CometShuffleExchangeExec => + exchange.shuffleDependency + .asInstanceOf[CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]] + } + .getOrElse(fail("Expected a native Comet shuffle exchange")) + + val frames = spark.sparkContext + .runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds) + writer.write(inputs) + writer.stop(success = true) + val dataFile = SparkEnv.get.shuffleManager.shuffleBlockResolver + .asInstanceOf[IndexShuffleBlockResolver] + .getDataFile(dependency.shuffleId, context.taskAttemptId()) + val bytes = Files.readAllBytes(dataFile.toPath) + var offset = 0 + writer.getPartitionLengths().flatMap { length => + val end = offset + length.toInt + val frame = bytes.slice(offset, end) + offset = end + if (frame.nonEmpty) Some(frame) else None + } + }) + .flatten + + assert(frames.nonEmpty) + run(dependency, frames.head) + } + } + } + + test("fresh Celeborn readers initialize the broadcast reducer-file-group decoder") { + assert(!RecordingCelebornRawClient.broadcastDecoderRegistered()) + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + client.requiresBroadcastDecoder = true + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, new ByteArrayInputStream(Array[Byte](9))) + val backend = new RecordingReaderBackend(client) + val api = new RecordingReaderApi + val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) + val manager = + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + + val remote = manager + .getReader[Int, ColumnarBatch]( + handle, + 0, + Int.MaxValue, + 0, + 1, + context, + context.taskMetrics.createTempShuffleReadMetrics()) + .asInstanceOf[CometShuffleReader[Int, ColumnarBatch]] + + assert(RecordingCelebornRawClient.broadcastDecoderRegistered()) + assert(remote.readAsRawStream().readAllBytes().toSeq == Seq[Byte](9)) + assert(client.updateFileGroupCalls == 1) + context.markTaskCompleted(None) + } + + test("the manager routes native Celeborn handles through its reflected raw reader") { + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + client.fileGroups.partitionGroups.put(1, location("worker")) + client.streams.put(1, new ByteArrayInputStream(Array[Byte](3, 4))) + val backend = new RecordingReaderBackend(client) + val api = new RecordingReaderApi + val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) + val manager = + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + + val remote = manager + .getReader[Int, ColumnarBatch]( + handle, + 2, + 6, + 1, + 3, + context, + context.taskMetrics.createTempShuffleReadMetrics()) + .asInstanceOf[CometShuffleReader[Int, ColumnarBatch]] + + assert(backend.readRange.contains((2, 6, 1, 3))) + assert(api.resolvedClient eq client) + assert(api.resolvedHandle eq handle) + assert(api.resolvedContext eq context) + assert(api.resolvedAsWriter.contains(false)) + assert(remote.readAsRawStream().readAllBytes().toSeq == Seq[Byte](3, 4)) + assert(client.requests.size() == 1) + assert(!client.requests.get(0).needDecompress) + assert(client.requests.get(0).startMapIndex == 2) + assert(client.requests.get(0).endMapIndex == 6) + + assert(ownedClients(manager).containsKey(client)) + assert(manager.unregisterShuffle(17)) + assert(backend.unregistered.contains(17)) + assert(client.cleanupCalls == 1) + context.markTaskCompleted(None) + } + + test("the manager fails closed on disabled stage reruns and unavailable optional reader APIs") { + Seq(false, true).foreach { unavailableApi => + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + val backend = new RecordingReaderBackend(client) + val api = new RecordingReaderApi + if (unavailableApi) api.generationFailure = new NoSuchMethodException("celebornShuffleId") + val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle( + 17, + dependency, + stageRerunEnabled = unavailableApi) + val manager = + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + + val failure = intercept[IllegalStateException] { + manager.getReader[Int, ColumnarBatch]( + handle, + 0, + Int.MaxValue, + 0, + 1, + context, + context.taskMetrics.createTempShuffleReadMetrics()) + } + + if (unavailableApi) { + assert(failure.getMessage.contains("reader API")) + assert(failure.getCause.isInstanceOf[NoSuchMethodException]) + } else { + assert(failure.getMessage.contains("stage reruns")) + assert(api.resolvedAsWriter.isEmpty) + } + assert(ownedClients(manager).containsKey(client)) + assert(client.failureReports == 0) + } + } + + test("reader clients remain owned when Celeborn retry configuration cannot be loaded") { + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + val backend = new RecordingReaderBackend(client) + val api = new RecordingReaderApi + val expected = new IllegalStateException("Celeborn retry settings are unavailable") + api.retryLimitFailure = expected + val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) + val manager = + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + + val failure = intercept[IllegalStateException] { + manager.getReader[Int, ColumnarBatch]( + handle, + 0, + Int.MaxValue, + 0, + 1, + context, + context.taskMetrics.createTempShuffleReadMetrics()) + } + + assert(failure eq expected) + assert(ownedClients(manager).containsKey(client)) + assert(api.resolvedAsWriter.isEmpty) + } + + test("a genuine reducer-generation resolution failure becomes a Spark fetch failure") { + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + val backend = new RecordingReaderBackend(client) + val api = new RecordingReaderApi + val expected = + new org.apache.celeborn.common.exception.CelebornRuntimeException( + "the reducer generation expired") + api.generationFailure = expected + val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) + val manager = + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + + val failure = intercept[FetchFailedException] { + manager.getReader[Int, ColumnarBatch]( + handle, + 0, + Int.MaxValue, + 0, + 1, + context, + context.taskMetrics.createTempShuffleReadMetrics()) + } + + assert(failure.getCause eq expected) + assert(api.resolvedAsWriter.contains(false)) + assert(ownedClients(manager).containsKey(client)) + } + + test("reader adapter bugs are never misclassified as retryable generation failures") { + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + val backend = new RecordingReaderBackend(client) + val api = new RecordingReaderApi + val expected = new ClassCastException("the optional client exposed an incompatible API") + api.generationFailure = expected + val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) + val manager = + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + + val failure = intercept[ClassCastException] { + manager.getReader[Int, ColumnarBatch]( + handle, + 0, + Int.MaxValue, + 0, + 1, + context, + context.taskMetrics.createTempShuffleReadMetrics()) + } + + assert(failure eq expected) + assert(ownedClients(manager).containsKey(client)) + assert(client.failureReports == 0) + } + + test("raw fetch preserves generation, reducer order, map ranges, and one file-group snapshot") { + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + val first = new TrackingInputStream(Array[Byte](1, 2)) + val third = new TrackingInputStream(Array[Byte](3, 4, 5)) + client.fileGroups.partitionGroups.put(1, location("worker-b", "worker-c")) + client.fileGroups.partitionGroups.put(2, location()) + client.fileGroups.partitionGroups.put(3, location("worker-a")) + client.fileGroups.mapAttempts = Array(3, 1, 4) + client.fileGroups.pushFailedBatches.put("worker-b", "replayed-batch") + client.streams.put(1, first) + client.streams.put(3, third) + + val input = rawReader(client, context, 4, 9, 1, 4).openPartitions() + assert(client.updateFileGroupCalls == 0) + assert(input.readAllBytes().toSeq == Seq[Byte](1, 2, 3, 4, 5)) + assert(client.updateFileGroupCalls == 1) + assert(client.requests.asScala.map(_.partitionId).toSeq == Seq(1, 3)) + + client.requests.asScala.foreach { request => + assert(request.shuffleId == 91) + assert(request.appShuffleId == 17) + assert(request.startMapIndex == 4) + assert(request.endMapIndex == 9) + assert(request.taskId == context.taskAttemptId) + assert(!request.needDecompress) + assert(request.exceptionMaker == null) + assert(request.streamHandlers == null) + assert(request.chunksRange == null) + assert(request.coalescedPartitionInfos == null) + assert(request.pushFailedBatches eq client.fileGroups.pushFailedBatches) + assert(request.mapAttempts eq client.fileGroups.mapAttempts) + } + assert(client.requests.get(0).locations.asScala.toSeq == Seq("worker-b", "worker-c")) + assert(first.closeCalls == 1) + assert(third.closeCalls == 1) + input.close() + context.markTaskCompleted(None) + } + + test("raw fetch forwards Celeborn byte, block, and wait metrics") { + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, new ByteArrayInputStream(Array[Byte](7))) + val input = rawReader(client, context).openPartitions() + assert(input.read() == 7) + + val callback = client.requests.get(0).metricsCallback + callback.incBytesRead(29) + callback.incBytesRead(11) + callback.incReadTime(13) + callback.incRemoteReadRetryCount(1) + callback.recordRemoteReadWorker("worker:9097") + callback.recordRemoteReadWorker("worker:9097") + context.taskMetrics.mergeShuffleReadMetrics() + + assert(context.taskMetrics.shuffleReadMetrics.remoteBytesRead == 40) + assert(context.taskMetrics.shuffleReadMetrics.remoteBlocksFetched == 2) + assert(context.taskMetrics.shuffleReadMetrics.fetchWaitTime >= 13) + input.close() + context.markTaskCompleted(None) + } + + test("native raw consumption merges task shuffle metrics at EOF and close only once") { + Seq(false, true).foreach { closeEarly => + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, new ByteArrayInputStream(Array[Byte](7, 8))) + val input = reader(simpleDependency(), client, context).readAsRawStream() + assert(input.read() == 7) + client.requests.get(0).metricsCallback.incBytesRead(19) + + if (closeEarly) input.close() + else assert(input.readAllBytes().toSeq == Seq[Byte](8)) + input.close() + context.markTaskCompleted(None) + + assert(context.taskMetrics.shuffleReadMetrics.remoteBytesRead == 19) + assert(context.taskMetrics.shuffleReadMetrics.remoteBlocksFetched == 1) + } + } + + test("broken optional Celeborn metrics are disabled without failing remote reads") { + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, new ByteArrayInputStream(Array[Byte](7))) + val underlying = context.taskMetrics.createTempShuffleReadMetrics() + var optionalUpdates = 0 + val reporter = Proxy + .newProxyInstance( + getClass.getClassLoader, + Array( + classOf[ShuffleReadMetricsReporter], + classOf[RecordingCelebornRawClient.OptionalMetricsReporter]), + new InvocationHandler { + override def invoke(proxy: AnyRef, method: Method, arguments: Array[AnyRef]): AnyRef = { + if (method.getName == "incCelebornRemoteReadRetryCount") { + optionalUpdates += 1 + throw new IllegalStateException("an optional reporter is unavailable") + } + method.invoke(underlying, arguments: _*) + } + }) + .asInstanceOf[ShuffleReadMetricsReporter] + val partitions = + new CelebornRawPartitionReader(client, 17, 91, 0, Int.MaxValue, 0, 1, context, reporter, 2) + + val input = partitions.openPartitions() + assert(input.read() == 7) + client.requests.get(0).metricsCallback.incRemoteReadRetryCount(1) + client.requests.get(0).metricsCallback.incRemoteReadRetryCount(1) + + assert(optionalUpdates == 1) + input.close() + context.markTaskCompleted(None) + } + + test("coalesced reducer reads retain only one task-completion stream listener") { + val reducers = 256 + val context = TaskContext.empty() + val dependency = simpleDependency(reducers) + val client = new RecordingCelebornRawClient + val streams = (0 until reducers).map { partition => + val stream = new TrackingInputStream(Array(partition.toByte)) + client.fileGroups.partitionGroups.put(partition, location(s"worker-$partition")) + client.streams.put(partition, stream) + stream + } + + val input = reader(dependency, client, context).readAsRawStream() + assert(completionListenerCount(context) == 1) + assert(input.readAllBytes().length == reducers) + assert(client.requests.size() == reducers) + assert(completionListenerCount(context) == 1) + assert(streams.forall(_.closeCalls == 1)) + + context.markTaskCompleted(None) + assert(streams.forall(_.closeCalls == 1)) + } + + test("closing an unread or partially consumed stream never opens another reducer") { + val untouchedContext = TaskContext.empty() + val untouchedClient = new RecordingCelebornRawClient + rawReader(untouchedClient, untouchedContext).openPartitions().close() + assert(untouchedClient.updateFileGroupCalls == 0) + + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + val first = new TrackingInputStream(Array[Byte](1, 2)) + val second = new TrackingInputStream(Array[Byte](3)) + client.fileGroups.partitionGroups.put(0, location("worker-a")) + client.fileGroups.partitionGroups.put(1, location("worker-b")) + client.streams.put(0, first) + client.streams.put(1, second) + + val input = rawReader(client, context).openPartitions() + assert(input.read() == 1) + input.close() + input.close() + + assert(first.closeCalls == 1) + assert(second.closeCalls == 0) + assert(client.requests.size() == 1) + context.markTaskCompleted(None) + assert(first.closeCalls == 1) + assert(second.closeCalls == 0) + } + + test("task completion closes an opened reducer stream") { + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + val stream = new TrackingInputStream(Array[Byte](1, 2)) + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, stream) + + val input = reader(dependency, client, context).readAsRawStream() + assert(input.read() == 1) + context.markTaskCompleted(None) + + assert(stream.closeCalls == 1) + } + + test("task completion racing reducer-stream creation closes the unpublished stream") { + val context = TaskContext.empty() + val dependency = simpleDependency() + val client = new RecordingCelebornRawClient + val stream = new TrackingInputStream(Array[Byte](1)) + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, stream) + client.readPartitionStarted = new CountDownLatch(1) + client.allowReadPartition = new CountDownLatch(1) + val input = reader(dependency, client, context).readAsRawStream() + val result = new AtomicInteger(Int.MinValue) + val failure = new AtomicReference[Throwable]() + val worker = new Thread(() => { + try result.set(input.read()) + catch { case caught: Throwable => failure.set(caught) } + }) + + worker.start() + try { + assert(client.readPartitionStarted.await(5, TimeUnit.SECONDS)) + context.markTaskCompleted(None) + } finally { + client.allowReadPartition.countDown() + worker.join(5000) + } + + assert(!worker.isAlive) + assert(failure.get() == null) + assert(result.get() == -1) + assert(stream.closeCalls == 1) + assert(client.requests.size() == 1) + assert(client.failureReports == 0) + } + + test("file-group RPC timeouts retry only while the map stage is still running") { + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + client.stageEnded = false + client.timeoutFailures = 2 + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, new ByteArrayInputStream(Array[Byte](9))) + + val input = rawReader(client, context, retries = 2).openPartitions() + assert(input.read() == 9) + assert(client.updateFileGroupCalls == 3) + assert(client.stageEndChecks == 3) + assert(client.failureReports == 0) + input.close() + context.markTaskCompleted(None) + } + + test("exhausted file-group timeouts propagate without invalidating a shuffle generation") { + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + client.stageEnded = false + client.timeoutFailures = 4 + + val failure = intercept[IOException] { + rawReader(client, context, retries = 2).openPartitions().read() + } + + assert(failure.getCause.isInstanceOf[TimeoutException]) + assert(client.updateFileGroupCalls == 3) + assert(client.failureReports == 0) + } + + test("file-group and partition-open failures invalidate and become Spark fetch failures") { + Seq(false, true).foreach { failAtOpen => + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + val expected = new IOException(if (failAtOpen) "worker open failed" else "metadata lost") + if (failAtOpen) { + client.fileGroups.partitionGroups.put(2, location("worker")) + client.readPartitionFailure = expected + } else { + client.updateFileGroupFailure = expected + } + + val failure = intercept[FetchFailedException] { + rawReader(client, context, startPartition = 2, endPartition = 3).openPartitions().read() + } + + assert(failure.getCause eq expected) + assert(failure.getMessage.contains("17/91")) + assert(client.failureReports == 1) + } + } + + test("lazy stream failures invalidate once and interrupted tasks never invalidate") { + Seq(false, true).foreach { interrupted => + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + val expected = new IOException("worker failed after stream creation") + client.fileGroups.partitionGroups.put(1, location("worker")) + client.streams.put( + 1, + new InputStream { + override def read(): Int = throw expected + }) + val input = + rawReader(client, context, startPartition = 1, endPartition = 2).openPartitions() + if (interrupted) context.markInterrupted("the reduce task was cancelled") + + val failure = intercept[Throwable](input.read()) + if (interrupted) { + assert(failure eq expected) + assert(client.failureReports == 0) + } else { + assert(failure.isInstanceOf[FetchFailedException]) + assert(failure.getCause eq expected) + assert(client.failureReports == 1) + } + } + } + + test("physical-skew encoded map ranges fail closed") { + val failure = intercept[IllegalArgumentException] { + rawReader(new RecordingCelebornRawClient, TaskContext.empty(), startMap = 9, endMap = 4) + } + assert(failure.getMessage.contains("physical-skew")) + } + + test("reader attempt identities match writer packing and reject wrapped attempt numbers") { + assert(CelebornRawPartitionReader.encodeAttemptNumber(0, 0) == 0) + assert(CelebornRawPartitionReader.encodeAttemptNumber(1, 7) == 65543) + assert(CelebornRawPartitionReader.encodeAttemptNumber(32767, 65535) == Int.MaxValue) + + Seq((-1, 0), (32768, 0), (0, -1), (0, 65536)).foreach { case (stageAttempt, taskAttempt) => + intercept[IllegalArgumentException] { + CelebornRawPartitionReader.encodeAttemptNumber(stageAttempt, taskAttempt) + } + } + } + + test("empty reducers are skipped and null mapper-attempt metadata fails closed") { + val emptyContext = TaskContext.empty() + val emptyClient = new RecordingCelebornRawClient + assert(rawReader(emptyClient, emptyContext).openPartitions().read() == -1) + assert(emptyClient.requests.isEmpty) + + val brokenContext = TaskContext.empty() + val brokenClient = new RecordingCelebornRawClient + brokenClient.fileGroups.mapAttempts = null + val failure = intercept[IllegalStateException] { + rawReader(brokenClient, brokenContext).openPartitions().read() + } + assert(failure.getMessage.contains("mapper-attempt")) + assert(brokenClient.requests.isEmpty) + } + + test("native ShuffleScan consumes real remote shuffle frames without decompression") { + withNativeFrame { (dependency, frame) => + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, new ByteArrayInputStream(frame)) + val iterator = + new CometShuffleBlockIterator(reader(dependency, client, context).readAsRawStream()) + + val expectedLength = ByteBuffer.wrap(frame).order(ByteOrder.LITTLE_ENDIAN).getLong.toInt - 8 + assert(iterator.hasNext() == expectedLength) + val actual = new Array[Byte](expectedLength) + val buffer = iterator.getBuffer.duplicate() + buffer.position(0) + buffer.get(actual) + assert(actual.toSeq == frame.slice(16, 16 + expectedLength).toSeq) + assert(iterator.hasNext() == -1) + assert(!client.requests.get(0).needDecompress) + context.markTaskCompleted(None) + } + } + + test("JVM consumers decode real remote shuffle frames and count rows") { + withNativeFrame { (dependency, frame) => + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + client.fileGroups.partitionGroups.put(0, location("worker")) + client.streams.put(0, new ByteArrayInputStream(frame)) + + val batches = reader(dependency, client, context).read() + var rows = 0 + while (batches.hasNext) { + rows += batches.next()._2.numRows() + } + + assert(rows > 0) + assert(context.taskMetrics.shuffleReadMetrics.recordsRead == rows) + assert(!client.requests.get(0).needDecompress) + context.markTaskCompleted(None) + } + } + + test("a remote shuffle reader cannot be consumed through both paths") { + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient + val remote = reader(simpleDependency(), client, context) + remote.readAsRawStream() + + val failure = intercept[IllegalStateException](remote.readAsRawStream()) + assert(failure.getMessage.contains("only be consumed once")) + context.markTaskCompleted(None) + } +} From 26b10243804f2f21ba6a3b074360f2b5267858ef Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:19:56 +0000 Subject: [PATCH 09/12] feat: enable native-only Celeborn shuffle planning --- .../comet/CometSparkSessionExtensions.scala | 14 +- .../shuffle/CometCelebornShuffleManager.scala | 9 +- .../shuffle/CometShuffleExchangeExec.scala | 34 +++- .../CometSparkSessionExtensionsSuite.scala | 177 +++++++++++++++--- 4 files changed, 197 insertions(+), 37 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 6232b330f25..b345e76d29e 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -177,11 +177,17 @@ object CometSparkSessionExtensions extends Logging { } } - // The Celeborn manager is valid for loading Comet, but its native shuffle writer and reader - // are not wired yet. Keep exchanges on Celeborn's existing Spark shuffle path until they are. def isCometShuffleEnabled(conf: SQLConf): Boolean = - COMET_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) && - conf.getConfString(SHUFFLE_MANAGER_KEY) != classOf[CometCelebornShuffleManager].getName + COMET_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) && { + // Explicit native mode opts out of Celeborn's local fallback. Keep this restriction at the + // common gate because CollectLimit and TakeOrdered bypass ordinary exchange planning. + !isCometCelebornShuffleManagerEnabled(conf) || + (COMET_EXEC_ENABLED.get(conf) && COMET_SHUFFLE_MODE.get(conf) == "native") + } + + def isCometCelebornShuffleManagerEnabled(conf: SQLConf): Boolean = + conf.contains(SHUFFLE_MANAGER_KEY) && + conf.getConfString(SHUFFLE_MANAGER_KEY) == classOf[CometCelebornShuffleManager].getName def isCometShuffleManagerEnabled(conf: SQLConf): Boolean = { conf.contains(SHUFFLE_MANAGER_KEY) && { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index 533ce59a7cc..2934975eb14 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -38,9 +38,9 @@ import org.apache.comet.util.ClassLoaders /** * Lets Comet execution coexist with the application's existing Celeborn shuffle manager. * - * Ordinary Spark shuffle dependencies are owned entirely by Celeborn. Native Comet map tasks can - * write to and read from Celeborn, while query planning keeps native shuffle disabled until its - * remaining production-readiness work is complete. JVM Comet shuffle remains unsupported. + * Ordinary Spark shuffle dependencies are owned entirely by Celeborn. Native Comet shuffle writes + * to and reads from Celeborn, while unsupported exchanges retain the existing Spark shuffle path. + * JVM Comet shuffle remains unsupported. * * Celeborn is loaded reflectively because its client is an optional, application-provided * dependency rather than part of Comet's compile-time or runtime distribution. @@ -395,8 +395,7 @@ class CometCelebornShuffleManager private[shuffle] ( private def rejectCometShuffle(): Nothing = { throw new UnsupportedOperationException( - "Comet shuffle over Celeborn is not supported yet; its remote writer, reader, " + - "and task lifecycle must be integrated before Comet shuffle can be enabled") + "Comet shuffle over Celeborn is not supported for JVM shuffle or non-Celeborn handles") } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 778c3013172..0e9813b2a59 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -51,7 +51,7 @@ import com.google.common.base.Objects import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometConf.{COMET_SHUFFLE_ENABLED, COMET_SHUFFLE_MODE} -import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isCometShuffleEnabled, isCometShuffleManagerEnabled, withFallbackReasons} +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isCometCelebornShuffleManagerEnabled, isCometShuffleEnabled, isCometShuffleManagerEnabled, withFallbackReasons} import org.apache.comet.serde.{Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.{CometTypeShim, ShimCometShuffleExchangeExec} @@ -347,14 +347,29 @@ object CometShuffleExchangeExec return None } - // Native path is only eligible when the child is already a Comet plan; otherwise skip it - // silently (no reason to surface) and let columnar take over. + val usesCelebornShuffleManager = isCometCelebornShuffleManagerEnabled(s.conf) + + // Native path is only eligible when the child is already a Comet plan; otherwise a local + // manager can fall back to columnar shuffle, but Celeborn must keep Spark's existing shuffle. + val nativeChild = isCometPlan(s.child) val nativeReasons: Seq[String] = - if (isCometPlan(s.child)) nativeShuffleFailureReasons(s) else Seq.empty - if (isCometPlan(s.child) && nativeReasons.isEmpty) { + if (nativeChild) nativeShuffleFailureReasons(s) else Seq.empty + if (nativeChild && nativeReasons.isEmpty) { return Some(CometNativeShuffle) } + if (usesCelebornShuffleManager) { + val reasons = if (nativeChild) { + nativeReasons + } else { + Seq("Celeborn native shuffle requires a Comet child") + } + val columnarReason = + "Comet columnar shuffle is not supported by the Celeborn shuffle manager" + withFallbackReasons(s, (reasons :+ columnarReason).toSet) + return None + } + if (!isCometPlan(s.child) && !CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(s.conf)) { withFallbackReasons( @@ -681,8 +696,13 @@ object CometShuffleExchangeExec s"${classOf[CometShuffleManager].getName} or " + classOf[CometCelebornShuffleManager].getName) } else if (!isCometShuffleEnabled(op.conf)) { - Some( - "Celeborn-backed Comet shuffle is unavailable until its native writer and reader are wired") + if (!CometConf.COMET_EXEC_ENABLED.get(op.conf)) { + Some("Celeborn-backed Comet shuffle requires Comet native execution to be enabled") + } else if (COMET_SHUFFLE_MODE.get(op.conf) == "jvm") { + Some("Celeborn-backed Comet shuffle does not support spark.comet.shuffle.mode=jvm") + } else { + Some("Celeborn-backed Comet shuffle requires spark.comet.shuffle.mode=native") + } } else { None } diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index d4a2586b7a4..8f7a5c858ed 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -21,15 +21,47 @@ package org.apache.comet import org.apache.spark.SparkConf import org.apache.spark.sql._ -import org.apache.spark.sql.catalyst.plans.physical.SinglePartition -import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleExchangeExec, CometShuffleManager} +import org.apache.spark.sql.catalyst.plans.physical.{RoundRobinPartitioning, SinglePartition} +import org.apache.spark.sql.comet.{CometNativeExec, CometSinkPlaceHolder} +import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec, CometShuffleManager} import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.internal.SQLConf +import org.apache.comet.serde.OperatorOuterClass + class CometSparkSessionExtensionsSuite extends CometTestBase { import CometSparkSessionExtensions._ + private def withShuffleManagerSession( + manager: String, + mode: String = "auto", + nativeExecution: Boolean = true)(f: SQLConf => Unit): Unit = { + val session = spark.newSession() + val conf = session.sessionState.conf + conf.setConfString(CometConf.COMET_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_SHUFFLE_MODE.key, mode) + conf.setConfString(CometConf.COMET_EXEC_ENABLED.key, nativeExecution.toString) + conf.setConfString("spark.shuffle.manager", manager) + + val previousActiveSession = SparkSession.getActiveSession + try { + SparkSession.setActiveSession(session) + f(conf) + } finally { + previousActiveSession match { + case Some(previousSession) => SparkSession.setActiveSession(previousSession) + case None => SparkSession.clearActiveSession() + } + } + } + + private def nativeShuffleChild(): CometNativeExec = { + val child = spark.emptyDataFrame.queryExecution.executedPlan + CometSinkPlaceHolder(OperatorOuterClass.Operator.getDefaultInstance, child, child) + } + test("isCometLoaded") { val conf = new SQLConf // Disable Comet shuffle so this test can focus on other checks without needing @@ -75,44 +107,147 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { assert(isCometShuffleEnabled(conf)) } - test("Celeborn manager loads Comet without enabling its unfinished native shuffle transport") { + test("Celeborn manager enables Comet shuffle only with explicit native opt-in") { val conf = new SQLConf conf.setConfString(CometConf.COMET_ENABLED.key, "true") conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_EXEC_ENABLED.key, "true") conf.setConfString("spark.shuffle.manager", classOf[CometCelebornShuffleManager].getName) assert(isCometShuffleManagerEnabled(conf)) assert(isCometLoaded(conf)) + assert(!isCometShuffleEnabled(conf), "default auto mode must preserve Spark shuffle") + + Seq("auto", "jvm").foreach { mode => + conf.setConfString(CometConf.COMET_SHUFFLE_MODE.key, mode) + assert( + !isCometShuffleEnabled(conf), + s"Celeborn native shuffle must not be enabled for mode=$mode") + } + + conf.setConfString(CometConf.COMET_SHUFFLE_MODE.key, "native") + assert(isCometShuffleEnabled(conf)) + + conf.setConfString(CometConf.COMET_EXEC_ENABLED.key, "false") + assert(!isCometShuffleEnabled(conf)) + + conf.setConfString(CometConf.COMET_EXEC_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "false") assert(!isCometShuffleEnabled(conf)) } - test("Celeborn manager leaves shuffle exchanges on its existing Spark shuffle path") { - val child = spark.emptyDataFrame.queryExecution.executedPlan - val session = spark.newSession() - val conf = session.sessionState.conf - conf.setConfString(CometConf.COMET_ENABLED.key, "true") - conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") - conf.setConfString("spark.shuffle.manager", classOf[CometCelebornShuffleManager].getName) + test("Celeborn manager selects native shuffle for supported Comet native children") { + val child = nativeShuffleChild() - val previousActiveSession = SparkSession.getActiveSession - try { - SparkSession.setActiveSession(session) + withShuffleManagerSession(classOf[CometCelebornShuffleManager].getName, "native") { _ => val shuffle = ShuffleExchangeExec(SinglePartition, child) + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).contains(CometNativeShuffle)) + assert(shuffle.getTagValue(CometExplainInfo.FALLBACK_REASONS).isEmpty) + } + } + + test("Celeborn manager preserves Spark shuffle for default and explicit auto mode") { + val child = nativeShuffleChild() + + Seq(false, true).foreach { explicitAutoMode => + withShuffleManagerSession(classOf[CometCelebornShuffleManager].getName) { conf => + if (!explicitAutoMode) { + conf.unsetConf(CometConf.COMET_SHUFFLE_MODE.key) + } + assert(CometConf.COMET_SHUFFLE_MODE.get(conf) == "auto") + assert(!isCometShuffleEnabled(conf)) + + val shuffle = ShuffleExchangeExec(SinglePartition, child) + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).isEmpty) + assert( + shuffle + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + .exists(_.contains("requires spark.comet.shuffle.mode=native"))) + } + } + } + + test("Celeborn manager leaves non-native Spark children on the Spark shuffle path") { + val sparkChild = spark.emptyDataFrame.queryExecution.executedPlan + val nativeChild = nativeShuffleChild() + + withShuffleManagerSession(classOf[CometCelebornShuffleManager].getName, "native") { _ => + val shuffle = ShuffleExchangeExec(SinglePartition, sparkChild) + + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).isEmpty) + val reasons = shuffle + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + assert(reasons.exists(_.contains("requires a Comet child"))) + assert(reasons.exists(_.contains("columnar shuffle is not supported"))) + + // AQE may reshape the child later; a prior Spark-fallback decision must remain sticky. + val reshaped = shuffle.withNewChildren(Seq(nativeChild)).asInstanceOf[ShuffleExchangeExec] + assert(CometShuffleExchangeExec.shuffleSupported(reshaped).isEmpty) + } + } + + test("unsupported Celeborn native partitioning falls back to Spark instead of Comet columnar") { + val child = nativeShuffleChild() + + withShuffleManagerSession(classOf[CometCelebornShuffleManager].getName, "native") { _ => + val shuffle = ShuffleExchangeExec(RoundRobinPartitioning(2), child) + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).isEmpty) + val reasons = shuffle + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) assert( - shuffle - .getTagValue(CometExplainInfo.FALLBACK_REASONS) - .getOrElse(Set.empty[String]) - .exists(_.contains("Celeborn-backed Comet shuffle is unavailable"))) - } finally { - previousActiveSession match { - case Some(previousSession) => SparkSession.setActiveSession(previousSession) - case None => SparkSession.clearActiveSession() + reasons.exists( + _.contains(CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key))) + assert(reasons.exists(_.contains("columnar shuffle is not supported"))) + } + + withShuffleManagerSession(classOf[CometShuffleManager].getName) { _ => + val shuffle = ShuffleExchangeExec(RoundRobinPartitioning(2), child) + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).contains(CometColumnarShuffle)) + } + } + + test("Celeborn manager explains JVM mode and disabled native execution fallback") { + val child = nativeShuffleChild() + val scenarios = Seq( + ("jvm", true, "does not support spark.comet.shuffle.mode=jvm"), + ("auto", false, "requires Comet native execution to be enabled")) + + scenarios.foreach { case (mode, nativeExecution, expectedReason) => + withShuffleManagerSession( + classOf[CometCelebornShuffleManager].getName, + mode, + nativeExecution) { conf => + assert(!isCometShuffleEnabled(conf)) + + val shuffle = ShuffleExchangeExec(SinglePartition, child) + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).isEmpty) + assert( + shuffle + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + .exists(_.contains(expectedReason))) } } } + test("local Comet manager still supports columnar shuffle without native execution") { + val child = spark.emptyDataFrame.queryExecution.executedPlan + + withShuffleManagerSession( + classOf[CometShuffleManager].getName, + mode = "jvm", + nativeExecution = false) { conf => + assert(isCometShuffleEnabled(conf)) + val shuffle = ShuffleExchangeExec(SinglePartition, child) + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).contains(CometColumnarShuffle)) + } + } + test("stock Celeborn manager does not satisfy Comet shuffle-manager requirements") { val conf = new SQLConf conf.setConfString(CometConf.COMET_ENABLED.key, "true") From 6422794bee03e4d127191ac8aca856755670367a Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 25 Aug 2026 23:46:46 +0000 Subject: [PATCH 10/12] fix: make Celeborn remote shuffle compatible with upstream Apache releases --- docs/source/user-guide/latest/tuning.md | 26 +++ .../src/execution/rss_planner_jni_tests.rs | 2 + .../src/writers/rss/rss_partition_writer.rs | 3 + .../CelebornShufflePartitionPusher.java | 97 ++++++--- .../scala/org/apache/comet/CometConf.scala | 13 +- .../comet/CometSparkSessionExtensions.scala | 8 +- .../CelebornShufflePusherFactory.scala | 6 +- .../shuffle/CometCelebornShuffleManager.scala | 16 +- .../shuffle/CometCelebornShuffleReader.scala | 54 ++--- .../shuffle/CometNativeShuffleWriter.scala | 18 +- .../shuffle/CometShuffleExchangeExec.scala | 8 +- .../shuffle/RecordingCelebornRawClient.java | 44 +++- .../CometSparkSessionExtensionsSuite.scala | 26 +++ .../CelebornShufflePartitionPusherSuite.scala | 202 +++++++++++++++++- ...ometCelebornNativeShuffleWriterSuite.scala | 36 ++++ .../CometCelebornShuffleReaderSuite.scala | 21 ++ 16 files changed, 497 insertions(+), 83 deletions(-) diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 5499052464c..639c43dc55a 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -220,6 +220,32 @@ Comet provides a fully native shuffle implementation, which generally provides t supports `HashPartitioning`, `RangePartitioning` and `SinglePartitioning` but currently only supports primitive type partitioning keys. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays. +#### Apache Celeborn Remote Shuffle + +Comet native shuffle can write partition data to and read it from an existing Apache Celeborn remote shuffle +service. Provide an Apache Celeborn 0.7.0 or newer Spark client on the driver and executor classpaths, then +configure the composite shuffle manager and explicitly enable native shuffle: + +``` +spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager +spark.celeborn.master.endpoints=celeborn-master:9097 +spark.comet.enabled=true +spark.comet.exec.enabled=true +spark.comet.shuffle.enabled=true +spark.comet.shuffle.mode=native +``` + +The composite manager delegates ordinary Spark shuffles to Celeborn and uses native Comet shuffle only for +supported exchanges whose child is already a Comet plan. Unsupported exchanges, the default `auto` shuffle mode, +and JVM columnar shuffle remain on the existing Spark/Celeborn shuffle path. The Celeborn client is optional and +is not bundled with Comet. Spark I/O encryption (`spark.io.encryption.enabled=true`) is not supported by the +native remote shuffle path; encrypted applications continue to use the existing Spark/Celeborn shuffle path. + +Set `spark.comet.shuffle.celeborn.enabled=false` to disable Celeborn-backed partition pushers. The maximum +complete partition frame size defaults to 64 MiB and can be configured with +`spark.comet.shuffle.rss.maxFrameBytes`. The executor-wide in-flight push limit defaults to 256 MiB and can be +configured with `spark.comet.shuffle.rss.maxInFlightBytes`. + #### Columnar (JVM) Shuffle Comet Columnar shuffle is JVM-based and supports `HashPartitioning`, `RoundRobinPartitioning`, `RangePartitioning`, and diff --git a/native/core/src/execution/rss_planner_jni_tests.rs b/native/core/src/execution/rss_planner_jni_tests.rs index ed191c6de73..17172f0ffc1 100644 --- a/native/core/src/execution/rss_planner_jni_tests.rs +++ b/native/core/src/execution/rss_planner_jni_tests.rs @@ -110,6 +110,8 @@ fn execution_context( tracing_memory_metric_name: String::new(), tracing_event_name: String::new(), task_context: None, + class_loader: None, + memory_pool_registration: None, rss_pusher: None, }) } diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs index 15ecd2ccac2..c66b604f6c7 100644 --- a/native/shuffle/src/writers/rss/rss_partition_writer.rs +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -24,6 +24,7 @@ use arrow::array::{ }; use arrow::buffer::OffsetBuffer; use arrow::datatypes::DataType; +use arrow::ipc::writer::CompressionContext; use arrow::record_batch::RecordBatch; use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion::common::{DataFusionError, Result}; @@ -94,6 +95,7 @@ impl PartitionPusher for JavaShufflePartitionPusher { pub struct RssPartitionWriter { pusher: P, block_writer: ShuffleBlockWriter, + compression_context: CompressionContext, num_partitions: usize, max_frame_bytes: usize, next_partition: usize, @@ -112,6 +114,7 @@ impl RssPartitionWriter

{ Ok(Self { pusher, block_writer, + compression_context: CompressionContext::default(), num_partitions, max_frame_bytes, next_partition: 0, diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java index 19f87577a11..2dde5404cbb 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java +++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java @@ -226,9 +226,17 @@ public CelebornShufflePartitionPusher( "Celeborn raw-push API must be an instance method returning an int"); } - Method mapperEndMethod = - resolveLifecycleMethod( - shuffleClient, "mapperEnd", int.class, int.class, int.class, int.class); + Method mapperEndMethod; + try { + mapperEndMethod = + resolveLifecycleMethod( + shuffleClient, "mapperEnd", int.class, int.class, int.class, int.class, int.class); + } catch (IllegalArgumentException unsupportedStandardMapperEnd) { + // Older application-provided Celeborn clients did not require the reducer count. + mapperEndMethod = + resolveLifecycleMethod( + shuffleClient, "mapperEnd", int.class, int.class, int.class, int.class); + } Method cleanupMethod = resolveLifecycleMethod(shuffleClient, "cleanup", int.class, int.class, int.class); final Method getPushStateMethod; @@ -241,8 +249,15 @@ public CelebornShufflePartitionPusher( try { getPushStateMethod = shuffleClient.getClass().getMethod("getPushState", String.class); setPushMetricsCallbackMethod = resolvePushMetricsCallback(getPushStateMethod.getReturnType()); - pushMetricsCallbackClass = setPushMetricsCallbackMethod.getParameterTypes()[0]; - pushMetricsCallbackClass.getMethod("incPushDataCount", long.class); + if (setPushMetricsCallbackMethod == null) { + // Apache Celeborn does not expose push-completion callbacks. Its in-flight request + // tracker still records every accepted request and completion, so periodic reconciliation + // provides the same completion-backed admission without requiring a patched client. + pushMetricsCallbackClass = null; + } else { + pushMetricsCallbackClass = setPushMetricsCallbackMethod.getParameterTypes()[0]; + pushMetricsCallbackClass.getMethod("incPushDataCount", long.class); + } clientPushStatesField = resolveClientPushStates(shuffleClient.getClass()); inFlightRequestTrackerField = getPushStateMethod.getReturnType().getDeclaredField("inFlightRequestTracker"); @@ -265,9 +280,10 @@ public CelebornShufflePartitionPusher( e); } if (Modifier.isStatic(getPushStateMethod.getModifiers()) - || Modifier.isStatic(setPushMetricsCallbackMethod.getModifiers()) - || setPushMetricsCallbackMethod.getReturnType() != void.class - || !pushMetricsCallbackClass.isInterface()) { + || (setPushMetricsCallbackMethod != null + && (Modifier.isStatic(setPushMetricsCallbackMethod.getModifiers()) + || setPushMetricsCallbackMethod.getReturnType() != void.class + || !pushMetricsCallbackClass.isInterface()))) { throw new IllegalArgumentException("Celeborn push-state completion API is incompatible"); } @@ -293,14 +309,13 @@ public CelebornShufflePartitionPusher( this.partitionLengths = new AtomicLongArray(numPartitions); } - private static Method resolvePushMetricsCallback(Class pushStateClass) - throws NoSuchMethodException { + private static Method resolvePushMetricsCallback(Class pushStateClass) { for (Method method : pushStateClass.getMethods()) { if (method.getName().equals("setMetricsCallback") && method.getParameterCount() == 1) { return method; } } - throw new NoSuchMethodException("PushState.setMetricsCallback"); + return null; } private static Field resolveClientPushStates(Class clientClass) throws NoSuchFieldException { @@ -532,7 +547,12 @@ public long[] finish() throws IOException { mapperEndThread = Thread.currentThread(); } try { - mapperEnd.invoke(shuffleClient, shuffleId, mapId, encodedAttemptId, numMappers); + if (mapperEnd.getParameterCount() == 5) { + mapperEnd.invoke( + shuffleClient, shuffleId, mapId, encodedAttemptId, numMappers, numPartitions); + } else { + mapperEnd.invoke(shuffleClient, shuffleId, mapId, encodedAttemptId, numMappers); + } } finally { synchronized (lifecycleLock) { mapperEndThread = null; @@ -663,25 +683,27 @@ private ObservedPushState observePushState(Object pushState) new ObservedPushState( (LongAdder) totalInFlightRequests.get(tracker), (AtomicReference) pushStateException.get(pushState)); - Object callback = - Proxy.newProxyInstance( - pushMetricsCallbackClass.getClassLoader(), - new Class[] {pushMetricsCallbackClass}, - (proxy, method, arguments) -> { - if (method.getName().equals("incPushDataCount")) { - boolean metricsOnlyFailure = - created.inFlightRequests.sum() > 0 && isTerminalFailureCallback(); - completeAcceptedPushes(created, (long) arguments[0], metricsOnlyFailure); - } else if (method.getName().equals("hashCode")) { - return System.identityHashCode(proxy); - } else if (method.getName().equals("equals")) { - return proxy == arguments[0]; - } else if (method.getName().equals("toString")) { - return "Celeborn native shuffle push completion callback"; - } - return null; - }); - setPushMetricsCallback.invoke(pushState, callback); + if (setPushMetricsCallback != null) { + Object callback = + Proxy.newProxyInstance( + pushMetricsCallbackClass.getClassLoader(), + new Class[] {pushMetricsCallbackClass}, + (proxy, method, arguments) -> { + if (method.getName().equals("incPushDataCount")) { + boolean metricsOnlyFailure = + created.inFlightRequests.sum() > 0 && isTerminalFailureCallback(); + completeAcceptedPushes(created, (long) arguments[0], metricsOnlyFailure); + } else if (method.getName().equals("hashCode")) { + return System.identityHashCode(proxy); + } else if (method.getName().equals("equals")) { + return proxy == arguments[0]; + } else if (method.getName().equals("toString")) { + return "Celeborn native shuffle push completion callback"; + } + return null; + }); + setPushMetricsCallback.invoke(pushState, callback); + } created.failedBeforeObservation = created.exception.get() != null; observedPushStates.put(pushState, created); return created; @@ -769,6 +791,16 @@ private void reconcileAcceptedPushes() { final ArrayList completed = new ArrayList<>(); synchronized (lifecycleLock) { for (ObservedPushState pushState : observedPushStates.values()) { + if (setPushMetricsCallback == null + && (state == State.OPEN || state == State.FINISHING) + && pushState.exception.get() != null + && !pushState.recoveredMissedFailure) { + // Apache Celeborn records a terminal push failure without removing its transport batch. + // Its first exception proves exactly one request completed unsuccessfully. Cancellation + // also installs an exception, so never infer completion after this attempt was aborted. + pushState.metricsOnlyFailures++; + pushState.recoveredMissedFailure = true; + } long transportRequests = Math.max(0L, pushState.inFlightRequests.sum() - pushState.metricsOnlyFailures); long trackerCompletions = Math.max(0L, pushState.submittedPushes - transportRequests); @@ -894,6 +926,9 @@ private long[] snapshotPartitionLengths() { private void abortAndSuppress(Throwable original) { try { + // A stock client can report terminal failure synchronously without a completion callback. + // Reconcile its exception before abort() marks the attempt cancelled and hides that signal. + reconcileAcceptedPushes(); abort(); } catch (Throwable cleanupFailure) { if (cleanupFailure != original) { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 31a103a3998..f11e2119698 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -308,12 +308,23 @@ object CometConf extends ShimCometConf { .doc( "Whether to enable Comet native shuffle. " + "Note that this requires setting `spark.shuffle.manager` to " + - "`org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager`. " + + "`org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager` or " + + "`org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager`. " + "`spark.shuffle.manager` must be set before starting the Spark application and " + "cannot be changed during the application.") .booleanConf .createWithDefault(true) + val COMET_SHUFFLE_CELEBORN_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.shuffle.celeborn.enabled") + .category(CATEGORY_SHUFFLE) + .doc( + "Whether to allow Comet native shuffle to use an application-provided Apache " + + "Celeborn client when a Celeborn shuffle manager, shuffle data I/O plugin, " + + "or master endpoint is configured.") + .booleanConf + .createWithDefault(true) + val COMET_SHUFFLE_DIRECT_READ_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.directRead.enabled") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.directRead.enabled") diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index b345e76d29e..d9a85633f7f 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -123,6 +123,7 @@ object CometSparkSessionExtensions extends Logging { lazy val isBigEndian: Boolean = ByteOrder.nativeOrder().equals(ByteOrder.BIG_ENDIAN) private val SHUFFLE_MANAGER_KEY = "spark.shuffle.manager" + private val IO_ENCRYPTION_ENABLED_KEY = "spark.io.encryption.enabled" /** * Checks whether Comet extension should be loaded for Spark. @@ -182,9 +183,14 @@ object CometSparkSessionExtensions extends Logging { // Explicit native mode opts out of Celeborn's local fallback. Keep this restriction at the // common gate because CollectLimit and TakeOrdered bypass ordinary exchange planning. !isCometCelebornShuffleManagerEnabled(conf) || - (COMET_EXEC_ENABLED.get(conf) && COMET_SHUFFLE_MODE.get(conf) == "native") + (COMET_EXEC_ENABLED.get(conf) && COMET_SHUFFLE_MODE.get(conf) == "native" && + COMET_SHUFFLE_CELEBORN_ENABLED.get(conf) && + !isCometCelebornShuffleEncryptionEnabled(conf)) } + def isCometCelebornShuffleEncryptionEnabled(conf: SQLConf): Boolean = + conf.getConfString(IO_ENCRYPTION_ENABLED_KEY, "false").toBoolean + def isCometCelebornShuffleManagerEnabled(conf: SQLConf): Boolean = conf.contains(SHUFFLE_MANAGER_KEY) && conf.getConfString(SHUFFLE_MANAGER_KEY) == classOf[CometCelebornShuffleManager].getName diff --git a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala index 9dc48b40e12..5e3ca25c0d1 100644 --- a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala +++ b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala @@ -32,7 +32,8 @@ import org.apache.comet.util.ClassLoaders /** Creates task-owned Celeborn pushers using the application's existing Spark configuration. */ object CelebornShufflePusherFactory { - private val CELEBORN_ENABLED_KEY = "spark.comet.celeborn.enabled" + private val CELEBORN_ENABLED = CometConf.COMET_SHUFFLE_CELEBORN_ENABLED + private val IO_ENCRYPTION_ENABLED_KEY = "spark.io.encryption.enabled" private val SHUFFLE_MANAGER_KEY = "spark.shuffle.manager" private val SHUFFLE_DATA_IO_KEY = "spark.shuffle.sort.io.plugin.class" private val CELEBORN_MASTER_ENDPOINTS_KEY = "spark.celeborn.master.endpoints" @@ -55,7 +56,8 @@ object CelebornShufflePusherFactory { /** Detects resolved Celeborn configuration while honoring an explicit application opt-out. */ def isEnabled(conf: SparkConf): Boolean = { - conf.getBoolean(CELEBORN_ENABLED_KEY, true) && + conf.getBoolean(CELEBORN_ENABLED.key, CELEBORN_ENABLED.defaultValue.get) && + !conf.getBoolean(IO_ENCRYPTION_ENABLED_KEY, false) && (conf .getOption(SHUFFLE_MANAGER_KEY) .exists(manager => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index 2934975eb14..82820db5dda 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -553,12 +553,8 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( private def invalidateOwners(shuffleId: Int): Unit = { generationEpochs.update(shuffleId, currentEpoch(shuffleId) + 1L) - claimOwners.filterInPlace { case ((ownerShuffleId, _, _, _), _) => - ownerShuffleId != shuffleId - } - deniedAttempts.filterInPlace { case ((ownerShuffleId, _, _, _), _) => - ownerShuffleId != shuffleId - } + claimOwners.retain((key, _) => key._1 != shuffleId) + deniedAttempts.retain((key, _) => key._1 != shuffleId) } private def resetSparkCommitOwners(generation: PrepareCelebornShuffleGeneration): Unit = @@ -791,12 +787,8 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( generations.remove(shuffleId) invalidatedGenerations.remove(shuffleId) generationEpochs.remove(shuffleId) - claimOwners.filterInPlace { case ((ownerShuffleId, _, _, _), _) => - ownerShuffleId != shuffleId - } - deniedAttempts.filterInPlace { case ((ownerShuffleId, _, _, _), _) => - ownerShuffleId != shuffleId - } + claimOwners.retain((key, _) => key._1 != shuffleId) + deniedAttempts.retain((key, _) => key._1 != shuffleId) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala index ac3dcddfa05..3bc9b08c842 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala @@ -158,19 +158,26 @@ private[shuffle] final class CelebornRawPartitionReader( private val updateFileGroup = client.getClass.getMethod("updateFileGroup", java.lang.Integer.TYPE, java.lang.Integer.TYPE) private val readPartitionMethod = client.getClass.getMethods - .find { method => + .filter { method => val parameters = method.getParameterTypes - method.getName == "readPartition" && parameters.length == 16 && + method.getName == "readPartition" && (parameters.length == 15 || parameters.length == 16) && parameters(0) == java.lang.Integer.TYPE && parameters(4) == java.lang.Long.TYPE && - parameters(15) == java.lang.Boolean.TYPE && parameters(14).isInterface && + parameters(parameters.length - 1) == java.lang.Boolean.TYPE && + parameters(parameters.length - 2).isInterface && + parameters(parameters.length - 3) == classOf[Array[Int]] && classOf[InputStream].isAssignableFrom(method.getReturnType) } + // Prefer Apache Celeborn's public API when a client also retains an extended overload. + .sortBy(_.getParameterCount) + .headOption .getOrElse { throw new IllegalStateException( - "The Celeborn client does not expose its raw 16-argument partition reader") + "The Celeborn client does not expose its raw 15- or 16-argument partition reader") } private val metricsCallback = - createMetricsCallback(readPartitionMethod.getParameterTypes.apply(14), readMetrics) + createMetricsCallback( + readPartitionMethod.getParameterTypes.apply(readPartitionMethod.getParameterCount - 2), + readMetrics) private lazy val fileGroups: AnyRef = loadFileGroups() @@ -250,28 +257,25 @@ private[shuffle] final class CelebornRawPartitionReader( pushFailedBatches: AnyRef, mapAttempts: Array[Int]): InputStream = { val attempt = encodeAttemptNumber(context.stageAttemptNumber(), context.attemptNumber()) + val arguments = Array[AnyRef]( + Int.box(celebornShuffleId), + Int.box(sparkShuffleId), + Int.box(partition), + Int.box(attempt), + Long.box(context.taskAttemptId()), + Int.box(startMapIndex), + Int.box(endMapIndex), + null, + locations, + null, + pushFailedBatches, + null) ++ + (if (readPartitionMethod.getParameterCount == 16) Array[AnyRef](null) + else Array.empty[AnyRef]) ++ + Array[AnyRef](mapAttempts, metricsCallback, java.lang.Boolean.FALSE) val stream = try { - Option( - invoke( - readPartitionMethod, - client, - Int.box(celebornShuffleId), - Int.box(sparkShuffleId), - Int.box(partition), - Int.box(attempt), - Long.box(context.taskAttemptId()), - Int.box(startMapIndex), - Int.box(endMapIndex), - null, - locations, - null, - pushFailedBatches, - null, - null, - mapAttempts, - metricsCallback, - java.lang.Boolean.FALSE)) + Option(invoke(readPartitionMethod, client, arguments: _*)) .getOrElse { throw new IOException(s"Celeborn returned a null stream for reducer $partition") } 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 d411b56fb68..cc109fa82bf 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 @@ -523,16 +523,20 @@ private[shuffle] object CelebornNativeShuffleDestination { pusher: CelebornShufflePartitionPusher, reportFailure: Throwable => Unit): ScheduledFuture[_] = { val handled = new AtomicBoolean(false) + def abortOnce(): Unit = { + if (handled.compareAndSet(false, true)) { + try pusher.abort() + catch { + case failure: Throwable => reportFailure(failure) + } + } + } + + taskContext.addTaskFailureListener((_: TaskContext, _: Throwable) => abortOnce()) cancellationWatcher.scheduleWithFixedDelay( new Runnable { override def run(): Unit = { - if ((taskContext.isInterrupted() || taskContext.isFailed()) && - handled.compareAndSet(false, true)) { - try pusher.abort() - catch { - case failure: Throwable => reportFailure(failure) - } - } + if (taskContext.isInterrupted()) abortOnce() } }, 0L, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 0e9813b2a59..806017d5e74 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -51,7 +51,7 @@ import com.google.common.base.Objects import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometConf.{COMET_SHUFFLE_ENABLED, COMET_SHUFFLE_MODE} -import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isCometCelebornShuffleManagerEnabled, isCometShuffleEnabled, isCometShuffleManagerEnabled, withFallbackReasons} +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isCometCelebornShuffleEncryptionEnabled, isCometCelebornShuffleManagerEnabled, isCometShuffleEnabled, isCometShuffleManagerEnabled, withFallbackReasons} import org.apache.comet.serde.{Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.{CometTypeShim, ShimCometShuffleExchangeExec} @@ -698,6 +698,12 @@ object CometShuffleExchangeExec } else if (!isCometShuffleEnabled(op.conf)) { if (!CometConf.COMET_EXEC_ENABLED.get(op.conf)) { Some("Celeborn-backed Comet shuffle requires Comet native execution to be enabled") + } else if (!CometConf.COMET_SHUFFLE_CELEBORN_ENABLED.get(op.conf)) { + Some( + "Celeborn-backed Comet shuffle is disabled: " + + CometConf.COMET_SHUFFLE_CELEBORN_ENABLED.key) + } else if (isCometCelebornShuffleEncryptionEnabled(op.conf)) { + Some("Celeborn-backed Comet shuffle does not support spark.io.encryption.enabled=true") } else if (COMET_SHUFFLE_MODE.get(op.conf) == "jvm") { Some("Celeborn-backed Comet shuffle does not support spark.comet.shuffle.mode=jvm") } else { diff --git a/spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java b/spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java index 9833c6efbe8..44c8066fb01 100644 --- a/spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java +++ b/spark/src/test/java/org/apache/spark/sql/comet/execution/shuffle/RecordingCelebornRawClient.java @@ -37,7 +37,7 @@ import org.apache.spark.shuffle.ShuffleReader; /** Exposes the pinned Celeborn read contract without making Celeborn a Comet test dependency. */ -public final class RecordingCelebornRawClient { +public class RecordingCelebornRawClient { private static final AtomicBoolean BROADCAST_DECODER_REGISTERED = new AtomicBoolean(false); @@ -63,6 +63,48 @@ public interface OptionalMetricsReporter { void incCelebornRemoteReadRetryCount(long count); } + /** Also exposes the unmodified, 15-argument Apache Celeborn partition-reader API. */ + public static final class StockApiClient extends RecordingCelebornRawClient { + public int stockReadPartitionCalls; + + public InputStream readPartition( + int shuffleId, + int appShuffleId, + int partitionId, + int attemptNumber, + long taskId, + int startMapIndex, + int endMapIndex, + Object exceptionMaker, + ArrayList locations, + ArrayList streamHandlers, + Map pushFailedBatches, + Map chunksRange, + int[] mapAttempts, + MetricsCallback metricsCallback, + boolean needDecompress) + throws IOException { + stockReadPartitionCalls++; + return super.readPartition( + shuffleId, + appShuffleId, + partitionId, + attemptNumber, + taskId, + startMapIndex, + endMapIndex, + exceptionMaker, + locations, + streamHandlers, + pushFailedBatches, + chunksRange, + null, + mapAttempts, + metricsCallback, + needDecompress); + } + } + public static final class BackendReader implements ShuffleReader { private final RecordingCelebornRawClient client; diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 8f7a5c858ed..1e923c5d3fd 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -128,6 +128,15 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { conf.setConfString(CometConf.COMET_SHUFFLE_MODE.key, "native") assert(isCometShuffleEnabled(conf)) + conf.setConfString(CometConf.COMET_SHUFFLE_CELEBORN_ENABLED.key, "false") + assert(!isCometShuffleEnabled(conf)) + conf.setConfString(CometConf.COMET_SHUFFLE_CELEBORN_ENABLED.key, "true") + + conf.setConfString("spark.io.encryption.enabled", "true") + assert(!isCometShuffleEnabled(conf)) + conf.setConfString("spark.io.encryption.enabled", "false") + assert(isCometShuffleEnabled(conf)) + conf.setConfString(CometConf.COMET_EXEC_ENABLED.key, "false") assert(!isCometShuffleEnabled(conf)) @@ -147,6 +156,23 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { } } + test("Celeborn manager keeps encrypted Spark shuffle on the existing Celeborn path") { + val child = nativeShuffleChild() + + withShuffleManagerSession(classOf[CometCelebornShuffleManager].getName, "native") { conf => + conf.setConfString("spark.io.encryption.enabled", "true") + assert(!isCometShuffleEnabled(conf)) + + val shuffle = ShuffleExchangeExec(SinglePartition, child) + assert(CometShuffleExchangeExec.shuffleSupported(shuffle).isEmpty) + assert( + shuffle + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + .exists(_.contains("spark.io.encryption.enabled=true"))) + } + } + test("Celeborn manager preserves Spark shuffle for default and explicit auto mode") { val child = nativeShuffleChild() diff --git a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala index 420a895f6d4..edc9d78d7b5 100644 --- a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala +++ b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala @@ -28,7 +28,7 @@ import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.{SparkConf, TaskContext} import org.apache.spark.shuffle.ShuffleHandle -import org.apache.comet.CometExecIterator +import org.apache.comet.{CometConf, CometExecIterator} /** Matches the pinned Celeborn completion hook without adding its optional client dependency. */ trait RecordingCelebornPushMetricsCallback { @@ -303,6 +303,89 @@ final class RecordingCelebornPushState(client: RecordingCelebornShuffleClient) { } } +/** Mirrors Apache Celeborn's public client without a push-completion metrics callback. */ +final class StockCelebornShuffleClient { + @volatile var automaticallyCompletePushes = true + @volatile var lastMapperEnd: (Int, Int, Int, Int, Int) = _ + val pushCalls = new AtomicInteger() + private val pushStates = new ConcurrentHashMap[String, StockCelebornPushState]() + private val pendingPushCompletions = new ConcurrentLinkedDeque[StockCelebornPushState]() + + def getPushState(mapKey: String): StockCelebornPushState = + pushStates.computeIfAbsent(mapKey, _ => new StockCelebornPushState) + + def pushOrMergeData( + shuffleId: Int, + mapId: Int, + attemptId: Int, + partitionId: Int, + bytes: Array[Byte], + offset: Int, + length: Int, + numMappers: Int, + numPartitions: Int, + doPush: Boolean, + skipCompress: Boolean): Int = { + pushCalls.incrementAndGet() + val pushState = getPushState(s"$shuffleId-$mapId-$attemptId") + pushState.addPush() + pendingPushCompletions.add(pushState) + if (automaticallyCompletePushes) completeNextPush() + length + 16 + } + + def completeNextPush(): Boolean = { + val pushState = pendingPushCompletions.pollFirst() + if (pushState == null) false + else { + pushState.completePush() + true + } + } + + def failNextPush(): Boolean = { + val pushState = pendingPushCompletions.pollFirst() + if (pushState == null) false + else { + pushState.failPushWithoutRemovingBatch() + true + } + } + + def mapperEnd( + shuffleId: Int, + mapId: Int, + attemptId: Int, + numMappers: Int, + numPartitions: Int): Unit = { + Option(pushStates.get(s"$shuffleId-$mapId-$attemptId")).foreach(_.checkFailure()) + lastMapperEnd = (shuffleId, mapId, attemptId, numMappers, numPartitions) + } + + def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = { + Option(pushStates.remove(s"$shuffleId-$mapId-$attemptId")).foreach(_.cleanup()) + } +} + +final class StockCelebornPushState { + private val inFlightRequestTracker = new RecordingCelebornInFlightRequestTracker() + private val exception = new AtomicReference[IOException]() + + def addPush(): Unit = inFlightRequestTracker.addBatch() + + def completePush(): Unit = inFlightRequestTracker.removeBatch() + + def failPushWithoutRemovingBatch(): Unit = { + exception.compareAndSet(null, new IOException("the stock Celeborn push failed")) + } + + def cleanup(): Unit = { + exception.compareAndSet(null, new IOException("Cleaned Up")) + } + + def checkFailure(): Unit = Option(exception.get()).foreach(throw _) +} + final case class RecordedCelebornPush( shuffleId: Int, mapId: Int, @@ -341,7 +424,7 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { private val pluginKey = "spark.shuffle.sort.io.plugin.class" private val pluginClass = "org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO" private val endpointsKey = "spark.celeborn.master.endpoints" - private val enabledKey = "spark.comet.celeborn.enabled" + private val enabledKey = CometConf.COMET_SHUFFLE_CELEBORN_ENABLED.key private def enabledConf: SparkConf = new SparkConf(false).set(endpointsKey, "celeborn-master:9097") @@ -374,6 +457,112 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { assert(client.observedMapKey == s"19-3-${(4 << 16) | 7}") } + test("stock Apache Celeborn commits with its reducer-aware mapperEnd API") { + val client = new StockCelebornShuffleClient + val adapter = new CelebornShufflePartitionPusher(client, 19, 3, 7, 12, 9) + + assert(adapter.pushPartitionData(2, Array[Byte](1, 2, 3), 3) == 3) + assert(adapter.finish().sameElements(Array[Long](0, 0, 3, 0, 0, 0, 0, 0, 0))) + assert(client.lastMapperEnd == ((19, 3, 7, 12, 9))) + } + + test("stock Apache Celeborn restores byte admission without a push metrics callback") { + val client = new StockCelebornShuffleClient + client.automaticallyCompletePushes = false + val first = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val second = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val failure = new AtomicReference[Throwable]() + val bytes = Array.fill[Byte](32)(1) + + assert(first.pushPartitionData(0, bytes, bytes.length) == bytes.length) + val worker = new Thread(() => { + try second.pushPartitionData(0, bytes, bytes.length) + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + + Thread.sleep(100) + assert(worker.isAlive) + assert(client.pushCalls.get() == 1) + + assert(client.completeNextPush()) + worker.join(5000) + + if (worker.isAlive) { + second.abort() + worker.join(5000) + fail("stock Celeborn transport completion did not restore executor-wide admission") + } + assert(failure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + + test("stock Apache Celeborn terminal push failures restore shared byte admission once") { + val client = new StockCelebornShuffleClient + client.automaticallyCompletePushes = false + val failed = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val replacementFailure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + assert(failed.pushPartitionData(0, frame, frame.length) == frame.length) + val worker = new Thread(() => { + try replacement.pushPartitionData(0, frame, frame.length) + catch { case failure: Throwable => replacementFailure.set(failure) } + }) + worker.start() + + Thread.sleep(100) + assert(worker.isAlive) + assert(client.failNextPush()) + val failure = intercept[IOException](failed.finish()) + assert(failure.getMessage.contains("stock Celeborn push failed")) + worker.join(5000) + + if (worker.isAlive) { + replacement.abort() + worker.join(5000) + fail("a stock Celeborn terminal push failure permanently retained byte admission") + } + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + + test("stock Apache Celeborn cancellation does not release a live transport request") { + val client = new StockCelebornShuffleClient + client.automaticallyCompletePushes = false + val cancelled = new CelebornShufflePartitionPusher(client, 19, 3, 1, 12, 9, 48) + val replacement = new CelebornShufflePartitionPusher(client, 19, 4, 1, 12, 9, 48) + val replacementFailure = new AtomicReference[Throwable]() + val frame = Array.fill[Byte](32)(1) + + assert(cancelled.pushPartitionData(0, frame, frame.length) == frame.length) + cancelled.abort() + + val worker = new Thread(() => { + try replacement.pushPartitionData(0, frame, frame.length) + catch { case failure: Throwable => replacementFailure.set(failure) } + }) + worker.start() + + Thread.sleep(100) + assert(worker.isAlive) + assert(client.pushCalls.get() == 1) + assert(client.completeNextPush()) + worker.join(5000) + + if (worker.isAlive) { + replacement.abort() + worker.join(5000) + fail("a completed stock Celeborn request permanently retained byte admission") + } + assert(replacementFailure.get() == null) + assert(client.pushCalls.get() == 2) + assert(client.completeNextPush()) + } + test("raw Celeborn push requires exactly the payload plus its transport header") { val bytes = Array[Byte](1, 2, 3) @@ -1455,6 +1644,15 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { assert(!CelebornShufflePusherFactory.isEnabled(conf)) } + test("factory disables Celeborn native shuffle while Spark I/O encryption is enabled") { + val conf = enabledConf.set("spark.io.encryption.enabled", "true") + + assert(!CelebornShufflePusherFactory.isEnabled(conf)) + + conf.set("spark.io.encryption.enabled", "false") + assert(CelebornShufflePusherFactory.isEnabled(conf)) + } + test("an acquired client remains owned when Celeborn shuffle-generation resolution fails") { val client = new RecordingCelebornShuffleClient val expected = new IOException("Celeborn shuffle-generation lookup failed") 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 79f7c07b2e8..4c3274338b2 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 @@ -527,6 +527,42 @@ class CometCelebornNativeShuffleWriterSuite extends CometTestBase { } } + test("task failure proactively cleans up a blocked Celeborn native push") { + val client = new RecordingCelebornShuffleClient + client.pushStarted = new CountDownLatch(1) + client.allowPush = new CountDownLatch(1) + client.cleanupUnblocksPush = true + val taskContext = TaskContext.empty() + val pusher = + new org.apache.comet.shuffle.CelebornShufflePartitionPusher(client, 93, 0, 0, 1, 1) + val callbackFailure = new AtomicReference[Throwable]() + val cleanupFailure = new AtomicReference[Throwable]() + val watcher = CelebornNativeShuffleDestination.watchForCancellation( + taskContext, + pusher, + failure => cleanupFailure.set(failure)) + + try { + val worker = new Thread(() => { + try pusher.pushPartitionData(0, Array[Byte](1), 1) + catch { case failure: Throwable => callbackFailure.set(failure) } + }) + worker.start() + assert(client.pushStarted.await(5, TimeUnit.SECONDS)) + + taskContext.markTaskFailed(new IOException("the Spark task failed")) + worker.join(5000) + + assert(!worker.isAlive) + assert(callbackFailure.get().isInstanceOf[IOException]) + assert(cleanupFailure.get() == null) + assert(client.cleanupCalls.get() == 2) + } finally { + watcher.cancel(false) + pusher.abort() + } + } + test("task interruption proactively wakes blocked Celeborn map completion") { val client = new RecordingCelebornShuffleClient client.mapperEndStarted = new CountDownLatch(1) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index fcd93e53e5a..56bcd06db9b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -476,6 +476,27 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { context.markTaskCompleted(None) } + test("raw fetch prefers Apache Celeborn's public 15-argument partition reader") { + val context = TaskContext.empty() + val client = new RecordingCelebornRawClient.StockApiClient + client.fileGroups.partitionGroups.put(1, location("worker")) + client.fileGroups.mapAttempts = Array(2, 4) + client.streams.put(1, new ByteArrayInputStream(Array[Byte](3, 5))) + + val input = rawReader(client, context, startPartition = 1, endPartition = 2).openPartitions() + assert(input.readAllBytes().toSeq == Seq[Byte](3, 5)) + assert(client.stockReadPartitionCalls == 1) + assert(client.requests.size() == 1) + val request = client.requests.get(0) + assert(request.shuffleId == 91) + assert(request.appShuffleId == 17) + assert(request.coalescedPartitionInfos == null) + assert(request.mapAttempts eq client.fileGroups.mapAttempts) + assert(!request.needDecompress) + input.close() + context.markTaskCompleted(None) + } + test("raw fetch forwards Celeborn byte, block, and wait metrics") { val context = TaskContext.empty() val client = new RecordingCelebornRawClient From 48b069e4deea41a65ae154c66453e285b3b2d3f0 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 26 Aug 2026 00:09:45 +0000 Subject: [PATCH 11/12] ci: register Celeborn shuffle suites in PR build matrices --- .github/workflows/pr_build_linux.yml | 4 ++++ .github/workflows/pr_build_macos.yml | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index b7c88254c93..4743fed1913 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -327,6 +327,10 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 8210f91b7f9..a9d8f4a0bc8 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -143,6 +143,10 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite From f2e14b3f48f7290cf13e64e1db90dc3736c036ce Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 26 Aug 2026 00:23:21 +0000 Subject: [PATCH 12/12] fix: satisfy Scala lint and Rust Clippy checks for Celeborn shuffle --- native/shuffle/tests/shuffle_destinations.rs | 6 ++++-- .../apache/comet/shuffle/CelebornShufflePusherFactory.scala | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/native/shuffle/tests/shuffle_destinations.rs b/native/shuffle/tests/shuffle_destinations.rs index c19e55c8f50..1888811d488 100644 --- a/native/shuffle/tests/shuffle_destinations.rs +++ b/native/shuffle/tests/shuffle_destinations.rs @@ -206,8 +206,10 @@ fn read_local(output_paths: &(String, String), num_partitions: usize) -> Vec = index - .chunks_exact(8) - .map(|bytes| u64::from_le_bytes(bytes.try_into().unwrap()) as usize) + .as_chunks::<8>() + .0 + .iter() + .map(|bytes| u64::from_le_bytes(*bytes) as usize) .collect(); assert_eq!(offsets[0], 0); assert_eq!(offsets[num_partitions], data.len()); diff --git a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala index 5e3ca25c0d1..51f840fa881 100644 --- a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala +++ b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala @@ -320,7 +320,7 @@ object CelebornShufflePusherFactory { Long.box(-1L), Int.box(-1), Int.box(-1), - s"Retried Celeborn map attempt requires a new shuffle generation: " + + "Retried Celeborn map attempt requires a new shuffle generation: " + s"$sparkShuffleId/$celebornShuffleId", null) .asInstanceOf[Throwable]