You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by co...@apache.org on 2026/09/05 15:46:57 UTC
(datafusion-comet) branch main updated: fix: recover native Celeborn shuffle from oversized rows (#5668)
This is an automated email from the ASF dual-hosted git repository.
comphead pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 1c38a3ee2a fix: recover native Celeborn shuffle from oversized rows (#5668)
1c38a3ee2a is described below
commit 1c38a3ee2aaf1f146a73e88072274637064107fe
Author: Ping Zhang <pi...@gmail.com>
AuthorDate: Sat Sep 5 08:46:52 2026 -0700
fix: recover native Celeborn shuffle from oversized rows (#5668)
---
.github/workflows/pr_build_linux.yml | 4 +
.github/workflows/pr_build_macos.yml | 4 +
.../comet/CometShuffleSizeLimitException.java | 17 +-
docs/source/user-guide/latest/configs.md | 19 +
docs/source/user-guide/latest/tuning.md | 27 +-
native/jni-bridge/src/errors.rs | 63 +-
native/shuffle/src/writers/rss/mod.rs | 151 +++-
.../src/writers/rss/rss_partition_writer.rs | 80 +-
.../shuffle/CelebornShufflePartitionPusher.java | 2 +-
.../main/scala/org/apache/comet/CometConf.scala | 7 +-
.../shuffle/CometCelebornShuffleManager.scala | 254 ++++++-
.../CometCelebornShuffleMaterialization.scala | 359 +++++++++
.../shuffle/CometNativeShuffleInputRDD.scala | 14 +
.../shuffle/CometNativeShuffleWriter.scala | 18 +-
.../execution/shuffle/CometShuffleDependency.scala | 66 +-
.../shuffle/CometShuffleExchangeExec.scala | 31 +-
.../shuffle/CometShuffleOutputMetrics.scala | 69 ++
.../execution/shuffle/CometShuffledRowRDD.scala | 13 +-
.../shims/ShimCometShuffleMaterialization.scala} | 32 +-
.../shims/ShimCometShuffleMaterialization.scala | 47 ++
.../shims/ShimCometShuffleMaterialization.scala | 50 ++
.../scala/org/apache/comet/CometConfSuite.scala | 2 +-
.../comet/exec/CometNativeShuffleSuite.scala | 9 +-
.../shuffle/celeborn/CelebornShuffleHandle.scala | 4 +-
...metCelebornConcurrentMaterializationSuite.scala | 338 +++++++++
.../CometCelebornLocalFetchFailureSuite.scala | 139 ++++
.../CometCelebornNativeShuffleWriterSuite.scala | 301 +++++++-
.../CometCelebornShuffleFallbackSuite.scala | 808 +++++++++++++++++++++
.../shuffle/CometCelebornShuffleManagerSuite.scala | 240 ++++++
.../shuffle/CometCelebornShuffleReaderSuite.scala | 60 +-
.../CometCelebornShuffleStatisticsSuite.scala | 132 ++++
.../shuffle/CometNativeShuffleInputRDDSuite.scala | 69 +-
32 files changed, 3336 insertions(+), 93 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 5b8ed1f307..79fea7fb88 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -375,6 +375,10 @@ jobs:
org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleFallbackSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornConcurrentMaterializationSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleStatisticsSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornLocalFetchFailureSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index bbb0563f2f..5473d3988a 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -148,6 +148,10 @@ jobs:
org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleFallbackSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornConcurrentMaterializationSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleStatisticsSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornLocalFetchFailureSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
diff --git a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala b/common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java
similarity index 64%
copy from spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala
copy to common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java
index 6104c320b9..ff257fb0fb 100644
--- a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala
+++ b/common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java
@@ -17,14 +17,11 @@
* under the License.
*/
-package org.apache.spark.shuffle.celeborn
+package org.apache.comet;
-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)
+/** A remote shuffle frame or its encoding workspace cannot fit the configured byte limits. */
+public final class CometShuffleSizeLimitException extends CometNativeException {
+ public CometShuffleSizeLimitException(String message) {
+ super(message);
+ }
+}
diff --git a/docs/source/user-guide/latest/configs.md b/docs/source/user-guide/latest/configs.md
index a268691a3c..49ee6dd48d 100644
--- a/docs/source/user-guide/latest/configs.md
+++ b/docs/source/user-guide/latest/configs.md
@@ -45,6 +45,25 @@ These settings can be used to determine which parts of the plan are accelerated
## Shuffle Configuration Settings
+For native remote shuffle, `spark.comet.shuffle.rss.maxFrameBytes` limits one complete
+encoded frame, while `spark.comet.shuffle.rss.maxInFlightBytes` limits the memory reserved
+by map attempts sharing an executor's remote shuffle client. The reservation includes
+encoding workspace and overlapping frame copies. An ordinary uncompressed frame needs
+approximately seven times its size plus schema and transport overhead. The default 512 MiB
+reservation budget accommodates ordinary frames up to the default 64 MiB frame limit.
+Compressed frames still need workspace for their uncompressed data. Increase the reservation
+budget when larger rows or schemas need more workspace.
+
+If a row cannot fit the remote limits, Comet materializes a replacement shuffle using its
+local writer before publishing the exchange to downstream tasks. The replacement has a separate
+shuffle and scheduling identity, so late remote results or failures cannot affect local output.
+Independent exchanges can materialize concurrently, and runtime output statistics describe only
+the selected destination. Subsequent fetch failures retain Spark's normal recovery behavior.
+With dynamic allocation enabled, native Celeborn shuffle also requires either
+`spark.shuffle.service.enabled=true` or `spark.dynamicAllocation.shuffleTracking.enabled=true`
+(the Spark default) to preserve fallback files. Otherwise exchanges retain ordinary Spark/Celeborn
+shuffle, including applications that rely only on remote reliable storage or decommissioning.
+
<!--BEGIN:CONFIG_TABLE[shuffle]-->
<!--END:CONFIG_TABLE-->
diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md
index decb47a8c2..0ad8ec992c 100644
--- a/docs/source/user-guide/latest/tuning.md
+++ b/docs/source/user-guide/latest/tuning.md
@@ -257,7 +257,7 @@ even when both its parent and child are non-Comet operators.
Applications using Apache Celeborn can use Comet's composite shuffle manager to retain ordinary
Spark/Celeborn shuffle while accelerating other operators with Comet.
-Native shuffle requires reliable completion tracking for in-flight payloads. Released Celeborn
+Native shuffle also requires reliable completion tracking for in-flight payloads. Released Celeborn
0.6.0 and 0.7.0 clients do not provide the required guarantee, so these versions retain ordinary
Spark/Celeborn shuffle even when `spark.comet.shuffle.mode=native`. Native shuffle support for
these clients requires a safe Celeborn push-completion API. The following settings request
@@ -302,8 +302,29 @@ Celeborn to prohibit local fallback for ordinary Spark shuffles.
Native frames retain Comet's configured compression; the raw Celeborn client path bypasses
Celeborn's additional row compression and decompression. Use
`spark.comet.shuffle.rss.maxFrameBytes` and `spark.comet.shuffle.rss.maxInFlightBytes` to bound
-encoded frame size and executor-side push admission. These limits include framing and overlapping
-native/JNI/client copies; a frame that cannot fit is rejected rather than split across requests.
+encoded frame size and executor-side push admission. The defaults are 64 MiB and 512 MiB,
+respectively. Admission includes Arrow encoding workspace as well as overlapping native, JNI,
+and client frame copies. An uncompressed frame needs roughly seven times its size plus schema
+and codec overhead. Compression reduces the transmitted bytes but still needs uncompressed
+encoding workspace.
+
+Comet splits large batches between rows. If a single row, its schema, or its encoding workspace
+cannot fit the remote limits, Comet abandons the remote shuffle and materializes a replacement
+using its local shuffle writer before downstream tasks can consume the exchange. The replacement
+has a separate shuffle and scheduling identity, so late remote results cannot overwrite or skip
+local map output, and remote stage failures cannot abort the replacement. Independent exchanges
+can materialize concurrently; readers wait for their storage decisions before execution. Runtime
+output statistics count only the selected destination. All reads and retries for the replacement
+use local files and Spark's block transfer
+service, including normal recovery after later fetch failures. Native operators and Comet's
+Arrow shuffle format are preserved, and remote admission limits remain enforced. Once remote
+output has been published, subsequent failures use the existing Spark/Celeborn recovery path;
+Comet does not change that shuffle's destination. Local fallback uses executor disk. When `spark.dynamicAllocation.enabled=true`, native Celeborn shuffle requires
+`spark.shuffle.service.enabled=true` or `spark.dynamicAllocation.shuffleTracking.enabled=true`
+(the Spark default) so those files remain available. Applications using dynamic allocation with
+both settings disabled retain ordinary Spark/Celeborn shuffle, even if remote reliable storage or
+decommissioning enables dynamic allocation. Executor shutdown preserves fallback files for the
+external shuffle service; explicit shuffle unregister retains the normal local cleanup behavior.
AQE reducer coalescing and mapper-range reads are supported, but Celeborn physical-skew chunk reads
are not.
diff --git a/native/jni-bridge/src/errors.rs b/native/jni-bridge/src/errors.rs
index 4b8cb439a1..a072be95ed 100644
--- a/native/jni-bridge/src/errors.rs
+++ b/native/jni-bridge/src/errors.rs
@@ -90,6 +90,11 @@ pub enum CometError {
#[error("Comet Internal Error: {0}")]
Internal(String),
+ /// A remote shuffle frame or its encoding workspace cannot fit the configured limits.
+ /// Preserve this classification so Spark can restart the shuffle with a local writer.
+ #[error("{0}")]
+ ShuffleSizeLimit(String),
+
#[error(transparent)]
Arrow {
#[from]
@@ -214,7 +219,9 @@ impl From<CometError> for DataFusionError {
// own codegen inside the JVM UDF kernel) as an `External` error so it survives the trip
// back through DataFusion and can be re-thrown with its exact type at the JNI boundary.
// Flattening it to a string here would surface it as a generic CometNativeException.
- value @ CometError::JavaException { .. } => DataFusionError::External(Box::new(value)),
+ value @ (CometError::JavaException { .. } | CometError::ShuffleSizeLimit(_)) => {
+ DataFusionError::External(Box::new(value))
+ }
_ => DataFusionError::Execution(value.to_string()),
}
}
@@ -338,6 +345,10 @@ impl jni::errors::ToException for CometError {
class: spark_err.exception_class().to_string(),
msg: spark_err.to_string(),
},
+ CometError::ShuffleSizeLimit(message) => Exception {
+ class: "org/apache/comet/CometShuffleSizeLimitException".to_string(),
+ msg: message.clone(),
+ },
_other => Exception {
class: "org/apache/comet/CometNativeException".to_string(),
msg: self.to_string(),
@@ -473,6 +484,15 @@ pub fn unwrap_or_throw_default<T: JNIDefault>(
fn throw_exception(env: &mut Env, error: &CometError, backtrace: Option<String>) {
// If there isn't already an exception?
if !env.exception_check() {
+ // DataFusion operators can wrap the original failure in Context, Shared, or External
+ // errors. Keep capacity failures typed across those wrappers and the JNI boundary.
+ if let Some(message) = shuffle_size_limit_message(error) {
+ let _ = env.throw_new(
+ jni::jni_str!("org/apache/comet/CometShuffleSizeLimitException"),
+ JNIString::new(message),
+ );
+ return;
+ }
// ... then throw new exception
// Note: in jni 0.22.x, throw/throw_new return Err(JavaException) on success
// (to signal the pending exception to Rust callers via `?`). We discard the
@@ -554,6 +574,17 @@ fn throw_exception(env: &mut Env, error: &CometError, backtrace: Option<String>)
}
}
+fn shuffle_size_limit_message<'a>(error: &'a (dyn std::error::Error + 'static)) -> Option<&'a str> {
+ let mut cause = Some(error);
+ while let Some(error) = cause {
+ if let Some(CometError::ShuffleSizeLimit(message)) = error.downcast_ref::<CometError>() {
+ return Some(message);
+ }
+ cause = error.source();
+ }
+ None
+}
+
/// Generic fallback throw for an error that isn't a structured `SparkError`. Recognises a
/// file-not-found arriving through non-typed wrapping paths and duplicate-field errors; otherwise
/// throws the error's natural JVM exception (with the captured backtrace when available).
@@ -912,6 +943,36 @@ mod tests {
}
}
+ #[test]
+ #[cfg_attr(miri, ignore)] // miri cannot create a JVM.
+ fn shuffle_size_limit_survives_datafusion_wrappers_and_jni() {
+ let message = "Remote shuffle exceeds spark.comet.shuffle.rss.maxInFlightBytes";
+ let error = DataFusionError::from(CometError::ShuffleSizeLimit(message.to_string()));
+ let error = DataFusionError::Shared(Arc::new(DataFusionError::Context(
+ "executing shuffle writer".to_string(),
+ Box::new(error),
+ )));
+ jvm()
+ .attach_current_thread(|env| -> jni::errors::Result<()> {
+ unwrap_or_throw_default::<()>(env, Err(CometError::from(error)));
+ assert_pending_java_exception_detailed(
+ env,
+ Some("org/apache/comet/CometShuffleSizeLimitException"),
+ Some(message),
+ );
+ Ok(())
+ })
+ .unwrap();
+ }
+
+ #[test]
+ fn shuffle_size_limit_is_not_inferred_from_error_text() {
+ let error = CometError::from(DataFusionError::Execution(
+ "Remote shuffle exceeds spark.comet.shuffle.rss.maxInFlightBytes".to_string(),
+ ));
+ assert!(shuffle_size_limit_message(&error).is_none());
+ }
+
#[test]
#[cfg_attr(miri, ignore)] // miri can't call foreign function `dlopen`
pub fn error_from_panic() {
diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs
index a151b6831f..bf7cbaa464 100644
--- a/native/shuffle/src/writers/rss/mod.rs
+++ b/native/shuffle/src/writers/rss/mod.rs
@@ -33,6 +33,7 @@ mod tests {
use arrow::record_batch::RecordBatch;
use datafusion::common::{DataFusionError, Result};
use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time};
+ use datafusion_comet_jni_bridge::errors::CometError;
use datafusion_comet_jni_bridge::ShufflePartitionPusher;
use std::collections::HashMap;
use std::io::{self, Cursor};
@@ -411,6 +412,28 @@ mod tests {
.unwrap()
}
+ fn single_string_row(bytes: usize) -> RecordBatch {
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "value",
+ DataType::Utf8,
+ false,
+ )])),
+ vec![Arc::new(StringArray::from(vec!["x".repeat(bytes)]))],
+ )
+ .unwrap()
+ }
+
+ fn size_limit_message(error: &DataFusionError) -> &str {
+ let DataFusionError::External(cause) = error else {
+ panic!("expected a typed shuffle size limit, got {error}");
+ };
+ let Some(CometError::ShuffleSizeLimit(message)) = cause.downcast_ref::<CometError>() else {
+ panic!("expected a typed shuffle size limit, got {error}");
+ };
+ message
+ }
+
#[test]
#[cfg_attr(miri, ignore)]
fn round_trips_all_supported_compression_codecs() {
@@ -633,7 +656,15 @@ mod tests {
frame_size - 1,
);
- assert!(write_batches(&mut writer, 0, vec![batch], &metrics()).is_err());
+ let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err();
+ let message = size_limit_message(&error);
+ assert!(message.contains(&format!(
+ "effective frame limit from spark.comet.shuffle.rss.maxFrameBytes \
+ and spark.comet.shuffle.rss.maxInFlightBytes is {} bytes",
+ frame_size - 1
+ )));
+ assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes"));
+ assert!(message.contains(&format!("frame needs at least {frame_size} bytes")));
assert!(
pusher.frames().is_empty(),
"a single oversized Arrow IPC row must never be fragmented"
@@ -786,7 +817,7 @@ mod tests {
#[test]
fn small_batches_encode_concurrently_with_default_frame_and_admission_limits() {
let frame_limit = 64 * 1024 * 1024;
- let admission_limit = 256 * 1024 * 1024;
+ let admission_limit = 512 * 1024 * 1024;
let pusher = Arc::new(ConcurrentAdmissionPusher::new(admission_limit));
thread::scope(|scope| {
@@ -815,6 +846,115 @@ mod tests {
assert_eq!(pusher.state.lock().unwrap().total, 0);
}
+ #[test]
+ #[cfg_attr(miri, ignore)]
+ fn large_single_rows_fit_default_frame_and_admission_limits() {
+ let frame_limit = 64 * 1024 * 1024;
+ // Celeborn reserves its 16-byte request header separately from the native reservation.
+ let reservation_limit = 512 * 1024 * 1024 - 16;
+ for row_size_mib in [38, 50, 63] {
+ let batch = single_string_row(row_size_mib * 1024 * 1024);
+ let pusher = Arc::new(ReservationRecordingPusher {
+ max_reservation: Some(reservation_limit),
+ capacity: Some(reservation_limit),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(
+ &batch,
+ CompressionCodec::None,
+ pusher.clone(),
+ 1,
+ frame_limit,
+ );
+ let metrics = metrics();
+
+ finish_partition(&mut writer, 0, vec![batch.clone()], &metrics).unwrap();
+ writer.finish_all(&metrics).unwrap();
+
+ let frames = pusher.frames.lock().unwrap();
+ assert_eq!(frames.len(), 1);
+ assert!(frames[0].1.len() <= frame_limit);
+ assert_eq!(decode_frame(&frames[0].1), batch);
+ assert!(pusher.outstanding.lock().unwrap().is_none());
+ assert!(pusher
+ .reservations
+ .lock()
+ .unwrap()
+ .iter()
+ .all(|bytes| *bytes <= reservation_limit));
+ }
+ }
+
+ #[test]
+ #[cfg_attr(miri, ignore)]
+ fn reports_admission_limit_when_a_single_row_fits_the_configured_frame_limit() {
+ let batch = single_string_row(38 * 1024 * 1024);
+ let frame_limit = 64 * 1024 * 1024;
+ let reservation_limit = 256 * 1024 * 1024 - 16;
+ let pusher = Arc::new(ReservationRecordingPusher {
+ max_reservation: Some(reservation_limit),
+ capacity: Some(reservation_limit),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(
+ &batch,
+ CompressionCodec::None,
+ pusher.clone(),
+ 1,
+ frame_limit,
+ );
+ assert!(encoded_frame_size(&batch) < frame_limit);
+
+ let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err();
+
+ let message = size_limit_message(&error);
+ assert!(
+ message.contains("effective frame limit from spark.comet.shuffle.rss.maxFrameBytes")
+ );
+ assert!(message.contains(&format!(": {frame_limit} bytes)")));
+ assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes"));
+ assert!(message.contains(&format!("{reservation_limit}-byte reservation")));
+ assert!(message.contains("encoded bytes fit alongside"));
+ assert!(message.contains("required reservation is at least"));
+ assert!(pusher.frames.lock().unwrap().is_empty());
+ assert!(pusher.outstanding.lock().unwrap().is_none());
+ }
+
+ #[test]
+ #[cfg_attr(miri, ignore)]
+ fn reports_workspace_limit_before_encoding_a_single_row() {
+ let batch = single_string_row(20 * 1024 * 1024);
+ let frame_limit = 64 * 1024 * 1024;
+ let reservation_limit = 64 * 1024 * 1024 - 16;
+ let pusher = Arc::new(ReservationRecordingPusher {
+ max_reservation: Some(reservation_limit),
+ capacity: Some(reservation_limit),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(
+ &batch,
+ CompressionCodec::None,
+ pusher.clone(),
+ 1,
+ frame_limit,
+ );
+
+ let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err();
+
+ let message = size_limit_message(&error);
+ assert!(message.contains("encoding workspace for a single row requires"));
+ assert!(message.contains("bytes of reservation are needed"));
+ assert!(message.contains(&format!("only {reservation_limit} bytes per frame")));
+ assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes"));
+ assert!(
+ message.contains("effective frame limit from spark.comet.shuffle.rss.maxFrameBytes")
+ );
+ assert!(message.contains(&format!("is {frame_limit} bytes")));
+ assert_eq!(pusher.reservations.lock().unwrap().len(), 1);
+ assert!(pusher.frames.lock().unwrap().is_empty());
+ assert!(pusher.outstanding.lock().unwrap().is_none());
+ }
+
#[test]
fn codec_workspace_is_rejected_before_encoding_under_a_tiny_budget() {
let batch = sample_batch(0, 1);
@@ -830,7 +970,12 @@ mod tests {
});
let mut writer = writer(&batch, codec, pusher.clone(), 1, 5_456);
let error = write_batches(&mut writer, 0, vec![batch.clone()], &metrics()).unwrap_err();
- assert!(error.to_string().contains("encoding workspace"));
+ let message = size_limit_message(&error);
+ assert!(message.contains("encoding workspace"));
+ assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes"));
+ assert!(message
+ .contains("effective frame limit from spark.comet.shuffle.rss.maxFrameBytes"));
+ assert!(message.contains("is 5456 bytes"));
assert!(pusher.reservations.lock().unwrap().is_empty());
assert!(pusher.frames.lock().unwrap().is_empty());
}
diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs
index 76503011db..a7e374f702 100644
--- a/native/shuffle/src/writers/rss/rss_partition_writer.rs
+++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs
@@ -29,6 +29,7 @@ use arrow::datatypes::{DataType, Field, Int16Type, Int32Type, Int64Type};
use arrow::ipc::writer::CompressionContext;
use arrow_select::dictionary::garbage_collect_any_dictionary;
use datafusion::common::{DataFusionError, Result};
+use datafusion_comet_jni_bridge::errors::CometError;
use datafusion_comet_jni_bridge::ShufflePartitionPusher;
use std::io::{self, Cursor, Seek, SeekFrom, Write};
use std::sync::Arc;
@@ -160,14 +161,19 @@ impl RssPartitionWriter {
let (metadata_scratch, planning_scratch) = Self::estimated_ipc_metadata_scratch(batch);
let codec_workspace = self.block_writer.rss_codec_workspace()?;
let reservation_limit = self.pusher.max_reservation_size();
- if metadata_scratch
+ let minimum_reservation = metadata_scratch
.saturating_add(codec_workspace)
- .saturating_add(60)
- > reservation_limit
- {
+ .saturating_add(60);
+ if minimum_reservation > reservation_limit {
// Neither schema descriptors nor codec workspace gets smaller when rows are split.
- return Err(DataFusionError::Execution(format!(
- "Remote shuffle frame schema and encoding workspace exceed the byte admission budget of {reservation_limit} bytes"
+ return Err(Self::size_limit_error(format!(
+ "Remote shuffle schema, encoding workspace, and minimum frame copies require at least \
+ {minimum_reservation} bytes of reservation, but only {reservation_limit} \
+ bytes per frame are available under spark.comet.shuffle.rss.maxInFlightBytes. \
+ Increase that setting to allow the required reservation plus transport overhead. \
+ The effective frame limit from spark.comet.shuffle.rss.maxFrameBytes and \
+ spark.comet.shuffle.rss.maxInFlightBytes is {} bytes",
+ self.max_frame_size
)));
}
@@ -199,7 +205,7 @@ impl RssPartitionWriter {
.and_then(|bytes| bytes.checked_add(compaction_scratch))
.and_then(|bytes| bytes.checked_add(codec_workspace))
.ok_or_else(|| {
- DataFusionError::Execution(
+ Self::size_limit_error(
"Remote shuffle encoding workspace exceeds the native integer limit"
.to_string(),
)
@@ -210,8 +216,16 @@ impl RssPartitionWriter {
if batch.num_rows() > 1 {
return self.push_split_batch(partition_id, batch, metrics);
}
- return Err(DataFusionError::Execution(format!(
- "Remote shuffle frame for a single row and encoding workspace exceed the byte admission budget of {reservation_limit} bytes"
+ return Err(Self::size_limit_error(format!(
+ "Remote shuffle encoding workspace for a single row requires {ipc_scratch} bytes; \
+ at least {} bytes of reservation are needed including a minimum frame, but only \
+ {reservation_limit} bytes per frame are available under \
+ spark.comet.shuffle.rss.maxInFlightBytes. Increase that setting to allow the \
+ required reservation plus transport overhead. \
+ The effective frame limit from spark.comet.shuffle.rss.maxFrameBytes and \
+ spark.comet.shuffle.rss.maxInFlightBytes is {} bytes",
+ ipc_scratch.saturating_add(60),
+ self.max_frame_size
)));
}
@@ -231,7 +245,7 @@ impl RssPartitionWriter {
.checked_mul(3)
.and_then(|bytes| bytes.checked_add(ipc_scratch))
.ok_or_else(|| {
- DataFusionError::Execution(
+ Self::size_limit_error(
"Remote shuffle frame-copy reservation exceeds the native integer limit"
.to_string(),
)
@@ -277,6 +291,7 @@ impl RssPartitionWriter {
&metrics.encode_time,
) {
let exceeded = output.exceeded;
+ let minimum_frame_size = output.minimum_size;
drop(output);
drop(compacted);
self.pusher.release_partition_data_reservation()?;
@@ -288,9 +303,32 @@ impl RssPartitionWriter {
continue;
}
if batch.num_rows() <= 1 {
- return Err(DataFusionError::Execution(format!(
- "Remote shuffle frame exceeds its configured maximum: a single row exceeds {} bytes",
- admitted_frame_limit
+ let minimum_reservation =
+ ipc_scratch.saturating_add(minimum_frame_size.saturating_mul(3));
+ let reason = if admitted_frame_limit < self.max_frame_size {
+ format!(
+ "only {admitted_frame_limit} encoded bytes fit alongside the \
+ {ipc_scratch}-byte encoding workspace in the {reservation_limit}-byte \
+ reservation available per frame under \
+ spark.comet.shuffle.rss.maxInFlightBytes \
+ (effective frame limit from spark.comet.shuffle.rss.maxFrameBytes \
+ and spark.comet.shuffle.rss.maxInFlightBytes: {} bytes)",
+ self.max_frame_size
+ )
+ } else {
+ format!(
+ "the effective frame limit from spark.comet.shuffle.rss.maxFrameBytes \
+ and spark.comet.shuffle.rss.maxInFlightBytes is {} bytes; encoding \
+ workspace requires {ipc_scratch} bytes and the reservation available \
+ per frame is {reservation_limit} bytes",
+ self.max_frame_size
+ )
+ };
+ return Err(Self::size_limit_error(format!(
+ "Remote shuffle cannot encode a single row: the frame needs at least \
+ {minimum_frame_size} bytes, but {reason}. Increase the limiting setting; \
+ the required reservation is at least {minimum_reservation} bytes plus \
+ transport overhead"
)));
}
return self.push_split_batch(partition_id, batch, metrics);
@@ -316,6 +354,10 @@ impl RssPartitionWriter {
}
}
+ fn size_limit_error(message: String) -> DataFusionError {
+ DataFusionError::External(Box::new(CometError::ShuffleSizeLimit(message)))
+ }
+
fn push_split_batch(
&mut self,
partition_id: i32,
@@ -1092,6 +1134,8 @@ struct BoundedBuffer {
inner: Cursor<Vec<u8>>,
limit: usize,
exceeded: bool,
+ // The first rejected write/seek establishes a lower bound without allocating more output.
+ minimum_size: usize,
}
impl BoundedBuffer {
@@ -1113,6 +1157,7 @@ impl BoundedBuffer {
inner: Cursor::new(bytes),
limit,
exceeded: false,
+ minimum_size: 0,
})
}
@@ -1128,6 +1173,11 @@ impl Write for BoundedBuffer {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
let end = self.inner.position().checked_add(bytes.len() as u64);
if end.is_none_or(|end| end > self.limit as u64) {
+ if !self.exceeded {
+ self.minimum_size = end
+ .and_then(|end| usize::try_from(end).ok())
+ .unwrap_or(usize::MAX);
+ }
self.exceeded = true;
return Err(Self::limit_error());
}
@@ -1145,6 +1195,9 @@ impl Seek for BoundedBuffer {
let next = self.inner.seek(position)?;
if next > self.limit as u64 {
self.inner.set_position(previous);
+ if !self.exceeded {
+ self.minimum_size = usize::try_from(next).unwrap_or(usize::MAX);
+ }
self.exceeded = true;
return Err(Self::limit_error());
}
@@ -1186,6 +1239,7 @@ mod buffer_tests {
assert_eq!(output.inner.get_ref().capacity(), 67_996);
assert!(output.write(&[1]).is_err());
assert!(output.exceeded);
+ assert_eq!(output.minimum_size, 67_997);
assert_eq!(output.inner.get_ref().capacity(), 67_996);
}
}
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 4240755412..c634b28e85 100644
--- a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java
+++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java
@@ -45,7 +45,7 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus
private static final int CELEBORN_BATCH_HEADER_BYTES = 4 * Integer.BYTES;
private static final int MINIMUM_COMET_FRAME_BYTES = 2 * Long.BYTES;
private static final int MAX_JVM_ARRAY_BYTES = Integer.MAX_VALUE - 8;
- private static final int DEFAULT_MAX_IN_FLIGHT_BYTES = 256 * 1024 * 1024;
+ private static final int DEFAULT_MAX_IN_FLIGHT_BYTES = 512 * 1024 * 1024;
private static final long RECONCILIATION_INTERVAL_MILLIS = 10;
// This daemon never owns client state: reconciliations are cancelled and removed once the
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 401adaf5ca..cd75909a91 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -659,14 +659,17 @@ object CometConf extends ShimCometConf {
"an executor-side remote shuffle client. Admission includes native encoding " +
"scratch and overlapping native, JNI, and remote shuffle frame copies. " +
"A frame must fit its codec and Arrow workspace as well as its encoded bytes; " +
- "too-small limits fail before encoding. Encrypted native RSS is not supported; " +
+ "ordinary uncompressed frames need approximately seven times their size plus " +
+ "schema and transport overhead. Compressed frames also reserve workspace for " +
+ "their uncompressed data. Admission is acquired before encoding. " +
+ "Encrypted native RSS is not supported; " +
"use ordinary Spark shuffle when spark.io.encryption.enabled is true.")
.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)
+ .createWithDefault(512L * 1024 * 1024)
val COMET_DEBUG_ENABLED: ConfigEntry[Boolean] =
conf("spark.comet.debug.enabled")
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 5b0516d076..71fec591ee 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
@@ -24,17 +24,19 @@ import java.util.Locale
import java.util.concurrent.ConcurrentHashMap
import scala.collection.mutable
+import scala.concurrent.ExecutionContext
import scala.jdk.CollectionConverters._
import scala.util.control.NonFatal
import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext}
+import org.apache.spark.internal.config.{DYN_ALLOCATION_ENABLED, DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED, SHUFFLE_SERVICE_ENABLED}
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.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory, ResolvedCelebornShufflePusher}
import org.apache.comet.util.ClassLoaders
/**
@@ -51,7 +53,8 @@ class CometCelebornShuffleManager private[shuffle] (
backendFactory: (SparkConf, Boolean) => ShuffleManager,
readerApi: CelebornRawPartitionReader.Api = CelebornRawPartitionReader.reflectedApi,
planningSupportFactory: SparkConf => CelebornNativeShufflePlanningSupport = conf =>
- CometCelebornShuffleManager.nativeShufflePlanningSupport(conf))
+ CometCelebornShuffleManager.nativeShufflePlanningSupport(conf),
+ localManagerFactory: SparkConf => ShuffleManager = conf => new CometShuffleManager(conf))
extends ShuffleManager {
/** Constructor used by Spark on both the driver and executors. */
@@ -72,16 +75,35 @@ class CometCelebornShuffleManager private[shuffle] (
def nativeShuffleFallbackReason(numPartitions: Int): Option[String] =
nativePlanningSupport.fallbackReason(numPartitions)
+ protected[shuffle] def materializationExecutionContext: ExecutionContext =
+ ExecutionContext.global
+
private val nativeShuffleClients =
new ConcurrentHashMap[Int, ConcurrentHashMap[Int, AnyRef]]()
+ private val sizeLimitFallbacks = new ConcurrentHashMap[Int, () => Boolean]()
+ private val localShuffleIds = ConcurrentHashMap.newKeySet[Integer]()
private val ownedNativeClients = new ConcurrentHashMap[AnyRef, java.lang.Boolean]()
@volatile private var nativeGenerationCoordinator: CelebornShuffleGenerationCoordinator = _
@volatile private var nativeGenerationEndpoint: RpcEndpointRef = _
+ @volatile private var localShuffleManagerInstance: ShuffleManager = _
+
+ private def localShuffleManager: ShuffleManager = synchronized {
+ if (localShuffleManagerInstance == null) {
+ localShuffleManagerInstance = Option(localManagerFactory(conf)).getOrElse {
+ throw new IllegalStateException("Local Comet shuffle manager factory returned null")
+ }
+ }
+ localShuffleManagerInstance
+ }
override def registerShuffle[K, V, C](
shuffleId: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
dependency match {
+ case native: CometShuffleDependency[_, _, _]
+ if native.shuffleType == CometNativeShuffle && native.useLocalShuffle =>
+ localShuffleIds.add(shuffleId)
+ localShuffleManager.registerShuffle(shuffleId, dependency)
case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle =>
val handle = backend.registerShuffle(shuffleId, dependency)
if (!CometCelebornShuffleManager.isCelebornHandle(handle)) {
@@ -107,6 +129,10 @@ class CometCelebornShuffleManager private[shuffle] (
mapId: Long,
context: TaskContext,
metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ if (isLocalNativeHandle(handle)) {
+ localShuffleIds.add(handle.shuffleId)
+ return localShuffleManager.getWriter(handle, mapId, context, metrics)
+ }
nativeDependency(handle) match {
case Some(dependency) =>
val earlyClaim = claimNativeShuffleAttempt(handle.shuffleId, context)
@@ -115,11 +141,9 @@ class CometCelebornShuffleManager private[shuffle] (
throw CelebornShufflePusherFactory.commitDenied(context)
}
var preparedClaim = earlyClaim
- val resolved = CelebornShufflePusherFactory.createFromHandle(
- conf,
+ val resolved = createRemotePusher(
handle,
context,
- client => ownedNativeClients.put(client, java.lang.Boolean.TRUE),
(celebornShuffleId, numMappers) =>
preparedClaim = prepareNativeShuffleGeneration(
handle.shuffleId,
@@ -171,7 +195,15 @@ class CometCelebornShuffleManager private[shuffle] (
handle.shuffleId,
resolved.celebornShuffleId,
context,
- preparedClaim))))
+ preparedClaim),
+ onSizeLimitExceeded = failure =>
+ requestLocalShuffle(
+ handle.shuffleId,
+ resolved.celebornShuffleId,
+ resolved.client,
+ context,
+ preparedClaim,
+ failure))))
case None =>
rejectCometHandle(handle)
@@ -188,6 +220,17 @@ class CometCelebornShuffleManager private[shuffle] (
endPartition: Int,
context: TaskContext,
metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ if (isLocalNativeHandle(handle)) {
+ localShuffleIds.add(handle.shuffleId)
+ return localShuffleManager.getReader(
+ handle,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ }
nativeDependency(handle) match {
case Some(dependency) =>
if (startMapIndex > endMapIndex) {
@@ -238,32 +281,60 @@ class CometCelebornShuffleManager private[shuffle] (
override def shuffleBlockResolver: ShuffleBlockResolver = backend.shuffleBlockResolver
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))
+ sizeLimitFallbacks.remove(shuffleId)
+ if (localShuffleIds.remove(shuffleId)) {
+ return localShuffleManager.unregisterShuffle(shuffleId)
}
- backend.unregisterShuffle(shuffleId)
+ val generations = Option(nativeShuffleClients.remove(shuffleId)).toSeq
+ .flatMap(_.asScala.toSeq)
+ var removed = false
+ val cleanup = Seq[() => Unit](() =>
+ Option(localShuffleManagerInstance).foreach(_.unregisterShuffle(shuffleId))) ++
+ generations.map { case (celebornShuffleId, client) =>
+ () => CelebornShufflePusherFactory.cleanupShuffle(client, celebornShuffleId)
+ } ++ Seq[() => Unit](
+ () => {
+ if (isDriver) {
+ Option(nativeGenerationCoordinator).foreach(_.unregisterShuffle(shuffleId))
+ }
+ },
+ () => removed = backend.unregisterShuffle(shuffleId))
+ cleanupAll(cleanup)
+ removed
}
override def stop(): Unit = {
- try backend.stop()
- finally {
- try {
- if (isDriver) {
- Option(nativeGenerationEndpoint).foreach { endpoint =>
- SparkEnv.get.rpcEnv.stop(endpoint)
+ // Spark's disk block manager owns executor shutdown cleanup. In particular, an external
+ // shuffle service must retain local fallback output after this executor has stopped.
+ val localCleanup = Option(localShuffleManagerInstance).toSeq.map { local => () =>
+ local.stop()
+ }
+ sizeLimitFallbacks.clear()
+ cleanupAll(
+ localCleanup ++ Seq[() => Unit](
+ () => backend.stop(),
+ () => {
+ if (isDriver) {
+ Option(nativeGenerationEndpoint).foreach { endpoint =>
+ SparkEnv.get.rpcEnv.stop(endpoint)
+ }
}
- }
- } finally {
- ownedNativeClients.keySet().asScala.foreach(CelebornShufflePusherFactory.releaseClient)
- ownedNativeClients.clear()
- nativeShuffleClients.clear()
+ }) ++ ownedNativeClients.keySet().asScala.toSeq.map { client => () =>
+ CelebornShufflePusherFactory.releaseClient(client)
+ } ++ Seq[() => Unit](() => ownedNativeClients.clear(), () => nativeShuffleClients.clear()))
+ }
+
+ private def cleanupAll(actions: Seq[() => Unit]): Unit = {
+ var failure: Throwable = null
+ actions.foreach { action =>
+ try action()
+ catch {
+ case cleanupFailure: Throwable =>
+ if (failure == null) failure = cleanupFailure
+ else if (failure ne cleanupFailure) failure.addSuppressed(cleanupFailure)
}
}
+ if (failure != null) throw failure
}
private def initializeNativeGenerationCoordinator(): Unit = synchronized {
@@ -273,7 +344,8 @@ class CometCelebornShuffleManager private[shuffle] (
}
val coordinator = new CelebornShuffleGenerationCoordinator(
env.outputCommitCoordinator,
- CelebornShufflePusherFactory.shouldReportShuffleFetchFailure)
+ shouldReportShuffleFetchFailure,
+ shuffleId => Option(sizeLimitFallbacks.get(shuffleId)).exists(callback => callback()))
val endpoint = env.rpcEnv.setupEndpoint(
CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT,
new CelebornShuffleGenerationEndpoint(env.rpcEnv, coordinator))
@@ -310,6 +382,91 @@ class CometCelebornShuffleManager private[shuffle] (
taskContext.partitionId(),
taskContext.attemptNumber()))
+ protected[shuffle] def createRemotePusher(
+ handle: ShuffleHandle,
+ context: TaskContext,
+ onGenerationResolved: (Int, Int) => Unit,
+ onGenerationInvalidated: (Int, Int) => Unit,
+ onInvalidationUnsafe: (Int, Int) => Boolean): ResolvedCelebornShufflePusher =
+ CelebornShufflePusherFactory.createFromHandle(
+ conf,
+ handle,
+ context,
+ client => ownedNativeClients.put(client, java.lang.Boolean.TRUE),
+ onGenerationResolved,
+ onGenerationInvalidated,
+ onInvalidationUnsafe)
+
+ protected[shuffle] def shouldReportShuffleFetchFailure(taskAttemptId: Long): Boolean =
+ CelebornShufflePusherFactory.shouldReportShuffleFetchFailure(taskAttemptId)
+
+ /** Registers the driver-owned materialization that can replace an unpublished shuffle. */
+ private[shuffle] def registerSizeLimitFallback(
+ shuffleId: Int,
+ callback: () => Boolean): Unit = {
+ require(isDriver, "Shuffle fallback belongs to the driver")
+ require(
+ sizeLimitFallbacks.putIfAbsent(shuffleId, callback) == null,
+ s"Shuffle $shuffleId already has a materialization")
+ }
+
+ private[shuffle] def removeSizeLimitFallback(shuffleId: Int): Unit = {
+ sizeLimitFallbacks.remove(shuffleId)
+ }
+
+ private def isLocalNativeHandle(handle: ShuffleHandle): Boolean = handle match {
+ case native: CometNativeShuffleHandle[_, _] =>
+ native.dependency match {
+ case dependency: CometShuffleDependency[_, _, _] => dependency.useLocalShuffle
+ case _ => false
+ }
+ case _ => false
+ }
+
+ private def requestLocalShuffle(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ client: AnyRef,
+ context: TaskContext,
+ claim: CelebornMapAttemptClaim,
+ failure: Throwable): Unit = {
+ val accepted = generationEndpoint.askSync[Boolean](
+ RequestLocalCometShuffle(
+ ValidateCelebornMapAttempt(
+ shuffleId,
+ celebornShuffleId,
+ context.stageId(),
+ context.stageAttemptNumber(),
+ context.partitionId(),
+ context.attemptNumber(),
+ claim.epoch),
+ context.taskAttemptId()))
+ if (!accepted) {
+ // The remote dependency may already have been published to readers. Preserve the typed
+ // failure instead of silently changing the identity of a shuffle they are consuming.
+ throw failure
+ }
+
+ // The driver is materializing a fresh local dependency. Cleanup is best effort: neither a
+ // lost executor nor a delayed remote success can affect that dependency's map outputs.
+ try {
+ client.getClass
+ .getMethod(
+ "reportShuffleFetchFailure",
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Long.TYPE)
+ .invoke(
+ client,
+ Int.box(shuffleId),
+ Int.box(celebornShuffleId),
+ Long.box(context.taskAttemptId()))
+ } catch {
+ case NonFatal(cleanupFailure) => failure.addSuppressed(cleanupFailure)
+ }
+ throw failure
+ }
+
private def prepareNativeShuffleGeneration(
shuffleId: Int,
celebornShuffleId: Int,
@@ -452,6 +609,16 @@ private[shuffle] object CometCelebornShuffleManager {
return CelebornNativeShufflePlanningSupport(
Some("Native Celeborn shuffle does not support spark.io.encryption.enabled=true"))
}
+ if (conf.get(DYN_ALLOCATION_ENABLED) && !conf.get(SHUFFLE_SERVICE_ENABLED) &&
+ !conf.get(DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED)) {
+ // A reliable remote ShuffleDataIO plugin can enable dynamic allocation without retaining
+ // executor-local output. Size-limit recovery must also preserve its local fallback files.
+ return CelebornNativeShufflePlanningSupport(
+ Some(
+ "Native Celeborn shuffle with dynamic allocation requires " +
+ "spark.shuffle.service.enabled=true or " +
+ "spark.dynamicAllocation.shuffleTracking.enabled=true for local fallback"))
+ }
try {
val completionUnavailable = pushCompletionUnavailableReason()
if (completionUnavailable.nonEmpty) {
@@ -577,12 +744,18 @@ private[shuffle] final case class AbandonCelebornMapAttempt(
taskAttemptId: Long)
extends Serializable
+private[shuffle] final case class RequestLocalCometShuffle(
+ validation: ValidateCelebornMapAttempt,
+ taskAttemptId: Long)
+ extends Serializable
+
/**
* Keeps Comet map admission aligned with Spark task failures and Celeborn shuffle generations.
*/
private[shuffle] final class CelebornShuffleGenerationCoordinator(
outputCommitCoordinator: OutputCommitCoordinator,
- shouldReportShuffleFetchFailure: Long => Boolean = _ => true) {
+ shouldReportShuffleFetchFailure: Long => Boolean = _ => true,
+ requestFallback: Int => Boolean = _ => false) {
private val generations = mutable.HashMap.empty[Int, PrepareCelebornShuffleGeneration]
private val invalidatedGenerations = mutable.HashSet.empty[Int]
@@ -590,6 +763,9 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator(
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]]
+ // An abandoned remote shuffle never becomes local. Its replacement owns a new shuffle ID,
+ // so even a success already queued in Spark cannot overwrite or skip a replacement map.
+ private val abandonedShuffles = mutable.HashSet.empty[Int]
private def currentEpoch(shuffleId: Int): Long = generationEpochs.getOrElse(shuffleId, 0L)
@@ -664,6 +840,9 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator(
}
def claimMapAttempt(claim: ClaimCelebornMapAttempt): CelebornMapAttemptClaim = synchronized {
+ if (abandonedShuffles.contains(claim.shuffleId)) {
+ return CelebornMapAttemptClaim(false, currentEpoch(claim.shuffleId))
+ }
val previousGeneration = generations.get(claim.shuffleId)
val stale = previousGeneration.exists { generation =>
generation.stageId == claim.stageId &&
@@ -724,6 +903,9 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator(
def prepareGeneration(generation: PrepareCelebornShuffleGeneration): Boolean = synchronized {
require(generation.numMappers > 0, "Celeborn shuffle mapper count must be positive")
+ if (abandonedShuffles.contains(generation.shuffleId)) {
+ return false
+ }
generations.get(generation.shuffleId) match {
case Some(previous)
@@ -867,7 +1049,23 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator(
abandon
}
+ def requestLocalShuffle(request: RequestLocalCometShuffle): Boolean = synchronized {
+ val validation = request.validation
+ if (abandonedShuffles.contains(validation.shuffleId) ||
+ !validateMapAttempt(validation) ||
+ !shouldReportShuffleFetchFailure(request.taskAttemptId) ||
+ !requestFallback(validation.shuffleId)) {
+ return false
+ }
+
+ abandonedShuffles.add(validation.shuffleId)
+ invalidatedGenerations.add(validation.shuffleId)
+ invalidateOwners(validation.shuffleId)
+ true
+ }
+
def unregisterShuffle(shuffleId: Int): Unit = synchronized {
+ abandonedShuffles.remove(shuffleId)
generations.remove(shuffleId)
invalidatedGenerations.remove(shuffleId)
generationEpochs.remove(shuffleId)
@@ -894,5 +1092,7 @@ private[shuffle] final class CelebornShuffleGenerationEndpoint(
context.reply(coordinator.invalidateGeneration(invalidation))
case abandoned: AbandonCelebornMapAttempt =>
context.reply(coordinator.abandonMapAttempt(abandoned))
+ case request: RequestLocalCometShuffle =>
+ context.reply(coordinator.requestLocalShuffle(request))
}
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala
new file mode 100644
index 0000000000..f9f6770b40
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala
@@ -0,0 +1,359 @@
+/*
+ * 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.util.Properties
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.concurrent.{CanAwait, ExecutionContext, Future, Promise}
+import scala.concurrent.duration.Duration
+import scala.util.{Failure, Success, Try}
+import scala.util.control.NonFatal
+
+import org.apache.spark.{FutureAction, MapOutputStatistics, ShuffleDependency, SparkException}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart}
+import org.apache.spark.sql.comet.shims.ShimCometShuffleMaterialization
+import org.apache.spark.util.ThreadUtils
+
+/**
+ * Driver-owned materialization of one native Celeborn shuffle. Storage is chosen before any
+ * downstream RDD can depend on its output. A size-limit failure cancels the remote map-stage job
+ * and materializes a fresh local dependency with an independent scheduling RDD. Their separate
+ * RDD, shuffle and stage identities isolate remote failures and late task completions from the
+ * replacement. Once output is published, its destination is fixed.
+ */
+private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C](
+ remoteDependency: CometShuffleDependency[K, V, C],
+ manager: CometCelebornShuffleManager)
+ extends FutureAction[MapOutputStatistics]
+ with ShimCometShuffleMaterialization {
+
+ import CometCelebornShuffleMaterialization._
+
+ private val sparkContext = remoteDependency.rdd.context
+ private val executionContext = manager.materializationExecutionContext
+ private val localProperties = sparkContext.getLocalProperties.clone().asInstanceOf[Properties]
+ private val contextClassLoader = Thread.currentThread().getContextClassLoader
+ private val completion = Promise[MapOutputStatistics]()
+ private val actions = ArrayBuffer.empty[FutureAction[MapOutputStatistics]]
+ private var state: State = RunningRemote
+ private var activeAction: Option[FutureAction[MapOutputStatistics]] = None
+ private var remoteAction: Option[FutureAction[MapOutputStatistics]] = None
+ private var remoteCancellationIssued = false
+ private var selected: Option[CometShuffleDependency[K, V, C]] = None
+
+ private val jobListener = new SparkListener {
+ override def onJobStart(event: SparkListenerJobStart): Unit = {
+ localJobStarted(event.jobId)
+ }
+ }
+ sparkContext.addSparkListener(jobListener)
+ try {
+ manager.registerSizeLimitFallback(remoteDependency.shuffleId, () => requestLocalShuffle())
+ submit(remoteDependency, RunningRemote)
+ } catch {
+ case NonFatal(failure) =>
+ removeCallbacks()
+ throw failure
+ }
+
+ private def withCapturedProperties[T](body: => T): T = withCapturedSession {
+ val thread = Thread.currentThread()
+ val previousProperties = sparkContext.getLocalProperties
+ val previousClassLoader = thread.getContextClassLoader
+ sparkContext.setLocalProperties(localProperties.clone().asInstanceOf[Properties])
+ thread.setContextClassLoader(contextClassLoader)
+ try body
+ finally {
+ thread.setContextClassLoader(previousClassLoader)
+ sparkContext.setLocalProperties(previousProperties)
+ }
+ }
+
+ private def submit(dependency: CometShuffleDependency[K, V, C], expected: State): Unit = {
+ implicit val executor: ExecutionContext = executionContext
+ val ready = Future {
+ withCapturedProperties(inputMaterializations(dependency.rdd))
+ }.flatMap(upstream => Future.sequence(upstream))
+ ready.onComplete {
+ case Failure(failure) =>
+ fail(expected, failure)
+ cancelActions(failure)
+ case Success(_) =>
+ try
+ withCapturedProperties {
+ // Upstream storage decisions have completed without occupying preparation workers.
+ // Cache the graph before taking our lock: submitMapStage also traverses it eagerly.
+ prepareInput(dependency.rdd)
+ val action = synchronized {
+ if (state != expected) {
+ None
+ } else {
+ if (expected == RunningLocal) remoteFailure.foreach(failure => throw failure)
+ // Keep job admission and action assignment atomic with cancellation and size
+ // reports.
+ val submitted = sparkContext.submitMapStage(dependency)
+ activeAction = Some(submitted)
+ if (expected == RunningRemote) remoteAction = Some(submitted)
+ actions += submitted
+ Some(submitted)
+ }
+ }
+ action.foreach(_.onComplete(result => finish(dependency, expected, result)))
+ }
+ catch {
+ case NonFatal(failure) =>
+ fail(expected, failure)
+ cancelActions(failure)
+ }
+ }
+ }
+
+ private def inputMaterializations(input: RDD[_]): Seq[Future[MapOutputStatistics]] = {
+ val visited = mutable.Set.empty[Int]
+ val upstream = mutable.Set.empty[Future[MapOutputStatistics]]
+ val pending = mutable.Stack[RDD[_]](input)
+ while (pending.nonEmpty && synchronized { state != Finished && state != Cancelled }) {
+ val next = pending.pop()
+ if (visited.add(next.id)) {
+ next match {
+ // Calling this RDD's dependencies would wait under Spark's RDD lock. Compose its
+ // completion instead; its materialization already owns preparation of its ancestors.
+ case shuffled: CometShuffledBatchRDD if !shuffled.isCheckpointed =>
+ forDependency(shuffled.dependency) match {
+ case Some(materialization) => upstream += materialization
+ case None => pending.push(shuffled.dependency.rdd)
+ }
+ case other => other.dependencies.foreach(dependency => pending.push(dependency.rdd))
+ }
+ }
+ }
+ upstream.toSeq
+ }
+
+ private def prepareInput(input: RDD[_]): Unit = {
+ val visited = mutable.Set.empty[Int]
+ val pending = mutable.Stack[RDD[_]](input)
+ while (pending.nonEmpty && synchronized { state != Finished && state != Cancelled }) {
+ val next = pending.pop()
+ if (visited.add(next.id)) {
+ next.partitions
+ next.dependencies.foreach(dependency => pending.push(dependency.rdd))
+ }
+ }
+ }
+
+ private def finish(
+ dependency: CometShuffleDependency[K, V, C],
+ expected: State,
+ result: Try[MapOutputStatistics]): Unit = {
+ val finished = synchronized {
+ if (state == RunningLocal && expected == RunningRemote &&
+ !remoteCancellationIssued && result.failed.toOption.exists(isNonSizeLimitFailure)) {
+ // A group cancellation processed before the replacement job starts must not disappear
+ // just because fallback has already won the storage decision.
+ state = Finished
+ completion.tryComplete(result)
+ cancelActions(result.failed.get)
+ true
+ } else if (state != expected) {
+ false
+ } else {
+ // JobStart listeners are asynchronous. If local output finishes before that notification,
+ // it still proves the replacement was active and lets us retire the old job here.
+ val outcome = if (expected == RunningLocal && result.isSuccess) {
+ cancelRemoteAfterLocalStarted().map(Failure(_)).getOrElse(result)
+ } else {
+ result
+ }
+ val published = outcome.map { statistics =>
+ dependency.outputMetrics.foreach(_.publish(sparkContext))
+ statistics
+ }
+ published match {
+ case Success(_) => selected = Some(dependency)
+ case _ =>
+ }
+ state = Finished
+ completion.tryComplete(published)
+ published.failed.toOption.foreach(cancelActions)
+ true
+ }
+ }
+ if (finished) removeCallbacks()
+ }
+
+ private def fail(expected: State, failure: Throwable): Unit =
+ finish(remoteDependency, expected, Failure(failure))
+
+ private def cancelActions(failure: Throwable): Unit = synchronized {
+ actions.foreach { action =>
+ try action.cancel()
+ catch {
+ case NonFatal(cancellationFailure) =>
+ if (failure ne cancellationFailure) failure.addSuppressed(cancellationFailure)
+ }
+ }
+ }
+
+ private def removeCallbacks(): Unit = {
+ try manager.removeSizeLimitFallback(remoteDependency.shuffleId)
+ finally sparkContext.removeSparkListener(jobListener)
+ }
+
+ private def isNonSizeLimitFailure(failure: Throwable): Boolean =
+ !CometNativeShuffleWriter.isSizeLimitFailure(failure)
+
+ private def remoteFailure: Option[Throwable] =
+ remoteAction.flatMap(_.value).flatMap(_.failed.toOption).filter(isNonSizeLimitFailure)
+
+ // Called with this materialization's lock held, after Spark has made the local job active.
+ private def cancelRemoteAfterLocalStarted(): Option[Throwable] = {
+ if (remoteCancellationIssued) return None
+ val previousFailure = remoteFailure
+ if (previousFailure.nonEmpty) return previousFailure
+ remoteCancellationIssued = true
+ try {
+ remoteAction.foreach(_.cancel())
+ None
+ } catch {
+ case NonFatal(failure) => Some(failure)
+ }
+ }
+
+ private def localJobStarted(jobId: Int): Unit = synchronized {
+ if (state == RunningLocal && activeAction.exists(_.jobIds.contains(jobId))) {
+ cancelRemoteAfterLocalStarted().foreach { failure =>
+ fail(RunningLocal, failure)
+ cancelActions(failure)
+ }
+ }
+ }
+
+ private def requestLocalShuffle(): Boolean = synchronized {
+ if (state != RunningRemote || activeAction.exists(_.isCompleted)) {
+ false
+ } else {
+ state = RunningLocal
+ try {
+ // Register the replacement before cancelling the old job. Its upstream inputs were
+ // already materialized for the remote submission.
+ // The JobStart listener retires the remote job only after Spark has made the replacement
+ // active. Until then, any external cancellation of the remote job also cancels fallback.
+ val localDependency =
+ withCapturedProperties(remoteDependency.createLocalShuffleDependency())
+ remoteFailure match {
+ case Some(failure) =>
+ fail(RunningLocal, failure)
+ cancelActions(failure)
+ case None => submit(localDependency, RunningLocal)
+ }
+ true
+ } catch {
+ case NonFatal(failure) =>
+ fail(RunningLocal, failure)
+ cancelActions(failure)
+ // The storage transition won even if creating the replacement failed. The coordinator
+ // must still fence the abandoned remote ID; the materialization retains this failure.
+ true
+ }
+ }
+ }
+
+ def completedDependency: Option[CometShuffleDependency[K, V, C]] = synchronized { selected }
+
+ def selectedDependency: CometShuffleDependency[K, V, C] = {
+ try ThreadUtils.awaitResult(this, Duration.Inf)
+ catch {
+ case interrupted: InterruptedException =>
+ cancel()
+ Thread.currentThread().interrupt()
+ throw interrupted
+ }
+ synchronized {
+ selected.getOrElse(throw new IllegalStateException("Shuffle materialization has no output"))
+ }
+ }
+
+ override def cancel(): Unit = cancel(None)
+
+ // Spark 4 calls this overload; Spark 3's FutureAction only declares cancel().
+ def cancel(reason: Option[String]): Unit = {
+ val cancelled = synchronized {
+ if (state == Finished || state == Cancelled) {
+ false
+ } else {
+ state = Cancelled
+ val failure =
+ new SparkException(reason.getOrElse("Comet shuffle materialization cancelled"))
+ cancelActions(failure)
+ completion.tryFailure(failure)
+ true
+ }
+ }
+ if (cancelled) removeCallbacks()
+ }
+
+ override def isCancelled: Boolean = synchronized { state == Cancelled }
+
+ override def jobIds: Seq[Int] = synchronized { actions.iterator.flatMap(_.jobIds).toVector }
+
+ override def isCompleted: Boolean = completion.isCompleted
+
+ override def value: Option[Try[MapOutputStatistics]] = completion.future.value
+
+ override def ready(atMost: Duration)(implicit permit: CanAwait): this.type = {
+ completion.future.ready(atMost)
+ this
+ }
+
+ override def result(atMost: Duration)(implicit permit: CanAwait): MapOutputStatistics =
+ completion.future.result(atMost)
+
+ override def onComplete[U](f: Try[MapOutputStatistics] => U)(implicit
+ executor: ExecutionContext): Unit = completion.future.onComplete(f)
+
+ override def transform[S](f: Try[MapOutputStatistics] => Try[S])(implicit
+ executor: ExecutionContext): Future[S] = completion.future.transform(f)
+
+ override def transformWith[S](f: Try[MapOutputStatistics] => Future[S])(implicit
+ executor: ExecutionContext): Future[S] = completion.future.transformWith(f)
+}
+
+private[shuffle] object CometCelebornShuffleMaterialization {
+ private sealed trait State
+ private case object RunningRemote extends State
+ private case object RunningLocal extends State
+ private case object Finished extends State
+ private case object Cancelled extends State
+
+ def forDependency(dependency: ShuffleDependency[_, _, _])
+ : Option[CometCelebornShuffleMaterialization[_, _, _]] = dependency match {
+ case comet: CometShuffleDependency[_, _, _] => comet.materialization
+ case _ => None
+ }
+
+ def selectForRead(dependency: ShuffleDependency[Int, _, _]): ShuffleDependency[Int, _, _] =
+ forDependency(dependency)
+ .map(_.selectedDependency.asInstanceOf[ShuffleDependency[Int, _, _]])
+ .getOrElse(dependency)
+}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
index 41d1d6cc6a..bced55b34d 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
@@ -45,6 +45,20 @@ private[shuffle] class CometNativeShuffleInputRDD(
sc,
inputRDDs.map(rdd => new OneToOneDependency(rdd))) {
+ /**
+ * Give local fallback its own scheduling RDD over the original upstream inputs. Spark aborts
+ * jobs whose RDD ancestry contains the failed stage's RDD, so reusing this instance or wrapping
+ * it in a narrow dependency lets a late remote failure abort the local replacement as well.
+ */
+ private[shuffle] def copyForLocalShuffle(): CometNativeShuffleInputRDD =
+ new CometNativeShuffleInputRDD(
+ context,
+ inputRDDs,
+ numPartitionsParam,
+ shuffleScanIndices,
+ spillMetricNode,
+ perPartitionByKey)
+
override protected def getPartitions: Array[Partition] =
(0 until numPartitionsParam).map { i =>
// Resolve leaf-RDD partitions on the driver here (where their @transient fields are still
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 89be100dbd..21dd768630 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
@@ -39,7 +39,7 @@ import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.types.StructField
import org.apache.spark.util.{ThreadUtils, Utils}
-import org.apache.comet.{CometConf, CometExecIterator}
+import org.apache.comet.{CometConf, CometExecIterator, CometShuffleSizeLimitException}
import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde}
import org.apache.comet.serde.OperatorOuterClass.{CompressionCodec, Operator}
import org.apache.comet.serde.operator.schema2Proto
@@ -103,6 +103,9 @@ class CometNativeShuffleWriter[K, V](
catch {
case cleanupFailure: Throwable => failure.addSuppressed(cleanupFailure)
}
+ if (CometNativeShuffleWriter.isSizeLimitFailure(failure)) {
+ destination.onSizeLimitExceeded(failure)
+ }
}
throw failure
}
@@ -484,6 +487,16 @@ class CometNativeShuffleWriter[K, V](
}
private[shuffle] object CometNativeShuffleWriter {
+ private[shuffle] def isSizeLimitFailure(failure: Throwable): Boolean = {
+ var cause = failure
+ val visited = new java.util.IdentityHashMap[Throwable, java.lang.Boolean]()
+ while (cause != null && visited.put(cause, java.lang.Boolean.TRUE) == null) {
+ if (cause.isInstanceOf[CometShuffleSizeLimitException]) return true
+ cause = cause.getCause
+ }
+ false
+ }
+
def drainAndClose(iterator: Iterator[_], close: () => Unit): Unit = {
Utils.tryWithSafeFinally {
while (iterator.hasNext) {
@@ -503,7 +516,8 @@ private[shuffle] final case class CelebornNativeShuffleDestination(
maxFrameBytes: Int,
numPartitions: Int,
commitAuthorized: Boolean = false,
- commitValidator: () => Boolean = () => true) {
+ commitValidator: () => Boolean = () => true,
+ onSizeLimitExceeded: Throwable => Unit = _ => ()) {
require(pusher != null, "The Celeborn shuffle partition pusher must not be null")
require(maxFrameBytes >= 20, "The Celeborn shuffle frame limit must fit a complete frame")
require(
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala
index 2a05843007..34f1407fc9 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala
@@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.plans.physical.Partitioning
import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext}
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.types.StructType
@@ -68,7 +69,9 @@ class CometShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag](
val shuffleWriteMetrics: Map[String, SQLMetric] = Map.empty,
val numParts: Int = 0,
val rangePartitionBounds: Option[Seq[InternalRow]] = None,
- val nativeShuffleSpec: Option[NativeShuffleSpec] = None)
+ val nativeShuffleSpec: Option[NativeShuffleSpec] = None,
+ val useLocalShuffle: Boolean = false,
+ private[shuffle] val outputMetrics: Option[CometShuffleOutputMetrics] = None)
extends ShuffleDependency[K, V, C](
_rdd,
partitioner,
@@ -76,7 +79,66 @@ class CometShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag](
keyOrdering,
aggregator,
mapSideCombine,
- shuffleWriterProcessor) {}
+ shuffleWriterProcessor) {
+
+ @transient @volatile private var materializationInstance
+ : CometCelebornShuffleMaterialization[K, V, C] = _
+
+ // One driver-owned materialization selects storage before any reader depends on this shuffle.
+ // Local Comet shuffle and ordinary Spark shuffle retain their existing lazy scheduling.
+ @transient private[shuffle] lazy val materialization
+ : Option[CometCelebornShuffleMaterialization[K, V, C]] = {
+ if (shuffleType == CometNativeShuffle && !useLocalShuffle && rdd.getNumPartitions > 0) {
+ rdd.context.env.shuffleManager match {
+ case manager: CometCelebornShuffleManager =>
+ val instance = new CometCelebornShuffleMaterialization(this, manager)
+ materializationInstance = instance
+ Some(instance)
+ case _ => None
+ }
+ } else {
+ None
+ }
+ }
+
+ private[shuffle] def currentShuffleDependency: CometShuffleDependency[K, V, C] =
+ Option(materializationInstance).flatMap(_.completedDependency).getOrElse(this)
+
+ private[shuffle] def createLocalShuffleDependency(): CometShuffleDependency[K, V, C] = {
+ // Spark aborts dependent jobs by input RDD identity, not just shuffle or stage identity.
+ // Use a sibling scheduling RDD so a failure in the abandoned remote stage cannot abort
+ // the replacement. The upstream inputs remain shared.
+ val localRDD = rdd
+ .asInstanceOf[CometNativeShuffleInputRDD]
+ .copyForLocalShuffle()
+ .asInstanceOf[RDD[_ <: Product2[K, V]]]
+ val localOutputMetrics = outputMetrics.map(_.newDestination(rdd.context))
+ val localWriteMetrics = shuffleWriteMetrics ++ localOutputMetrics.toSeq.flatMap(_.metrics)
+ new CometShuffleDependency[K, V, C](
+ localRDD,
+ partitioner,
+ serializer,
+ keyOrdering,
+ aggregator,
+ mapSideCombine,
+ if (localOutputMetrics.nonEmpty) {
+ ShuffleExchangeExec.createShuffleWriteProcessor(localWriteMetrics)
+ } else {
+ shuffleWriterProcessor
+ },
+ shuffleType,
+ schema,
+ decodeTime,
+ outputPartitioning,
+ outputAttributes,
+ localWriteMetrics,
+ numParts,
+ rangePartitionBounds,
+ nativeShuffleSpec,
+ useLocalShuffle = true,
+ outputMetrics = localOutputMetrics)
+ }
+}
/** Indicates shuffle type */
sealed trait ShuffleType
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 46fe621e48..39cd5186a7 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
@@ -148,7 +148,10 @@ case class CometShuffleExchangeExec(
if (inputRDD.getNumPartitions == 0) {
Future.successful(null)
} else {
- sparkContext.submitMapStage(shuffleDependency)
+ CometCelebornShuffleMaterialization.forDependency(shuffleDependency) match {
+ case Some(materialization) => materialization
+ case None => sparkContext.submitMapStage(shuffleDependency)
+ }
}
}
@@ -167,7 +170,13 @@ case class CometShuffleExchangeExec(
}
// TODO: add `override` keyword after dropping Spark-3.x supports
- def shuffleId: Int = getShuffleId(shuffleDependency)
+ def shuffleId: Int = {
+ val current = shuffleDependency match {
+ case comet: CometShuffleDependency[Int @unchecked, _, _] => comet.currentShuffleDependency
+ case other => other
+ }
+ getShuffleId(current)
+ }
/**
* A [[ShuffleDependency]] that will partition rows of its child based on the partitioning
@@ -882,19 +891,31 @@ object CometShuffleExchangeExec
None)
}
+ // The remote stage can finish some maps before an oversized row requires a complete local
+ // replacement. Keep output statistics separate from the public counters so neither those
+ // completed maps nor late remote updates inflate the selected shuffle's AQE statistics.
+ val outputMetrics = thinRDD.context.env.shuffleManager match {
+ case _: CometCelebornShuffleManager if numParts > 0 =>
+ Some(CometShuffleOutputMetrics(thinRDD.context, metrics))
+ case _ => None
+ }
+ val destinationMetrics = metrics ++ outputMetrics.toSeq.flatMap(_.metrics)
+
new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch](
thinRDD,
serializer = serializer,
- shuffleWriterProcessor = ShuffleExchangeExec.createShuffleWriteProcessor(metrics),
+ shuffleWriterProcessor =
+ ShuffleExchangeExec.createShuffleWriteProcessor(destinationMetrics),
shuffleType = CometNativeShuffle,
partitioner = partitioner,
decodeTime = metrics("decode_time"),
outputPartitioning = Some(outputPartitioning),
outputAttributes = outputAttributes,
- shuffleWriteMetrics = metrics,
+ shuffleWriteMetrics = destinationMetrics,
numParts = numParts,
rangePartitionBounds = rangePartitionBounds,
- nativeShuffleSpec = Some(augmentedSpec))
+ nativeShuffleSpec = Some(augmentedSpec),
+ outputMetrics = outputMetrics)
}
/**
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleOutputMetrics.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleOutputMetrics.scala
new file mode 100644
index 0000000000..029220812d
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleOutputMetrics.scala
@@ -0,0 +1,69 @@
+/*
+ * 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.apache.spark.SparkContext
+import org.apache.spark.sql.execution.SQLExecution
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleWriteMetricsReporter}
+
+/**
+ * Output statistics for one shuffle destination. Remote and replacement tasks update distinct
+ * accumulators; only the successful destination publishes output statistics to the exchange.
+ * Timing and operator metrics still account for the work done by both destinations.
+ */
+private[shuffle] final class CometShuffleOutputMetrics private (
+ val metrics: Map[String, SQLMetric],
+ @transient private val publishedMetrics: Map[String, SQLMetric],
+ private val executionId: String)
+ extends Serializable {
+
+ def newDestination(sc: SparkContext): CometShuffleOutputMetrics =
+ CometShuffleOutputMetrics.create(sc, publishedMetrics, executionId)
+
+ def publish(sc: SparkContext): Unit = {
+ metrics.foreach { case (name, metric) => publishedMetrics(name).set(metric.value) }
+ if (executionId != null) {
+ SQLMetrics.postDriverMetricUpdates(sc, executionId, publishedMetrics.values.toSeq)
+ }
+ }
+}
+
+private[shuffle] object CometShuffleOutputMetrics {
+ import SQLShuffleWriteMetricsReporter.{SHUFFLE_BYTES_WRITTEN, SHUFFLE_RECORDS_WRITTEN}
+
+ def apply(sc: SparkContext, metrics: Map[String, SQLMetric]): CometShuffleOutputMetrics = {
+ val outputNames = Set("dataSize", SHUFFLE_BYTES_WRITTEN, SHUFFLE_RECORDS_WRITTEN)
+ create(
+ sc,
+ metrics.filter { case (name, _) => outputNames.contains(name) },
+ sc.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
+ }
+
+ private def create(
+ sc: SparkContext,
+ publishedMetrics: Map[String, SQLMetric],
+ executionId: String): CometShuffleOutputMetrics = {
+ val metrics = Map(
+ "dataSize" -> SQLMetrics.createSizeMetric(sc, "data size"),
+ SHUFFLE_BYTES_WRITTEN -> SQLMetrics.createSizeMetric(sc, "shuffle bytes written"),
+ SHUFFLE_RECORDS_WRITTEN -> SQLMetrics.createMetric(sc, "shuffle records written"))
+ new CometShuffleOutputMetrics(metrics, publishedMetrics, executionId)
+ }
+}
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 2a598a7872..17f84ed662 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
@@ -51,7 +51,18 @@ class CometShuffledBatchRDD(
SortShuffleManager.FETCH_SHUFFLE_BLOCKS_IN_BATCH_ENABLED_KEY,
SQLConf.get.fetchShuffleBlocksInBatch.toString)
- override def getDependencies: Seq[Dependency[_]] = List(dependency)
+ // Start without waiting so constructing a join or union can start every independent input.
+ // Spark resolves dependencies on the submitting thread before queueing a downstream job;
+ // freeze the selected destination there, before Spark caches this RDD's dependency graph.
+ // Cache fixed partition metadata first: parents such as UnionRDD read it while discovering
+ // dependencies, and must not wait for the RDD lock held by a concurrent getDependencies call.
+ partitions
+ CometCelebornShuffleMaterialization.forDependency(dependency)
+
+ override def getDependencies: Seq[Dependency[_]] = {
+ dependency = CometCelebornShuffleMaterialization.selectForRead(dependency)
+ List(dependency)
+ }
override val partitioner: Option[Partitioner] =
if (partitionSpecs.forall(_.isInstanceOf[CoalescedPartitionSpec])) {
diff --git a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala
similarity index 51%
copy from spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala
copy to spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala
index 6104c320b9..26ba626a66 100644
--- a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala
+++ b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala
@@ -17,14 +17,28 @@
* under the License.
*/
-package org.apache.spark.shuffle.celeborn
+package org.apache.spark.sql.comet.shims
-import org.apache.spark.ShuffleDependency
-import org.apache.spark.shuffle.BaseShuffleHandle
+import org.apache.spark.sql.SparkSession
-/** 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)
+/**
+ * Restores the session context captured before shuffle materialization leaves the caller thread.
+ */
+trait ShimCometShuffleMaterialization {
+ private val activeSession = SparkSession.getActiveSession
+
+ protected def withCapturedSession[T](body: => T): T = {
+ val previousSession = SparkSession.getActiveSession
+ activeSession match {
+ case Some(session) => SparkSession.setActiveSession(session)
+ case None => SparkSession.clearActiveSession()
+ }
+ try body
+ finally {
+ previousSession match {
+ case Some(session) => SparkSession.setActiveSession(session)
+ case None => SparkSession.clearActiveSession()
+ }
+ }
+ }
+}
diff --git a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala
new file mode 100644
index 0000000000..440212dcdf
--- /dev/null
+++ b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala
@@ -0,0 +1,47 @@
+/*
+ * 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.shims
+
+import org.apache.spark.JobArtifactSet
+import org.apache.spark.sql.SparkSession
+
+/**
+ * Restores the session context captured before shuffle materialization leaves the caller thread.
+ */
+trait ShimCometShuffleMaterialization {
+ private val activeSession = SparkSession.getActiveSession
+ private val artifactState = JobArtifactSet.getCurrentJobArtifactState.orNull
+
+ protected def withCapturedSession[T](body: => T): T =
+ JobArtifactSet.withActiveJobArtifactState(artifactState) {
+ val previousSession = SparkSession.getActiveSession
+ activeSession match {
+ case Some(session) => SparkSession.setActiveSession(session)
+ case None => SparkSession.clearActiveSession()
+ }
+ try body
+ finally {
+ previousSession match {
+ case Some(session) => SparkSession.setActiveSession(session)
+ case None => SparkSession.clearActiveSession()
+ }
+ }
+ }
+}
diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala
new file mode 100644
index 0000000000..81fa37f9f6
--- /dev/null
+++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala
@@ -0,0 +1,50 @@
+/*
+ * 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.shims
+
+import org.apache.spark.JobArtifactSet
+import org.apache.spark.sql.classic.SparkSession
+
+/**
+ * Restores the session context captured before shuffle materialization leaves the caller thread.
+ */
+trait ShimCometShuffleMaterialization {
+ private val activeSession = SparkSession.getActiveSession
+ // Spark 4 also associates artifacts with ordinary, non-Connect sessions.
+ private val artifactState = JobArtifactSet.getCurrentJobArtifactState
+ .orElse(activeSession.map(_.artifactManager.state))
+ .orNull
+
+ protected def withCapturedSession[T](body: => T): T =
+ JobArtifactSet.withActiveJobArtifactState(artifactState) {
+ val previousSession = SparkSession.getActiveSession
+ activeSession match {
+ case Some(session) => SparkSession.setActiveSession(session)
+ case None => SparkSession.clearActiveSession()
+ }
+ try body
+ finally {
+ previousSession match {
+ case Some(session) => SparkSession.setActiveSession(session)
+ case None => SparkSession.clearActiveSession()
+ }
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
index d6db22b282..0265f30481 100644
--- a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
@@ -161,7 +161,7 @@ class CometConfSuite extends AnyFunSuite {
val conf = new SQLConf
assert(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.get(conf) == 64L * 1024 * 1024)
- assert(CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.get(conf) == 256L * 1024 * 1024)
+ assert(CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.get(conf) == 512L * 1024 * 1024)
}
test("remote shuffle frame limit rejects incomplete frames and oversized JVM requests") {
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
index 10cdbc7932..7012ae1ba5 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
@@ -43,7 +43,7 @@ import org.apache.spark.sql.functions.{col, count, sum}
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
-import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, Native}
+import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native}
import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass}
import org.apache.comet.shuffle.ShufflePartitionPusher
@@ -249,8 +249,11 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
iterator.hasNext
Iterator.single((false, callbacks))
} catch {
- case failure: Exception =>
- Iterator.single((failure.getMessage.contains("admission budget"), callbacks))
+ case failure: CometShuffleSizeLimitException =>
+ Iterator.single(
+ (
+ failure.getMessage.contains("spark.comet.shuffle.rss.maxInFlightBytes"),
+ callbacks))
}
} finally {
iterator.close()
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
index 6104c320b9..d2f825764c 100644
--- a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala
+++ b/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala
@@ -27,4 +27,6 @@ class CelebornShuffleHandle[K, V, C](
shuffleId: Int,
dependency: ShuffleDependency[K, V, C],
val stageRerunEnabled: Boolean = true)
- extends BaseShuffleHandle[K, V, C](shuffleId, dependency)
+ extends BaseShuffleHandle[K, V, C](shuffleId, dependency) {
+ val numMappers: Int = dependency.rdd.getNumPartitions
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala
new file mode 100644
index 0000000000..60132c984f
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala
@@ -0,0 +1,338 @@
+/*
+ * 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.util.concurrent.{CountDownLatch, TimeUnit}
+
+import scala.collection.mutable
+import scala.concurrent.{Await, ExecutionContext, Future}
+import scala.concurrent.duration.DurationInt
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
+import org.apache.spark.shuffle.{ShuffleHandle, ShuffleWriteMetricsReporter, ShuffleWriter}
+import org.apache.spark.shuffle.celeborn.CelebornShuffleHandle
+import org.apache.spark.sql.{CometTestBase, DataFrame, SparkSession}
+import org.apache.spark.sql.comet.shims.ShimCometShuffleMaterialization
+import org.apache.spark.util.ThreadUtils
+
+import org.apache.comet.CometConf
+
+class CometCelebornConcurrentMaterializationSuite extends CometTestBase {
+
+ import testImplicits._
+
+ override protected val shuffleManager: String =
+ classOf[CometCelebornConcurrentMaterializationTestManager].getName
+
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(CometConf.COMET_SHUFFLE_MODE.key, "native")
+ .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false")
+ .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k")
+ .set("spark.sql.adaptive.enabled", "false")
+ .set("spark.shuffle.compress", "false")
+ .set("spark.io.encryption.enabled", "false")
+ .set("spark.celeborn.client.spark.stageRerun.enabled", "true")
+
+ private def manager: CometCelebornConcurrentMaterializationTestManager =
+ SparkEnv.get.shuffleManager.asInstanceOf[CometCelebornConcurrentMaterializationTestManager]
+
+ private def branch(start: Long, oversized: Boolean): DataFrame = {
+ // Keep the large value dependent on the input, so Spark cannot replace the inner shuffle's
+ // payload with a small key and reconstruct a folded string above that exchange.
+ val payload = if (oversized) "concat(CAST(id AS STRING), repeat('x', 1048576))" else "'small'"
+ spark
+ .range(start, start + 1, 1, 1)
+ .selectExpr("id", s"$payload AS payload")
+ .repartition(2, $"id")
+ }
+
+ test("materialization restores the captured session without changing the worker's session") {
+ class CapturedSession extends ShimCometShuffleMaterialization {
+ def current: Option[SparkSession] = withCapturedSession(SparkSession.getActiveSession)
+ }
+ val captured = spark.withActive { new CapturedSession }
+ val worker = Future {
+ val previous = SparkSession.getActiveSession
+ SparkSession.clearActiveSession()
+ try {
+ assert(captured.current.contains(spark))
+ assert(SparkSession.getActiveSession.isEmpty)
+ } finally {
+ previous.foreach(SparkSession.setActiveSession)
+ }
+ }(ExecutionContext.global)
+ Await.result(worker, 20.seconds)
+ }
+
+ test("non-AQE sibling shuffles start before either remote materialization finishes") {
+ val gate = new ConcurrentRemoteMapStarts
+ manager.remoteMapStarts = Some(gate)
+ val query = branch(0, oversized = false).union(branch(100, oversized = false))
+ val plan = query.queryExecution.executedPlan
+ val exchanges = collect(plan) { case exchange: CometShuffleExchangeExec => exchange }
+ assert(exchanges.size == 2)
+ val localBefore = manager.localAttempts.size()
+ val construction = Future(exchanges.map(_.executeColumnar()))(ExecutionContext.global)
+ try {
+ // Both maps remain blocked. Reaching this point proves independent branches overlap,
+ // without relying on elapsed task timings or allowing one branch to finish first.
+ assert(gate.bothStarted.await(20, TimeUnit.SECONDS))
+ Await.result(construction, 20.seconds)
+ assert(gate.shuffleCount == 2)
+ gate.release.countDown()
+ exchanges.foreach { exchange =>
+ val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ val materialization = original.materialization.get
+ val statistics = Await.result(materialization, 20.seconds)
+ assert(statistics.shuffleId == original.shuffleId)
+ assert(materialization.completedDependency.exists(_ eq original))
+ }
+ assert(manager.localAttempts.size() == localBefore)
+ } finally {
+ gate.release.countDown()
+ Await.ready(construction, 20.seconds)
+ manager.remoteMapStarts = None
+ exchanges.foreach { exchange =>
+ val dependency = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ dependency.materialization.foreach(_.cancel())
+ manager.unregisterShuffle(dependency.shuffleId)
+ }
+ }
+ }
+
+ test("cancelling while an upstream shuffle materializes prevents downstream job admission") {
+ val gate = new ConcurrentRemoteMapStarts
+ manager.remoteMapStarts = Some(gate)
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val query = branch(0, oversized = true)
+ .selectExpr("id + 10 AS id", "payload")
+ .repartition(3, $"id")
+ val exchanges = collect(query.queryExecution.executedPlan) {
+ case exchange: CometShuffleExchangeExec => exchange
+ }
+ assert(exchanges.size == 2)
+ val outer = exchanges.find(_.outputPartitioning.numPartitions == 3).get
+ val construction = Future(outer.executeColumnar())(ExecutionContext.global)
+ try {
+ assert(gate.firstStarted.await(20, TimeUnit.SECONDS))
+ Await.result(construction, 20.seconds)
+ val dependency = outer.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ val materialization = dependency.materialization.get
+ assert(!materialization.isCompleted)
+ assert(materialization.jobIds.isEmpty)
+
+ // The outer worker is waiting for the blocked inner shuffle. It must not hold the state
+ // lock during that wait, and cancellation must fence its not-yet-submitted map-stage job.
+ val cancellation = Future(materialization.cancel(Some("Cancelled before job admission")))(
+ ExecutionContext.global)
+ Await.result(cancellation, 5.seconds)
+ assert(materialization.isCancelled)
+ assert(
+ materialization.value.exists(_.failed.get.getMessage ==
+ "Cancelled before job admission"))
+ gate.release.countDown()
+ exchanges.filterNot(_ eq outer).foreach { inner =>
+ val original = inner.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ Await.result(original.materialization.get, 20.seconds)
+ }
+ assert(materialization.jobIds.isEmpty)
+ } finally {
+ gate.release.countDown()
+ Await.ready(construction, 20.seconds)
+ manager.remoteMapStarts = None
+ exchanges.foreach { exchange =>
+ exchange.shuffleDependency
+ .asInstanceOf[CometShuffleDependency[_, _, _]]
+ .materialization
+ .foreach(_.cancel())
+ }
+ manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ manager.localAttempts.asScala.drop(localBefore).foreach { attempt =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ }
+ }
+
+ test("shared upstream materialization does not exhaust a bounded execution context") {
+ val workers = ThreadUtils.newDaemonFixedThreadPool(2, "comet-materialization-test")
+ val executionContext = ExecutionContext.fromExecutorService(workers)
+ val gate = new ConcurrentRemoteMapStarts
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val branchCount = 6
+ val shared = branch(0, oversized = true)
+ val query = (1 to branchCount)
+ .map { index =>
+ shared.selectExpr(s"id + $index AS id", "payload").repartition(3, $"id")
+ }
+ .reduce(_.union(_))
+ val exchanges = collect(query.queryExecution.executedPlan) {
+ case exchange: CometShuffleExchangeExec => exchange
+ }
+ val materializations =
+ mutable.ArrayBuffer.empty[CometCelebornShuffleMaterialization[_, _, _]]
+ val completedJobs = new CompletedMaterializationJobs
+ spark.sparkContext.addSparkListener(completedJobs)
+ manager.testMaterializationExecutionContext = Some(executionContext)
+ manager.remoteMapStarts = Some(gate)
+ try {
+ // Exchange reuse gives every sibling the same unfinished upstream destination. There are
+ // more dependents than workers, so blocking even a small bounded pool would starve the
+ // upstream completion callback as well as unrelated work scheduled on that pool.
+ assert(exchanges.size == branchCount + 1)
+ exchanges.foreach { exchange =>
+ exchange.executeColumnar()
+ materializations += exchange.shuffleDependency
+ .asInstanceOf[CometShuffleDependency[_, _, _]]
+ .materialization
+ .get
+ }
+ assert(gate.firstStarted.await(20, TimeUnit.SECONDS))
+ assert(gate.shuffleCount == 1)
+ assert(materializations.forall(!_.isCompleted))
+ Await.result(Future(())(executionContext), 5.seconds)
+
+ gate.release.countDown()
+ materializations.foreach { materialization =>
+ Await.result(materialization, 30.seconds)
+ assert(materialization.completedDependency.exists(_.useLocalShuffle))
+ assert(materialization.jobIds.size == 2)
+ }
+ val actual = query.collect().map(row => (row.getLong(0), row.getString(1))).toSeq
+ val expected = (1 to branchCount).map(index => (index.toLong, "0" + "x" * 1048576))
+ assert(actual.sortBy(_._1) == expected)
+ } finally {
+ gate.release.countDown()
+ materializations.foreach(_.cancel())
+ completedJobs.awaitCompletion(materializations.flatMap(_.jobIds).toSeq)
+ Await.result(Future(())(executionContext), 20.seconds)
+ spark.sparkContext.removeSparkListener(completedJobs)
+ manager.remoteMapStarts = None
+ manager.testMaterializationExecutionContext = None
+ executionContext.shutdown()
+ assert(executionContext.awaitTermination(20, TimeUnit.SECONDS))
+ manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ manager.localAttempts.asScala.drop(localBefore).foreach { attempt =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ }
+ }
+
+ for (nested <- Seq(false, true)) {
+ test(s"non-AQE fallback preserves concurrent branches and downstream reads: nested=$nested") {
+ val gate = new ConcurrentRemoteMapStarts
+ manager.remoteMapStarts = Some(gate)
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val left = if (nested) {
+ // The changed partition key keeps a second shuffle above the first one. Submitting its
+ // materialization must wait for the first destination without blocking the right branch.
+ branch(0, oversized = true)
+ .selectExpr("id + 10 AS id", "payload")
+ .repartition(3, $"id")
+ } else {
+ branch(0, oversized = true)
+ }
+ val query = left.union(branch(100, oversized = true))
+ val exchanges = collect(query.queryExecution.executedPlan) {
+ case exchange: CometShuffleExchangeExec => exchange
+ }
+ assert(exchanges.size == (if (nested) 3 else 2))
+ assert(exchanges.forall(_.output.exists(_.name == "payload")))
+ val execution = Future(query.collect())(ExecutionContext.global)
+ try {
+ assert(gate.bothStarted.await(20, TimeUnit.SECONDS))
+ assert(gate.shuffleCount == 2)
+ gate.release.countDown()
+ val actual =
+ Await.result(execution, 30.seconds).map(row => (row.getLong(0), row.getString(1))).toSeq
+ val leftId = if (nested) 10L else 0L
+ assert(
+ actual
+ .sortBy(_._1) == Seq((leftId, "0" + "x" * 1048576), (100L, "100" + "x" * 1048576)))
+ assert(manager.localAttempts.size() > localBefore)
+ exchanges.foreach { exchange =>
+ val dependency =
+ exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ assert(dependency.materialization.get.completedDependency.exists(_.useLocalShuffle))
+ }
+ } finally {
+ gate.release.countDown()
+ Await.ready(execution, 30.seconds)
+ manager.remoteMapStarts = None
+ manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ manager.localAttempts.asScala.drop(localBefore).foreach { attempt =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ }
+ }
+ }
+}
+
+/** Holds two distinct remote map stages so the tests can prove their execution overlaps. */
+private[shuffle] class ConcurrentRemoteMapStarts {
+ val firstStarted = new CountDownLatch(1)
+ val bothStarted = new CountDownLatch(2)
+ val release = new CountDownLatch(1)
+ private val shuffleIds = mutable.Set.empty[Int]
+
+ def shuffleCount: Int = synchronized { shuffleIds.size }
+
+ def mapStarted(shuffleId: Int): Unit = {
+ synchronized {
+ if (shuffleIds.add(shuffleId)) {
+ firstStarted.countDown()
+ bothStarted.countDown()
+ }
+ }
+ require(release.await(20, TimeUnit.SECONDS), "The test must release concurrent remote maps")
+ }
+}
+
+class CometCelebornConcurrentMaterializationTestManager(conf: SparkConf, isDriver: Boolean)
+ extends CometCelebornFallbackTestShuffleManager(conf, isDriver) {
+
+ @volatile private[shuffle] var remoteMapStarts: Option[ConcurrentRemoteMapStarts] = None
+ @volatile private[shuffle] var testMaterializationExecutionContext: Option[ExecutionContext] =
+ None
+
+ override protected[shuffle] def materializationExecutionContext: ExecutionContext =
+ testMaterializationExecutionContext.getOrElse(super.materializationExecutionContext)
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]]) {
+ remoteMapStarts.foreach(_.mapStarted(handle.shuffleId))
+ }
+ super.getWriter(handle, mapId, context, metrics)
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornLocalFetchFailureSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornLocalFetchFailureSuite.scala
new file mode 100644
index 0000000000..85829401d4
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornLocalFetchFailureSuite.scala
@@ -0,0 +1,139 @@
+/*
+ * 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.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
+import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
+import org.apache.spark.shuffle.{FetchFailedException, ShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter}
+import org.apache.spark.sql.CometTestBase
+
+import org.apache.comet.CometConf
+
+class CometCelebornLocalFetchFailureSuite extends CometTestBase {
+
+ import testImplicits._
+
+ override protected val shuffleManager: String =
+ classOf[CometCelebornLocalFetchFailureTestManager].getName
+
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(CometConf.COMET_SHUFFLE_MODE.key, "native")
+ .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false")
+ .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k")
+ .set("spark.shuffle.compress", "false")
+ .set("spark.io.encryption.enabled", "false")
+ .set("spark.celeborn.client.spark.stageRerun.enabled", "true")
+ .set("spark.stage.maxConsecutiveAttempts", "4")
+
+ test("local fetch failure after fallback recovers a partially completed result stage") {
+ val manager = SparkEnv.get.shuffleManager
+ .asInstanceOf[CometCelebornLocalFetchFailureTestManager]
+ val listener = new SparkListener {
+ override def onTaskEnd(event: SparkListenerTaskEnd): Unit = {
+ if (event.taskType == "ResultTask" && event.taskInfo.successful) {
+ manager.completedResultTasks.incrementAndGet()
+ }
+ }
+ }
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ // Both deterministic input partitions exceed the remote frame limit, forcing fresh
+ // local materialization. Keys 1 and 2 reach different reducers, so the result partition
+ // that completes before the fetch failure contains an actual row.
+ val query = spark
+ .range(1, 3, 1, 2)
+ .selectExpr("CAST(id AS INT) AS key", "repeat('x', 1048576) AS payload")
+ .repartition(2, $"key")
+
+ val actual = query.collect().map(row => (row.getInt(0), row.getString(1))).toSeq
+
+ assert(actual.sortBy(_._1) == Seq((1, "x" * (1024 * 1024)), (2, "x" * (1024 * 1024))))
+ assert(manager.injectedFetchFailures.get() == 1)
+ assert(manager.completedResultsBeforeFailure.get() > 0)
+ assert(!manager.remoteAttempts.isEmpty)
+ val failedShuffleId = manager.failedLocalShuffleId.get()
+ assert(manager.remoteAttempts.asScala.forall(_._1.shuffleId != failedShuffleId))
+ assert(manager.localAttempts.asScala.exists { attempt =>
+ attempt.shuffleId == failedShuffleId && attempt.stageAttempt > 0
+ })
+ }
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+}
+
+/** Injects one local block loss only after Spark accepts another result partition. */
+class CometCelebornLocalFetchFailureTestManager(conf: SparkConf, isDriver: Boolean)
+ extends CometCelebornFallbackTestShuffleManager(conf, isDriver) {
+
+ // This local-mode fixture keeps state in its SparkContext-owned shuffle manager.
+ val completedResultTasks = new AtomicInteger()
+ val completedResultsBeforeFailure = new AtomicInteger()
+ val injectedFetchFailures = new AtomicInteger()
+ val failedLocalShuffleId = new AtomicInteger(-1)
+ private val failNextLocalFetch = new AtomicBoolean(true)
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ if (handle.isInstanceOf[CometNativeShuffleHandle[_, _]] &&
+ context.partitionId() == 1 &&
+ failNextLocalFetch.compareAndSet(true, false)) {
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20)
+ while (completedResultTasks.get() == 0 && System.nanoTime() < deadline) {
+ Thread.sleep(10)
+ }
+ val completed = completedResultTasks.get()
+ require(completed > 0, "Another result partition must finish before the local fetch fails")
+ completedResultsBeforeFailure.set(completed)
+ failedLocalShuffleId.set(handle.shuffleId)
+ injectedFetchFailures.incrementAndGet()
+ throw new FetchFailedException(
+ SparkEnv.get.blockManager.shuffleServerId,
+ handle.shuffleId,
+ -1L,
+ 0,
+ startPartition,
+ "A local shuffle block was lost after another result partition completed",
+ null)
+ }
+ super.getReader(
+ handle,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ }
+}
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 c3f291af85..a300bf8de9 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
@@ -28,7 +28,7 @@ import org.apache.spark.shuffle.{BaseShuffleHandle, FetchFailedException, IndexS
import org.apache.spark.sql.CometTestBase
import org.apache.spark.sql.vectorized.ColumnarBatch
-import org.apache.comet.CometConf
+import org.apache.comet.{CometConf, CometShuffleSizeLimitException}
import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory, RecordingCelebornPushClient}
/** Exercises real Spark map tasks, native RSS planning, and the Celeborn map lifecycle. */
@@ -105,6 +105,22 @@ class CometCelebornNativeShuffleWriterSuite extends CometTestBase {
}
}
+ test("default pushers and the configured factory can share executor byte admission") {
+ val client = new RecordingCelebornPushClient
+ val defaultPusher = new CelebornShufflePartitionPusher(client, 19, 3, 0, 12, 9)
+ val configuredPusher = CelebornShufflePusherFactory.create(
+ new SparkConf(false),
+ client,
+ 19,
+ 12,
+ 9,
+ TaskContext.empty())
+
+ assert(defaultPusher.maxReservationBytes() == configuredPusher.maxReservationBytes())
+ defaultPusher.abort()
+ configuredPusher.abort()
+ }
+
test("configured frame limits and executor admission both constrain native RSS callbacks") {
val context = TaskContext.empty()
val configured = new SparkConf(false)
@@ -272,6 +288,282 @@ class CometCelebornNativeShuffleWriterSuite extends CometTestBase {
}
}
+ test("native frame limit failures abort before requesting a shuffle fallback") {
+ withNativeShuffleDependency() { dependency =>
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val client = new RecordingCelebornPushClient
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf.clone().set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "20"),
+ client,
+ 96,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ val restart =
+ new FetchFailedException(
+ null,
+ dependency.shuffleId,
+ -1L,
+ -1,
+ 0,
+ "restart locally",
+ null)
+ var fallbackCalls = 0
+ var reportedFailure: Throwable = null
+ var cleanupBeforeFallback = false
+ val writer = CometCelebornNativeShuffleWriterSuite.newWriter(
+ dependency,
+ context,
+ pusher,
+ onSizeLimitExceeded = failure => {
+ fallbackCalls += 1
+ reportedFailure = failure
+ cleanupBeforeFallback = client.cleanupCalls.get() == 1
+ throw restart
+ })
+ val actual =
+ try {
+ writer.write(inputs)
+ null
+ } catch {
+ case failure: FetchFailedException => failure
+ }
+
+ (
+ actual eq restart,
+ fallbackCalls,
+ reportedFailure.isInstanceOf[CometShuffleSizeLimitException],
+ cleanupBeforeFallback,
+ client.pushCount,
+ client.mapperEndCalls.get(),
+ writer.mapStatus == null,
+ writer.stop(success = false).isEmpty)
+ })
+
+ assert(results.nonEmpty)
+ assert(results.forall(_._1))
+ assert(results.forall(_._2 == 1))
+ assert(results.forall(_._3))
+ assert(results.forall(_._4))
+ assert(results.forall(_._5 == 0))
+ assert(results.forall(_._6 == 0))
+ assert(results.forall(_._7))
+ assert(results.forall(_._8))
+ }
+ }
+
+ test("a size failure after a successful push discards the partial remote map") {
+ withNativeShuffleDependency() { dependency =>
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val expected = new CometShuffleSizeLimitException("single row exceeds RSS frame limit")
+ val client = new RecordingCelebornPushClient {
+ override 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 = {
+ if (pushCount == 1) failure = expected
+ super.pushOrMergeData(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ bytes,
+ offset,
+ length,
+ numMappers,
+ numPartitions,
+ doPush,
+ skipCompress)
+ }
+ }
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf,
+ client,
+ 97,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ var fallbackCalls = 0
+ var cleanupBeforeFallback = false
+ val writer = CometCelebornNativeShuffleWriterSuite.newWriter(
+ dependency,
+ context,
+ pusher,
+ onSizeLimitExceeded = failure => {
+ assert(failure eq expected)
+ fallbackCalls += 1
+ cleanupBeforeFallback = client.cleanupCalls.get() == 1
+ })
+ val actual =
+ try {
+ writer.write(inputs)
+ null
+ } catch {
+ case failure: CometShuffleSizeLimitException => failure
+ }
+
+ (
+ actual eq expected,
+ fallbackCalls,
+ cleanupBeforeFallback,
+ client.pushCount,
+ client.mapperEndCalls.get(),
+ writer.mapStatus == null,
+ writer.stop(success = false).isEmpty)
+ })
+
+ assert(results.nonEmpty)
+ assert(results.forall(_._1))
+ assert(results.forall(_._2 == 1))
+ assert(results.forall(_._3))
+ assert(results.forall(_._4 == 2))
+ assert(results.forall(_._5 == 0))
+ assert(results.forall(_._6))
+ assert(results.forall(_._7))
+ }
+ }
+
+ test("transport failures do not request fallback based on their error message") {
+ withNativeShuffleDependency() { dependency =>
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val expected = new IOException("single row exceeds RSS frame limit")
+ val client = new RecordingCelebornPushClient
+ client.failure = expected
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf,
+ client,
+ 98,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ var fallbackCalls = 0
+ val writer = CometCelebornNativeShuffleWriterSuite.newWriter(
+ dependency,
+ context,
+ pusher,
+ onSizeLimitExceeded = _ => fallbackCalls += 1)
+ val actual =
+ try {
+ writer.write(inputs)
+ null
+ } catch {
+ case failure: IOException => failure
+ }
+
+ (
+ actual eq expected,
+ fallbackCalls,
+ 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))
+ }
+ }
+
+ for (adaptive <- Seq(false, true)) {
+ test(
+ s"a row exceeding the RSS frame limit roundtrips through local Comet shuffle: AQE=$adaptive") {
+ // Exercise the reported 70 MiB row once; the AQE variant uses the same limit crossing at a
+ // smaller scale. Disabling compression makes the encoded size deterministic.
+ val payloadBytes = if (adaptive) 1024 * 1024 else 70 * 1024 * 1024
+ val payload = "x" * payloadBytes
+ val rows = Seq((0, "small"), (1, payload))
+ val runtimeConf = SparkEnv.get.conf
+ val previousCompression = runtimeConf.getOption("spark.shuffle.compress")
+ runtimeConf.set("spark.shuffle.compress", "false")
+ try {
+ withSQLConf(
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ "spark.sql.adaptive.enabled" -> adaptive.toString) {
+ withTempPath { path =>
+ rows
+ .toDF("key", "payload")
+ .coalesce(1)
+ .write
+ .option("parquet.enable.dictionary", "false")
+ .parquet(path.getCanonicalPath)
+ val shuffled = spark.read.parquet(path.getCanonicalPath).repartition(3, $"key")
+ val dependency = collect(shuffled.queryExecution.executedPlan) {
+ case exchange: CometShuffleExchangeExec =>
+ exchange.shuffleDependency
+ .asInstanceOf[CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]]
+ }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange"))
+ val numMappers = dependency.rdd.getNumPartitions
+ val frameBytes = if (adaptive) "32k" else "64m"
+ val remoteAttempts = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val client = new RecordingCelebornPushClient
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf
+ .clone()
+ .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, frameBytes),
+ client,
+ 99,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ var fallbackRequested = false
+ val writer = CometCelebornNativeShuffleWriterSuite.newWriter(
+ dependency,
+ context,
+ pusher,
+ onSizeLimitExceeded = _ => fallbackRequested = true)
+ try {
+ writer.write(inputs)
+ } catch {
+ case _: CometShuffleSizeLimitException =>
+ }
+ val failedWithoutCommit =
+ fallbackRequested && client.cleanupCalls.get() == 1 &&
+ client.mapperEndCalls.get() == 0 && writer.mapStatus == null
+ writer.stop(success = false)
+ failedWithoutCommit
+ })
+ assert(remoteAttempts.exists(identity))
+
+ // The dependency is registered with the suite's real local Comet manager. Running
+ // it again exercises native file output, Spark map-status publication, and the Comet
+ // reader, including a frame larger than the remote limit and normal AQE stage reads.
+ val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq
+ assert(actual.sortBy(_._1) == rows)
+ }
+ }
+ } finally {
+ previousCompression match {
+ case Some(value) => runtimeConf.set("spark.shuffle.compress", value)
+ case None => runtimeConf.remove("spark.shuffle.compress")
+ }
+ }
+ }
+ }
+
test("low-cardinality range shuffle plans use the actual reducer partition count") {
withNativeShuffleDependency(rangePartitioning = true) { dependency =>
val actualPartitions = dependency.partitioner.numPartitions
@@ -424,7 +716,9 @@ private[shuffle] object CometCelebornNativeShuffleWriterSuite {
context: TaskContext,
pusher: CelebornShufflePartitionPusher,
commitAuthorized: Boolean = false,
- commitValidator: () => Boolean = () => true): CometNativeShuffleWriter[Int, ColumnarBatch] =
+ commitValidator: () => Boolean = () => true,
+ onSizeLimitExceeded: Throwable => Unit = _ => ())
+ : CometNativeShuffleWriter[Int, ColumnarBatch] =
new CometNativeShuffleWriter[Int, ColumnarBatch](
dependency.nativeShuffleSpec.get,
dependency.outputPartitioning.get,
@@ -442,7 +736,8 @@ private[shuffle] object CometCelebornNativeShuffleWriterSuite {
pusher.maxFrameBytes(),
dependency.partitioner.numPartitions,
commitAuthorized,
- commitValidator)))
+ commitValidator,
+ onSizeLimitExceeded)))
final class LocalFallbackShuffleManager extends ShuffleManager {
var unregisteredShuffle: Option[Int] = None
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala
new file mode 100644
index 0000000000..826ac7e86d
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala
@@ -0,0 +1,808 @@
+/*
+ * 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.util.Properties
+import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.AtomicInteger
+
+import scala.collection.mutable
+import scala.concurrent.Await
+import scala.concurrent.duration.DurationInt
+import scala.jdk.CollectionConverters._
+import scala.language.existentials
+
+import org.apache.spark.{FutureAction, MapOutputStatistics, MapOutputTrackerMaster, ShuffleDependency, SparkConf, SparkEnv, TaskContext}
+import org.apache.spark.scheduler.{JobFailed, MapStatus, SparkListener, SparkListenerJobEnd, SparkListenerJobStart}
+import org.apache.spark.shuffle.{ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+import org.apache.spark.shuffle.celeborn.CelebornShuffleHandle
+import org.apache.spark.sql.CometTestBase
+
+import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
+import org.apache.comet.shuffle.{CelebornShufflePusherFactory, RecordingCelebornPushClient, ResolvedCelebornShufflePusher}
+
+/**
+ * Real Spark stage recovery and local Comet I/O; only the optional Celeborn client is replaced.
+ */
+class CometCelebornShuffleFallbackSuite extends CometTestBase {
+
+ import testImplicits._
+
+ override protected val shuffleManager: String =
+ classOf[CometCelebornFallbackTestShuffleManager].getName
+
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(CometConf.COMET_SHUFFLE_MODE.key, "native")
+ .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false")
+ .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k")
+ .set("spark.shuffle.compress", "false")
+ .set("spark.io.encryption.enabled", "false")
+ .set("spark.celeborn.client.spark.stageRerun.enabled", "true")
+ .set("spark.stage.maxConsecutiveAttempts", "3")
+
+ private def manager: CometCelebornFallbackTestShuffleManager =
+ SparkEnv.get.shuffleManager.asInstanceOf[CometCelebornFallbackTestShuffleManager]
+
+ for (adaptive <- Seq(false, true)) {
+ test(s"an oversized row materializes a fresh local shuffle and succeeds: AQE=$adaptive") {
+ withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) {
+ val rows = Seq((0, "small"), (1, "x" * (1024 * 1024)))
+ withTempPath { path =>
+ rows
+ .toDF("key", "payload")
+ .coalesce(1)
+ .write
+ .option("parquet.enable.dictionary", "false")
+ .parquet(path.getCanonicalPath)
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val readsBefore = manager.nativeReads.get()
+ val shuffled = spark.read.parquet(path.getCanonicalPath).repartition(3, $"key")
+
+ val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq
+
+ assert(actual.sortBy(_._1) == rows)
+ val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq
+ assert(remote.size == 1, s"expected one failed remote map attempt, got $remote")
+ val local = manager.localAttempts.asScala.drop(localBefore).toSeq
+ assert(local.size == 1, s"expected a complete local map stage: $local")
+ assert(local.head.shuffleId != remote.head._1.shuffleId)
+ assert(local.head.stageId != remote.head._1.stageId)
+ assert(local.head.stageAttempt == 0)
+ assert(manager.nativeReads.get() > readsBefore)
+ val client = remote.head._2
+ assert(client.cleanupCalls.get() == 1)
+ assert(client.mapperEndCalls.get() == 0)
+ assert(client.fetchFailureReports.get() == 1)
+
+ // Each destination has its own shuffle lifecycle after the successful replacement.
+ assert(manager.unregisterShuffle(remote.head._1.shuffleId))
+ assert(manager.unregisterShuffle(local.head.shuffleId))
+ assert(client.shuffleCleanupCalls.get() == 1)
+ }
+ }
+ }
+
+ test(s"fallback recomputes a previously completed remote map: AQE=$adaptive") {
+ withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) {
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val readsBefore = manager.nativeReads.get()
+ manager.remoteMapsBeforeFailure.set(0)
+ manager.waitForFirstRemoteMap = true
+ try {
+ // Range partition zero produces the small row; partition one produces the oversized
+ // row. Delay the second writer until Spark has registered the first remote map output.
+ val shuffled = spark
+ .range(0, 2, 1, 2)
+ .selectExpr(
+ "CAST(id AS INT) AS key",
+ "CASE WHEN id = 0 THEN 'small' ELSE repeat('x', 1048576) END AS payload")
+ .repartition(3, $"key")
+ val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq
+
+ assert(actual.sortBy(_._1) == Seq((0, "small"), (1, "x" * (1024 * 1024))))
+ assert(manager.remoteMapsBeforeFailure.get() == 1)
+ val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq
+ assert(remote.size == 2)
+ assert(remote.map(_._1.shuffleId).distinct.size == 1)
+ assert(remote.forall(_._1.stageAttempt == 0))
+ assert(remote.count(_._2.mapperEndCalls.get() == 1) == 1)
+ assert(remote.map(_._2.fetchFailureReports.get()).sum == 1)
+ val local = manager.localAttempts.asScala.drop(localBefore).toSeq
+ assert(local.map(_.partitionId).sorted == Seq(0, 1))
+ assert(local.map(_.shuffleId).distinct.size == 1)
+ assert(local.map(_.stageId).distinct.size == 1)
+ assert(local.forall(_.shuffleId != remote.head._1.shuffleId))
+ assert(local.forall(_.stageId != remote.head._1.stageId))
+ assert(local.forall(_.stageAttempt == 0))
+ assert(manager.nativeReads.get() > readsBefore)
+ assert(manager.unregisterShuffle(remote.head._1.shuffleId))
+ assert(manager.unregisterShuffle(local.head.shuffleId))
+ } finally {
+ manager.waitForFirstRemoteMap = false
+ }
+ }
+ }
+
+ test(s"delayed remote map completion cannot replace fresh local output: AQE=$adaptive") {
+ withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) {
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val delayed = new DelayedRemoteMapCompletion
+ manager.delayedRemoteCompletion = Some(delayed)
+ try {
+ val shuffled = spark
+ .range(0, 2, 1, 2)
+ .selectExpr(
+ "CAST(id AS INT) AS key",
+ "CASE WHEN id = 0 THEN 'small' ELSE repeat('x', 1048576) END AS payload")
+ .repartition(3, $"key")
+
+ val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq
+
+ assert(actual.sortBy(_._1) == Seq((0, "small"), (1, "x" * (1024 * 1024))))
+ assert(delayed.remoteReturned.await(20, TimeUnit.SECONDS))
+ val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq
+ val local = manager.localAttempts.asScala.drop(localBefore).toSeq
+ assert(remote.size == 2)
+ assert(remote.count(_._2.mapperEndCalls.get() == 1) == 1)
+ assert(local.map(_.partitionId).sorted == Seq(0, 1))
+ assert(local.map(_.shuffleId).distinct.size == 1)
+ assert(local.map(_.stageId).distinct.size == 1)
+ assert(local.forall(_.shuffleId != remote.head._1.shuffleId))
+ assert(local.forall(_.stageId != remote.head._1.stageId))
+ assert(local.forall(_.stageAttempt == 0))
+ assert(manager.unregisterShuffle(remote.head._1.shuffleId))
+ assert(manager.unregisterShuffle(local.head.shuffleId))
+ } finally {
+ // Unblock an old writer even if the query fails before the local stage starts.
+ delayed.localStageStarted.countDown()
+ manager.delayedRemoteCompletion = None
+ }
+ }
+ }
+ }
+
+ test("successful remote materialization publishes the original shuffle identity") {
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val readsBefore = manager.nativeReads.get()
+ val query = spark.range(0, 2, 1, 2).repartition(2, $"id")
+ val exchange = collect(query.queryExecution.executedPlan) {
+ case value: CometShuffleExchangeExec => value
+ }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange"))
+ val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ val action = exchange.mapOutputStatisticsFuture
+ .asInstanceOf[FutureAction[MapOutputStatistics]]
+ try {
+ val statistics = Await.result(action, 20.seconds)
+
+ assert(statistics.shuffleId == original.shuffleId)
+ assert(statistics.bytesByPartitionId.sum > 0L)
+ assert(original.currentShuffleDependency eq original)
+ // Spark 4 exposes this ID to AQE; the Spark 3 shim supplies an unused placeholder.
+ if (isSpark40Plus) assert(exchange.shuffleId == original.shuffleId)
+ assert(action.jobIds.size == 1)
+ val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq
+ assert(remote.map(_._1.partitionId).sorted == Seq(0, 1))
+ assert(remote.forall(_._1.shuffleId == original.shuffleId))
+ assert(remote.forall(_._2.mapperEndCalls.get() == 1))
+ assert(manager.localAttempts.size() == localBefore)
+ assert(manager.nativeReads.get() == readsBefore)
+ } finally {
+ action.cancel()
+ manager.unregisterShuffle(original.shuffleId)
+ }
+ }
+ }
+
+ test("a remote stage failure cannot abort an active local replacement") {
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ val sc = spark.sparkContext
+ val previousProperties = sc.getLocalProperties.clone().asInstanceOf[Properties]
+ val groupId = "comet-remote-failure-isolation"
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val paused = new PausedShuffleMapCompletion(local = true)
+ val jobStart = new PausedReplacementJobStart(groupId)
+ val completedJobs = new CompletedMaterializationJobs
+ manager.pausedMapCompletion = Some(paused)
+ manager.remoteFailureAfterLocalMap = Some(paused)
+ // Delay the materialization's JobStart listener, so the remote size failure reaches
+ // Spark before Comet can cancel that job. Observe job failures on an independent queue.
+ sc.addSparkListener(jobStart)
+ sc.listenerBus.addToQueue(completedJobs, groupId)
+ var active: Option[FutureAction[MapOutputStatistics]] = None
+ try {
+ sc.setJobGroup(groupId, "Isolate the local replacement from remote failure", true)
+ val query = spark
+ .range(0, 1, 1, 1)
+ .selectExpr("CAST(id AS INT) AS key", "repeat('x', 1048576) AS payload")
+ .repartition(2, $"key")
+ val exchange = collect(query.queryExecution.executedPlan) {
+ case value: CometShuffleExchangeExec => value
+ }.head
+ val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ val action = exchange.mapOutputStatisticsFuture
+ .asInstanceOf[FutureAction[MapOutputStatistics]]
+ active = Some(action)
+ assert(jobStart.entered.await(20, TimeUnit.SECONDS))
+ assert(paused.stopped.await(20, TimeUnit.SECONDS))
+ assert(action.jobIds.size == 2)
+ val remoteJobId = action.jobIds.head
+ completedJobs.awaitCompletion(Seq(remoteJobId))
+ assert(completedJobs.hasFailed(remoteJobId))
+
+ jobStart.release.countDown()
+ paused.release.countDown()
+ val statistics = Await.result(action, 20.seconds)
+
+ val selected = original.currentShuffleDependency
+ assert(selected.useLocalShuffle)
+ assert(selected.rdd ne original.rdd)
+ assert(statistics.shuffleId == selected.shuffleId)
+ completedJobs.awaitCompletion(action.jobIds)
+ assert(!completedJobs.hasFailed(action.jobIds.last))
+ assert(query.collect().length == 1)
+ } finally {
+ jobStart.release.countDown()
+ paused.release.countDown()
+ active.foreach(_.cancel())
+ sc.setLocalProperties(previousProperties)
+ manager.pausedMapCompletion = None
+ manager.remoteFailureAfterLocalMap = None
+ sc.removeSparkListener(jobStart)
+ sc.removeSparkListener(completedJobs)
+ manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ manager.localAttempts.asScala.drop(localBefore).foreach { attempt =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ }
+ }
+ }
+
+ test("job-group cancellation while registering fallback does not submit a replacement job") {
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ val sc = spark.sparkContext
+ val previousProperties = sc.getLocalProperties.clone().asInstanceOf[Properties]
+ val groupId = "comet-materialization-cancel-during-local-registration"
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val registration = new PausedLocalShuffleRegistration
+ val completedJobs = new CompletedMaterializationJobs
+ manager.pausedLocalRegistration = Some(registration)
+ // The materialization holds its lock during registration, so its own listener can block.
+ // An independent queue lets the test observe Spark's cancellation before releasing it.
+ sc.listenerBus.addToQueue(completedJobs, groupId)
+ var active: Option[FutureAction[MapOutputStatistics]] = None
+ try {
+ sc.setJobGroup(groupId, "Cancel while local shuffle registration is paused", true)
+ val query = spark
+ .range(0, 1, 1, 1)
+ .selectExpr("CAST(id AS INT) AS key", "repeat('x', 1048576) AS payload")
+ .repartition(2, $"key")
+ val exchange = collect(query.queryExecution.executedPlan) {
+ case value: CometShuffleExchangeExec => value
+ }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange"))
+ val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ val action = exchange.mapOutputStatisticsFuture
+ .asInstanceOf[FutureAction[MapOutputStatistics]]
+ active = Some(action)
+ assert(registration.entered.await(20, TimeUnit.SECONDS))
+ val remoteJobId = completedJobs.awaitJobInGroup(groupId)
+
+ // Fallback has won, but submitMapStage(local) has not happened. Spark's ordinary
+ // cancelJobGroup only cancels current jobs, so a later replacement must be fenced.
+ sc.cancelJobGroup(groupId)
+ completedJobs.awaitCompletion(Seq(remoteJobId))
+ assert(completedJobs.hasFailed(remoteJobId))
+ registration.release.countDown()
+ assert(registration.returned.await(20, TimeUnit.SECONDS))
+ Await.ready(action, 20.seconds)
+
+ assert(action.value.exists(_.isFailure))
+ assert(original.materialization.get.completedDependency.isEmpty)
+ assert(action.jobIds == Seq(remoteJobId))
+ assert(manager.localAttempts.size() == localBefore)
+ assert(manager.remoteAttempts.size() == remoteBefore + 1)
+ } finally {
+ registration.release.countDown()
+ active.foreach(_.cancel())
+ sc.setLocalProperties(previousProperties)
+ manager.pausedLocalRegistration = None
+ sc.removeSparkListener(completedJobs)
+ manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ if (registration.shuffleId.get() >= 0) {
+ manager.unregisterShuffle(registration.shuffleId.get())
+ }
+ }
+ }
+ }
+
+ for (cancelLocal <- Seq(false, true)) {
+ test(s"cancelling materialization ends every submitted job: local=$cancelLocal") {
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ val remoteBefore = manager.remoteAttempts.size()
+ val localBefore = manager.localAttempts.size()
+ val paused = new PausedShuffleMapCompletion(cancelLocal)
+ val completedJobs = new CompletedMaterializationJobs
+ manager.pausedMapCompletion = Some(paused)
+ spark.sparkContext.addSparkListener(completedJobs)
+ var active: Option[FutureAction[MapOutputStatistics]] = None
+ try {
+ val payload = if (cancelLocal) "repeat('x', 1048576)" else "'small'"
+ val query = spark
+ .range(0, 1, 1, 1)
+ .selectExpr("CAST(id AS INT) AS key", s"$payload AS payload")
+ .repartition(2, $"key")
+ val exchange = collect(query.queryExecution.executedPlan) {
+ case value: CometShuffleExchangeExec => value
+ }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange"))
+ val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ val action = exchange.mapOutputStatisticsFuture
+ .asInstanceOf[FutureAction[MapOutputStatistics]]
+ active = Some(action)
+ assert(paused.stopped.await(20, TimeUnit.SECONDS))
+
+ // stop(true) has already produced a real MapStatus, but Spark has not received it.
+ // Cancelling here must prevent that delayed completion from publishing a dependency.
+ if (cancelLocal) {
+ original.materialization.get.cancel(Some("Cancelled during local materialization"))
+ } else {
+ action.cancel()
+ }
+ paused.release.countDown()
+ assert(paused.returned.await(20, TimeUnit.SECONDS))
+ Await.ready(action, 20.seconds)
+ completedJobs.awaitCompletion(action.jobIds)
+
+ assert(action.isCancelled)
+ assert(action.value.exists(_.isFailure))
+ if (cancelLocal) {
+ assert(
+ action.value.get.failed.get.getMessage == "Cancelled during local materialization")
+ }
+ assert(original.materialization.get.completedDependency.isEmpty)
+ assert(action.jobIds.size == (if (cancelLocal) 2 else 1))
+ val local = manager.localAttempts.asScala.drop(localBefore).toSeq
+ assert(local.size == (if (cancelLocal) 1 else 0))
+ assert(local.forall(_.shuffleId != original.shuffleId))
+ val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq
+ assert(remote.size == 1)
+ } finally {
+ paused.release.countDown()
+ active.foreach(_.cancel())
+ manager.pausedMapCompletion = None
+ spark.sparkContext.removeSparkListener(completedJobs)
+ manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ manager.localAttempts.asScala.drop(localBefore).foreach { attempt =>
+ manager.unregisterShuffle(attempt.shuffleId)
+ }
+ }
+ }
+ }
+ }
+
+}
+
+private[shuffle] case class CometShuffleFallbackAttempt(
+ shuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ partitionId: Int)
+
+/** Synchronizes one old remote map's validated completion with the new local stage. */
+private[shuffle] class DelayedRemoteMapCompletion {
+ val remoteCommitted = new CountDownLatch(1)
+ val localStageStarted = new CountDownLatch(1)
+ val remoteReturned = new CountDownLatch(1)
+
+ def awaitRemoteCommit(): Unit =
+ require(
+ remoteCommitted.await(20, TimeUnit.SECONDS),
+ "The small remote map must commit before the oversized map starts")
+
+ def awaitLocalStage(): Unit = {
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20)
+ var interrupted = false
+ var started = false
+ try {
+ while (!started && System.nanoTime() < deadline) {
+ try {
+ started = localStageStarted.await(
+ math.max(1L, deadline - System.nanoTime()),
+ TimeUnit.NANOSECONDS)
+ } catch {
+ case _: InterruptedException => interrupted = true
+ }
+ }
+ require(started, "A fresh local stage must start before the old map result is returned")
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt()
+ }
+ }
+ }
+}
+
+/** Holds one already computed map result so cancellation races with a concrete completion. */
+private[shuffle] class PausedShuffleMapCompletion(val local: Boolean) {
+ val stopped = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ val returned = new CountDownLatch(1)
+
+ def awaitRelease(): Unit = {
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20)
+ var interrupted = false
+ var released = false
+ try {
+ while (!released && System.nanoTime() < deadline) {
+ try {
+ released =
+ release.await(math.max(1L, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)
+ } catch {
+ case _: InterruptedException => interrupted = true
+ }
+ }
+ require(released, "The test must release the completed map")
+ } finally {
+ if (interrupted) Thread.currentThread().interrupt()
+ }
+ }
+}
+
+/** Holds listener delivery while Spark processes the old remote stage's failure. */
+private[shuffle] class PausedReplacementJobStart(groupId: String) extends SparkListener {
+ val entered = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ private var startedJobs = 0
+
+ override def onJobStart(event: SparkListenerJobStart): Unit = {
+ if (Option(event.properties).exists(_.getProperty("spark.jobGroup.id") == groupId)) {
+ startedJobs += 1
+ if (startedJobs == 2) {
+ entered.countDown()
+ require(release.await(20, TimeUnit.SECONDS), "The test must release local JobStart")
+ }
+ }
+ }
+}
+
+/** Holds fallback registration before the materialization can submit its local map stage. */
+private[shuffle] class PausedLocalShuffleRegistration {
+ val shuffleId = new AtomicInteger(-1)
+ val entered = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ val returned = new CountDownLatch(1)
+}
+
+/** Waits for scheduler-confirmed job completion, including jobs cancelled during fallback. */
+private[shuffle] class CompletedMaterializationJobs extends SparkListener {
+ private val completed = mutable.Set.empty[Int]
+ private val failed = mutable.Set.empty[Int]
+ private val jobsByGroup = mutable.Map.empty[String, Int]
+
+ override def onJobStart(event: SparkListenerJobStart): Unit = synchronized {
+ Option(event.properties)
+ .flatMap(properties => Option(properties.getProperty("spark.jobGroup.id")))
+ .foreach(groupId => jobsByGroup(groupId) = event.jobId)
+ notifyAll()
+ }
+
+ override def onJobEnd(event: SparkListenerJobEnd): Unit = synchronized {
+ completed += event.jobId
+ if (event.jobResult.isInstanceOf[JobFailed]) failed += event.jobId
+ notifyAll()
+ }
+
+ def hasFailed(jobId: Int): Boolean = synchronized { failed.contains(jobId) }
+
+ def awaitJobInGroup(groupId: String): Int = synchronized {
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20)
+ while (!jobsByGroup.contains(groupId) && System.nanoTime() < deadline) {
+ wait(math.max(1L, TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime())))
+ }
+ require(jobsByGroup.contains(groupId), s"No materialization job started in group $groupId")
+ jobsByGroup(groupId)
+ }
+
+ def awaitCompletion(jobIds: Seq[Int]): Unit = synchronized {
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20)
+ while (!jobIds.forall(completed.contains) && System.nanoTime() < deadline) {
+ wait(math.max(1L, TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime())))
+ }
+ require(jobIds.forall(completed.contains), s"Materialization jobs still running: $jobIds")
+ }
+}
+
+class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean)
+ extends CometCelebornShuffleManager(
+ conf,
+ isDriver,
+ (configuration, _) => new CometCelebornFallbackTestBackend(configuration),
+ planningSupportFactory = _ => CelebornNativeShufflePlanningSupport()) {
+
+ private[shuffle] val remoteAttempts =
+ new ConcurrentLinkedQueue[(CometShuffleFallbackAttempt, CometCelebornFallbackTestClient)]()
+ private[shuffle] val localAttempts = new ConcurrentLinkedQueue[CometShuffleFallbackAttempt]()
+ val nativeReads = new AtomicInteger()
+ val remoteMapsBeforeFailure = new AtomicInteger()
+ @volatile var waitForFirstRemoteMap = false
+ @volatile private[shuffle] var delayedRemoteCompletion: Option[DelayedRemoteMapCompletion] =
+ None
+ @volatile private[shuffle] var pausedMapCompletion: Option[PausedShuffleMapCompletion] = None
+ @volatile private[shuffle] var remoteFailureAfterLocalMap: Option[PausedShuffleMapCompletion] =
+ None
+ @volatile private[shuffle] var pausedLocalRegistration: Option[PausedLocalShuffleRegistration] =
+ None
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ val registration = dependency match {
+ case native: CometShuffleDependency[_, _, _] if native.useLocalShuffle =>
+ pausedLocalRegistration
+ case _ => None
+ }
+ registration.foreach { paused =>
+ paused.shuffleId.set(shuffleId)
+ paused.entered.countDown()
+ require(
+ paused.release.await(20, TimeUnit.SECONDS),
+ "The test must release local shuffle registration")
+ }
+ try super.registerShuffle(shuffleId, dependency)
+ finally registration.foreach(_.returned.countDown())
+ }
+
+ override protected[shuffle] def shouldReportShuffleFetchFailure(taskAttemptId: Long): Boolean =
+ true
+
+ override protected[shuffle] def createRemotePusher(
+ handle: ShuffleHandle,
+ context: TaskContext,
+ onGenerationResolved: (Int, Int) => Unit,
+ onGenerationInvalidated: (Int, Int) => Unit,
+ onInvalidationUnsafe: (Int, Int) => Boolean): ResolvedCelebornShufflePusher = {
+ val remoteHandle = handle.asInstanceOf[CelebornShuffleHandle[_, _, _]]
+ val dependency = remoteHandle.dependency
+ val client = new CometCelebornFallbackTestClient
+ val generation = handle.shuffleId + 1000
+ val numMappers = remoteHandle.numMappers
+ onGenerationResolved(generation, numMappers)
+ remoteAttempts.add((recordAttempt(handle, context), client))
+ ResolvedCelebornShufflePusher(
+ CelebornShufflePusherFactory.create(
+ conf,
+ client,
+ generation,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context),
+ client,
+ generation)
+ }
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ if (waitForFirstRemoteMap && handle.isInstanceOf[CelebornShuffleHandle[_, _, _]] &&
+ context.stageAttemptNumber() == 0 && context.partitionId() == 1) {
+ // This fixture runs Spark locally, so the executor can observe the driver's tracker.
+ // Waiting for a recorded MapStatus proves the retry discards already published output.
+ val tracker = SparkEnv.get.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster]
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20)
+ while (tracker.getNumAvailableOutputs(handle.shuffleId) == 0 &&
+ System.nanoTime() < deadline) {
+ Thread.sleep(10)
+ }
+ val completed = tracker.getNumAvailableOutputs(handle.shuffleId)
+ require(
+ completed == 1,
+ s"Expected one completed remote map before size failure: $completed")
+ remoteMapsBeforeFailure.set(completed)
+ }
+ val delayed = delayedRemoteCompletion
+ if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]] && context.partitionId() == 1) {
+ delayed.foreach(_.awaitRemoteCommit())
+ }
+ val underlying = super.getWriter[K, V](handle, mapId, context, metrics)
+ val failureGate = remoteFailureAfterLocalMap.filter { _ =>
+ handle.isInstanceOf[CelebornShuffleHandle[_, _, _]]
+ }
+ val writer = if (failureGate.nonEmpty) {
+ new ShuffleWriter[K, V] {
+ override def write(records: Iterator[Product2[K, V]]): Unit = {
+ try underlying.write(records)
+ catch {
+ case failure if CometNativeShuffleWriter.isSizeLimitFailure(failure) =>
+ // The replacement must be active when Spark processes the abandoned stage's
+ // exception; otherwise the shared-RDD bug depends on scheduler timing.
+ require(failureGate.get.stopped.await(20, TimeUnit.SECONDS))
+ throw failure
+ }
+ }
+
+ override def getPartitionLengths(): Array[Long] = underlying.getPartitionLengths()
+
+ override def stop(success: Boolean): Option[MapStatus] = underlying.stop(success)
+ }
+ } else {
+ underlying
+ }
+ if (handle.isInstanceOf[CometNativeShuffleHandle[_, _]]) {
+ localAttempts.add(recordAttempt(handle, context))
+ delayed.foreach(_.localStageStarted.countDown())
+ }
+ val paused = pausedMapCompletion.filter { completion =>
+ completion.local == handle.isInstanceOf[CometNativeShuffleHandle[_, _]]
+ }
+ if (paused.nonEmpty) {
+ new ShuffleWriter[K, V] {
+ override def write(records: Iterator[Product2[K, V]]): Unit = writer.write(records)
+
+ override def getPartitionLengths(): Array[Long] = writer.getPartitionLengths()
+
+ override def stop(success: Boolean): Option[MapStatus] = {
+ val result = writer.stop(success)
+ if (success && result.nonEmpty) {
+ paused.get.stopped.countDown()
+ paused.get.awaitRelease()
+ paused.get.returned.countDown()
+ }
+ result
+ }
+ }
+ } else if (handle
+ .isInstanceOf[CelebornShuffleHandle[_, _, _]] && context.partitionId() == 0 &&
+ delayed.nonEmpty) {
+ new ShuffleWriter[K, V] {
+ override def write(records: Iterator[Product2[K, V]]): Unit = writer.write(records)
+
+ override def getPartitionLengths(): Array[Long] = writer.getPartitionLengths()
+
+ override def stop(success: Boolean): Option[MapStatus] = {
+ val result = writer.stop(success)
+ if (success && result.nonEmpty) {
+ delayed.get.remoteCommitted.countDown()
+ // Keep the already validated result until local writers have started. Cancellation
+ // may interrupt this worker, but must not let this old result replace local output.
+ delayed.get.awaitLocalStage()
+ delayed.get.remoteReturned.countDown()
+ }
+ result
+ }
+ }
+ } else {
+ writer
+ }
+ }
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ val reader = super.getReader[K, C](
+ handle,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ if (handle.isInstanceOf[CometNativeShuffleHandle[_, _]]) {
+ nativeReads.incrementAndGet()
+ }
+ reader
+ }
+
+ private def recordAttempt(
+ handle: ShuffleHandle,
+ context: TaskContext): CometShuffleFallbackAttempt =
+ CometShuffleFallbackAttempt(
+ handle.shuffleId,
+ context.stageId(),
+ context.stageAttemptNumber(),
+ context.partitionId())
+}
+
+class CometCelebornFallbackTestClient extends RecordingCelebornPushClient {
+ val fetchFailureReports = new AtomicInteger()
+ val shuffleCleanupCalls = new AtomicInteger()
+
+ def reportShuffleFetchFailure(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ taskAttemptId: Long): Boolean = {
+ fetchFailureReports.incrementAndGet()
+ true
+ }
+
+ def cleanupShuffle(shuffleId: Int): Unit = {
+ shuffleCleanupCalls.incrementAndGet()
+ }
+}
+
+private[shuffle] class CometCelebornFallbackTestBackend(conf: SparkConf) extends ShuffleManager {
+ private val ordinaryShuffle = new CometShuffleManager(conf)
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = dependency match {
+ case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle =>
+ new CelebornShuffleHandle(shuffleId, dependency)
+ case _ => ordinaryShuffle.registerShuffle(shuffleId, dependency)
+ }
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ require(
+ !handle.isInstanceOf[CelebornShuffleHandle[_, _, _]],
+ "Unexpected ordinary native writer")
+ ordinaryShuffle.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] = {
+ require(
+ !handle.isInstanceOf[CelebornShuffleHandle[_, _, _]],
+ "Unexpected remote native reader")
+ ordinaryShuffle.getReader(
+ handle,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ }
+
+ override def shuffleBlockResolver: ShuffleBlockResolver = ordinaryShuffle.shuffleBlockResolver
+
+ override def unregisterShuffle(shuffleId: Int): Boolean =
+ ordinaryShuffle.unregisterShuffle(shuffleId)
+
+ override def stop(): Unit = ordinaryShuffle.stop()
+}
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 538531d728..c22626a520 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
@@ -19,6 +19,8 @@
package org.apache.spark.sql.comet.execution.shuffle
+import java.util.concurrent.{CountDownLatch, Executors, TimeUnit}
+
import scala.collection.mutable
import org.scalatest.funsuite.AnyFunSuite
@@ -360,6 +362,157 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite {
assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 1, 0, 0)).authorized)
}
+ test("an accepted size fallback abandons the remote shuffle and invalidates every owner") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val requestedFallbacks = mutable.ArrayBuffer.empty[Int]
+ val coordinator = new CelebornShuffleGenerationCoordinator(
+ sparkCoordinator,
+ _ => true,
+ shuffleId => {
+ requestedFallbacks += shuffleId
+ true
+ })
+ val shuffleId = 13
+ val stageId = 77
+ startStage(sparkCoordinator, stageId, numMappers = 2)
+ val generation = PrepareCelebornShuffleGeneration(shuffleId, 99, stageId, 0, 2)
+ assert(coordinator.prepareGeneration(generation))
+ val validations = (0 until 2).map { mapId =>
+ val owner =
+ coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, mapId, 0))
+ ValidateCelebornMapAttempt(shuffleId, 99, stageId, 0, mapId, 0, owner.epoch)
+ }
+ assert(validations.forall(coordinator.validateMapAttempt))
+
+ assert(coordinator.requestLocalShuffle(RequestLocalCometShuffle(validations.head, 100L)))
+ assert(requestedFallbacks.toSeq == Seq(shuffleId))
+ assert(validations.forall(validation => !coordinator.validateMapAttempt(validation)))
+ assert(!coordinator.requestLocalShuffle(RequestLocalCometShuffle(validations.last, 101L)))
+ assert(
+ !coordinator.prepareGeneration(generation.copy(celebornShuffleId = 100, stageAttempt = 1)))
+ assert(
+ !coordinator
+ .claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, 0, 1))
+ .authorized)
+ assert(
+ !coordinator
+ .claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 1, 0, 0))
+ .authorized)
+ assert(requestedFallbacks.toSeq == Seq(shuffleId))
+
+ coordinator.unregisterShuffle(shuffleId)
+ assert(
+ coordinator.prepareGeneration(generation.copy(celebornShuffleId = 101, stageAttempt = 2)))
+ }
+
+ test("a declined size fallback preserves the remote generation and can be requested again") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ var acceptFallback = false
+ var requests = 0
+ val coordinator = new CelebornShuffleGenerationCoordinator(
+ sparkCoordinator,
+ _ => true,
+ _ => {
+ requests += 1
+ acceptFallback
+ })
+ val stageId = 80
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+ assert(
+ coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(16, 105, stageId, 0, 1)))
+ val remote = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(16, stageId, 0, 0, 0))
+ val validation = ValidateCelebornMapAttempt(16, 105, stageId, 0, 0, 0, remote.epoch)
+ val request = RequestLocalCometShuffle(validation, 105L)
+
+ assert(!coordinator.requestLocalShuffle(request))
+ assert(coordinator.validateMapAttempt(validation))
+ acceptFallback = true
+ assert(coordinator.requestLocalShuffle(request))
+ assert(!coordinator.validateMapAttempt(validation))
+ assert(requests == 2)
+ }
+
+ test("concurrent size reports invoke the active fallback once") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val requestedFallbacks = mutable.ArrayBuffer.empty[Int]
+ val coordinator = new CelebornShuffleGenerationCoordinator(
+ sparkCoordinator,
+ _ => true,
+ shuffleId => {
+ requestedFallbacks += shuffleId
+ true
+ })
+ val stageId = 78
+ startStage(sparkCoordinator, stageId, numMappers = 2)
+ assert(
+ coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(14, 102, stageId, 0, 2)))
+ val requests = (0 until 2).map { mapId =>
+ val owner = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(14, stageId, 0, mapId, 0))
+ RequestLocalCometShuffle(
+ ValidateCelebornMapAttempt(14, 102, stageId, 0, mapId, 0, owner.epoch),
+ 102L + mapId)
+ }
+ val executor = Executors.newFixedThreadPool(2)
+ val ready = new CountDownLatch(2)
+ val start = new CountDownLatch(1)
+ try {
+ val reports = requests.map { request =>
+ executor.submit(new java.util.concurrent.Callable[Boolean] {
+ override def call(): Boolean = {
+ ready.countDown()
+ assert(start.await(5, TimeUnit.SECONDS))
+ coordinator.requestLocalShuffle(request)
+ }
+ })
+ }
+ assert(ready.await(5, TimeUnit.SECONDS))
+ start.countDown()
+ assert(reports.map(_.get(5, TimeUnit.SECONDS)).count(identity) == 1)
+ assert(requestedFallbacks.toSeq == Seq(14))
+ } finally {
+ start.countDown()
+ executor.shutdownNow()
+ }
+ }
+
+ test("stale or unsafe size reports never invoke the active fallback") {
+ Seq(false, true).foreach { stale =>
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(
+ sparkCoordinator,
+ _ => stale,
+ _ => fail("A stale or unsafe map attempt must not request fallback"))
+ val stageId = 79
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+ val generation = PrepareCelebornShuffleGeneration(15, 103, stageId, 0, 1)
+ assert(coordinator.prepareGeneration(generation))
+ val owner = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(15, stageId, 0, 0, 0))
+ val validation = ValidateCelebornMapAttempt(15, 103, stageId, 0, 0, 0, owner.epoch)
+ if (stale) {
+ assert(
+ coordinator.prepareGeneration(
+ generation.copy(celebornShuffleId = 104, stageAttempt = 1)))
+ }
+
+ assert(!coordinator.requestLocalShuffle(RequestLocalCometShuffle(validation, 104L)))
+ if (!stale) assert(coordinator.validateMapAttempt(validation))
+ }
+ }
+
+ test("a size report without an active materialization preserves the remote generation") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator, _ => true)
+ val stageId = 81
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+ assert(
+ coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(17, 106, stageId, 0, 1)))
+ val remote = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(17, stageId, 0, 0, 0))
+ val validation = ValidateCelebornMapAttempt(17, 106, stageId, 0, 0, 0, remote.epoch)
+
+ assert(!coordinator.requestLocalShuffle(RequestLocalCometShuffle(validation, 106L)))
+ assert(coordinator.validateMapAttempt(validation))
+ }
+
test("manager preserves the application Spark configuration and driver identity") {
val conf = new SparkConf(false).set("spark.app.name", "existing-celeborn-application")
val backend = new RecordingShuffleManager
@@ -414,6 +567,24 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite {
assert(composite.registerShuffle[Any, Any, Any](31, null) eq backend.returnedHandle)
}
+ test("native planning uses client compatibility without a scheduler-version gate") {
+ var configurationLoads = 0
+ var completionProbes = 0
+ val support = CometCelebornShuffleManager.nativeShufflePlanningSupport(
+ new SparkConf(false),
+ _ => {
+ configurationLoads += 1
+ new ReflectedPlanningConf
+ },
+ () => {
+ completionProbes += 1
+ None
+ })
+ assert(support.fallbackReason(1).isEmpty)
+ assert(configurationLoads == 1)
+ assert(completionProbes == 1)
+ }
+
test("encrypted native planning falls back before loading the optional Celeborn API") {
val conf = new SparkConf(false).set("spark.io.encryption.enabled", "true")
var loads = 0
@@ -429,6 +600,75 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite {
assert(loads == 0)
}
+ test("native planning requires executor retention for local size-limit fallback") {
+ val cases = Seq(
+ ("fixed executors", false, false, false, true),
+ ("shuffle tracking", true, false, true, true),
+ ("external shuffle service", true, true, false, true),
+ ("both retention mechanisms", true, true, true, true),
+ ("remote storage alone", true, false, false, false))
+ cases.foreach { case (name, dynamicAllocation, shuffleService, tracking, supported) =>
+ val conf = new SparkConf(false)
+ .set("spark.dynamicAllocation.enabled", dynamicAllocation.toString)
+ .set("spark.shuffle.service.enabled", shuffleService.toString)
+ .set("spark.dynamicAllocation.shuffleTracking.enabled", tracking.toString)
+ var probes = 0
+ val support = CometCelebornShuffleManager.nativeShufflePlanningSupport(
+ conf,
+ _ => new ReflectedPlanningConf,
+ () => {
+ probes += 1
+ None
+ })
+ assert(support.fallbackReason(1).isEmpty == supported, name)
+ assert(probes == (if (supported) 1 else 0), name)
+ }
+
+ // Read Spark's ConfigEntry so omitting the setting agrees with the executor allocation
+ // manager's effective behavior on each supported Spark version.
+ val defaultTracking = CometCelebornShuffleManager.nativeShufflePlanningSupport(
+ new SparkConf(false).set("spark.dynamicAllocation.enabled", "true"),
+ _ => new ReflectedPlanningConf,
+ () => None)
+ assert(
+ defaultTracking.fallbackReason(1).isEmpty ==
+ new SparkConf(false).get(
+ org.apache.spark.internal.config.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED))
+ }
+
+ test("unprotected local fallback retains ordinary delegated shuffle") {
+ Seq(false, true).foreach { decommission =>
+ val conf = new SparkConf(false)
+ .set("spark.dynamicAllocation.enabled", "true")
+ .set("spark.shuffle.service.enabled", "false")
+ .set("spark.dynamicAllocation.shuffleTracking.enabled", "false")
+ .set("spark.decommission.enabled", decommission.toString)
+ .set("spark.storage.decommission.enabled", decommission.toString)
+ .set("spark.storage.decommission.shuffleBlocks.enabled", decommission.toString)
+ val backend = new RecordingShuffleManager
+ val composite = new CometCelebornShuffleManager(
+ conf,
+ false,
+ (_, _) => backend,
+ planningSupportFactory = actualConf =>
+ CometCelebornShuffleManager.nativeShufflePlanningSupport(
+ actualConf,
+ _ => fail("Unprotected local fallback must not load native Celeborn configuration"),
+ () => fail("Unprotected local fallback must not probe native push completion")))
+ try {
+ assert(composite.nativeShuffleFallbackReason(1).exists(_.contains("local fallback")))
+ val handle = composite.registerShuffle[Any, Any, Any](31, null)
+ assert(composite.getWriter[Any, Any](handle, 12L, null, null) == null)
+ assert(composite.getReader[Any, Any](handle, 0, 2, null, null) == null)
+ assert(backend.writerCall.contains((handle, 12L)))
+ assert(backend.rangedReaderCall.contains((handle, 0, Int.MaxValue, 0, 2)))
+ } finally {
+ composite.stop()
+ }
+ assert(backend.stopped)
+ }
+ }
+
test("native planning uses Celeborn's effective stage-rerun setting") {
// Celeborn resolves JVM properties and legacy aliases itself. Spark SQL settings must not
// override the resolved value used by the actual client and lifecycle manager.
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 db89c1840f..c1e186e0ea 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
@@ -24,7 +24,7 @@ import java.lang.reflect.{InvocationHandler, Method, Proxy}
import java.nio.{ByteBuffer, ByteOrder}
import java.nio.channels.Channels
import java.nio.charset.StandardCharsets
-import java.nio.file.Files
+import java.nio.file.{Files, Path}
import java.util.{LinkedHashSet, Set => JSet}
import java.util.concurrent.{CountDownLatch, TimeoutException, TimeUnit}
import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
@@ -81,11 +81,15 @@ class CometCelebornShuffleReaderSuite extends CometTestBase {
extends ShuffleManager {
var readRange: Option[(Int, Int, Int, Int)] = None
var unregistered: Option[Int] = None
+ var fileToRemove: Option[Path] = None
+ var stopped = false
override def registerShuffle[K, V, C](
shuffleId: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle =
- throw new UnsupportedOperationException("This fixture only reads existing native shuffles")
+ new CometNativeShuffleHandle[K, V](
+ shuffleId,
+ dependency.asInstanceOf[ShuffleDependency[K, V, V]])
override def getWriter[K, V](
handle: ShuffleHandle,
@@ -110,10 +114,11 @@ class CometCelebornShuffleReaderSuite extends CometTestBase {
override def unregisterShuffle(shuffleId: Int): Boolean = {
unregistered = Some(shuffleId)
+ fileToRemove.foreach(Files.deleteIfExists(_))
true
}
- override def stop(): Unit = ()
+ override def stop(): Unit = stopped = true
}
private final class RecordingReaderApi extends CelebornRawPartitionReader.Api {
@@ -169,13 +174,17 @@ class CometCelebornShuffleReaderSuite extends CometTestBase {
context.taskMetrics().createTempShuffleReadMetrics(),
retries)
- private def simpleDependency(partitions: Int = 3, attributes: Seq[Attribute] = Seq.empty)
+ private def simpleDependency(
+ partitions: Int = 3,
+ attributes: Seq[Attribute] = Seq.empty,
+ useLocalShuffle: Boolean = false)
: 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"),
- outputAttributes = attributes)
+ outputAttributes = attributes,
+ useLocalShuffle = useLocalShuffle)
private def ownedClients(manager: CometCelebornShuffleManager)
: java.util.concurrent.ConcurrentHashMap[AnyRef, java.lang.Boolean] = {
@@ -298,6 +307,47 @@ class CometCelebornShuffleReaderSuite extends CometTestBase {
context.markTaskCompleted(None)
}
+ test("local fallback shutdown preserves files until explicit shuffle unregister") {
+ Seq(false, true).foreach { unregister =>
+ val context = TaskContext.empty()
+ val dependency = simpleDependency(useLocalShuffle = true)
+ val backend = new RecordingReaderBackend(new RecordingCelebornRawClient)
+ val local = new RecordingReaderBackend(new RecordingCelebornRawClient)
+ val file = Files.createTempFile("comet-local-fallback", ".data")
+ local.fileToRemove = Some(file)
+ val manager = new CometCelebornShuffleManager(
+ new SparkConf(false).set("spark.shuffle.service.enabled", "true"),
+ false,
+ (_, _) => backend,
+ localManagerFactory = _ => local)
+ val handle = manager.registerShuffle(17, dependency)
+ assert(handle.isInstanceOf[CometNativeShuffleHandle[_, _]])
+ try {
+ manager.getReader[Int, ColumnarBatch](
+ handle,
+ 0,
+ 1,
+ context,
+ context.taskMetrics.createTempShuffleReadMetrics())
+ assert(local.readRange.contains((0, Int.MaxValue, 0, 1)))
+ assert(backend.readRange.isEmpty)
+ if (unregister) {
+ assert(manager.unregisterShuffle(17))
+ assert(local.unregistered.contains(17))
+ assert(backend.unregistered.isEmpty)
+ }
+ manager.stop()
+ assert(local.stopped)
+ assert(backend.stopped)
+ assert(Files.exists(file) == !unregister)
+ if (!unregister) assert(local.unregistered.isEmpty)
+ } finally {
+ Files.deleteIfExists(file)
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+
test("the manager routes native Celeborn handles through its reflected raw reader") {
val context = TaskContext.empty()
val dependency = simpleDependency()
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleStatisticsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleStatisticsSuite.scala
new file mode 100644
index 0000000000..f4be268036
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleStatisticsSuite.scala
@@ -0,0 +1,132 @@
+/*
+ * 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 scala.concurrent.Await
+import scala.concurrent.duration.DurationInt
+
+import org.apache.spark.{SparkConf, SparkEnv}
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter.{SHUFFLE_BYTES_WRITTEN, SHUFFLE_RECORDS_WRITTEN}
+
+import org.apache.comet.CometConf
+
+class CometCelebornShuffleStatisticsSuite extends CometTestBase {
+ import testImplicits._
+
+ override protected val shuffleManager: String =
+ classOf[CometCelebornFallbackTestShuffleManager].getName
+
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(CometConf.COMET_SHUFFLE_MODE.key, "native")
+ .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false")
+ .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k")
+ .set(CometConf.COMET_EXCHANGE_SIZE_MULTIPLIER.key, "1")
+ .set("spark.shuffle.compress", "false")
+ .set("spark.celeborn.client.spark.stageRerun.enabled", "true")
+
+ private def manager: CometCelebornFallbackTestShuffleManager =
+ SparkEnv.get.shuffleManager.asInstanceOf[CometCelebornFallbackTestShuffleManager]
+
+ for (adaptive <- Seq(false, true)) {
+ test(s"fallback publishes only the local destination's output statistics: AQE=$adaptive") {
+ withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) {
+ manager.remoteMapsBeforeFailure.set(0)
+ manager.waitForFirstRemoteMap = true
+ try {
+ val shuffled = spark
+ .range(0, 2, 1, 2)
+ .selectExpr(
+ "CAST(id AS INT) AS key",
+ "CASE WHEN id = 0 THEN 'small' ELSE repeat('x', 1048576) END AS payload")
+ .repartition(3, $"key")
+ val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq
+ val exchange = collect(shuffled.queryExecution.executedPlan) {
+ case value: CometShuffleExchangeExec => value
+ }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange"))
+ val remote = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ val local = remote.currentShuffleDependency
+
+ assert(actual.sortBy(_._1) == Seq((0, "small"), (1, "x" * (1024 * 1024))))
+ assert(manager.remoteMapsBeforeFailure.get() == 1)
+ assert(local.useLocalShuffle)
+ assert(remote.shuffleWriteMetrics(SHUFFLE_RECORDS_WRITTEN).value == 1L)
+ assert(exchange.runtimeStatistics.rowCount.contains(BigInt(2)))
+ assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value == 2L)
+ assert(remote.shuffleWriteMetrics("dataSize").value > 0L)
+ assert(
+ exchange.runtimeStatistics.sizeInBytes ==
+ BigInt(local.shuffleWriteMetrics("dataSize").value))
+ assert(
+ exchange.metrics("dataSize").value ==
+ local.shuffleWriteMetrics("dataSize").value)
+ assert(
+ exchange.metrics(SHUFFLE_BYTES_WRITTEN).value ==
+ local.shuffleWriteMetrics(SHUFFLE_BYTES_WRITTEN).value)
+
+ // Simulate another accumulator merge from the abandoned remote stage after publication.
+ // Resetting the original counters would allow that update to inflate AQE statistics.
+ val published = exchange.runtimeStatistics
+ val publishedBytes = exchange.metrics(SHUFFLE_BYTES_WRITTEN).value
+ remote.shuffleWriteMetrics("dataSize").add(4096L)
+ remote.shuffleWriteMetrics(SHUFFLE_RECORDS_WRITTEN).add(1L)
+ remote.shuffleWriteMetrics(SHUFFLE_BYTES_WRITTEN).add(4096L)
+ assert(exchange.runtimeStatistics == published)
+ assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value == 2L)
+ assert(exchange.metrics("dataSize").value == published.sizeInBytes.toLong)
+ assert(exchange.metrics(SHUFFLE_BYTES_WRITTEN).value == publishedBytes)
+
+ assert(manager.unregisterShuffle(remote.shuffleId))
+ assert(manager.unregisterShuffle(local.shuffleId))
+ } finally {
+ manager.waitForFirstRemoteMap = false
+ }
+ }
+ }
+ }
+
+ test("successful remote materialization publishes remote output statistics") {
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ val query = spark.range(0, 2, 1, 2).repartition(2, $"id")
+ val exchange = collect(query.queryExecution.executedPlan) {
+ case value: CometShuffleExchangeExec => value
+ }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange"))
+ val remote = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]]
+ try {
+ val statistics = Await.result(exchange.mapOutputStatisticsFuture, 20.seconds)
+ assert(remote.currentShuffleDependency eq remote)
+ assert(exchange.runtimeStatistics.rowCount.contains(BigInt(2)))
+ assert(exchange.runtimeStatistics.sizeInBytes > 0)
+ assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value == 2L)
+ assert(statistics.bytesByPartitionId.sum > 0L)
+ // MapOutputStatistics uses Spark's compressed MapStatus sizes, which are approximate.
+ assert(
+ exchange.metrics(SHUFFLE_BYTES_WRITTEN).value ==
+ remote.shuffleWriteMetrics(SHUFFLE_BYTES_WRITTEN).value)
+ assert(
+ exchange.metrics("dataSize").value ==
+ exchange.runtimeStatistics.sizeInBytes.toLong)
+ } finally {
+ manager.unregisterShuffle(remote.shuffleId)
+ }
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
index c3a19c8790..99d9fe93e0 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
@@ -20,7 +20,7 @@
package org.apache.spark.sql.comet.execution.shuffle
import org.apache.spark.{HashPartitioner, Partition, TaskContext}
-import org.apache.spark.rdd.RDD
+import org.apache.spark.rdd.{DeterministicLevel, RDD}
import org.apache.spark.serializer.JavaSerializer
import org.apache.spark.sql.CometTestBase
import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext}
@@ -45,6 +45,73 @@ import org.apache.comet.serde.OperatorOuterClass.Operator
*/
class CometNativeShuffleInputRDDSuite extends CometTestBase {
+ test("native shuffle input preserves its parents' determinism") {
+ Seq(
+ DeterministicLevel.DETERMINATE,
+ DeterministicLevel.UNORDERED,
+ DeterministicLevel.INDETERMINATE).foreach { level =>
+ val parent = new RDD[AnyRef](spark.sparkContext, Nil) {
+ override protected def getOutputDeterministicLevel: DeterministicLevel.Value = level
+ override protected def getPartitions: Array[Partition] = Array.empty
+ override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] =
+ Iterator.empty
+ }
+ val input = new CometNativeShuffleInputRDD(
+ spark.sparkContext,
+ Seq(parent),
+ 0,
+ Set.empty,
+ CometMetricNode(Map.empty))
+ assert(input.outputDeterministicLevel == level)
+ assert(input.copyForLocalShuffle().outputDeterministicLevel == level)
+ }
+ }
+
+ test("local shuffle input is an independent sibling with the same partition inputs") {
+ val upstream = new RDD[AnyRef](spark.sparkContext, Nil) {
+ override protected def getPartitions: Array[Partition] = Array.tabulate(2) { i =>
+ new Partition {
+ override def index: Int = i
+ }
+ }
+
+ override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] =
+ Iterator.single(null)
+
+ override def getPreferredLocations(split: Partition): Seq[String] =
+ Seq(s"host-${split.index}")
+ }
+ val planData = Map("scan-0" -> Array(Array[Byte](1), Array[Byte](2)))
+ val remote = new CometNativeShuffleInputRDD(
+ spark.sparkContext,
+ Seq(upstream),
+ 2,
+ Set.empty,
+ CometMetricNode(Map.empty),
+ planData)
+ val local = remote.copyForLocalShuffle()
+
+ assert(local.id != remote.id)
+ assert(local.dependencies.map(_.rdd) == Seq(upstream))
+ assert(local.getNumPartitions == remote.getNumPartitions)
+ local.partitions.foreach { part =>
+ val input = part.asInstanceOf[CometNativeShuffleInputPartition]
+ assert(input.inputPartitions.toSeq == Seq(upstream.partitions(part.index)))
+ assert(input.planDataByKey("scan-0").sameElements(planData("scan-0")(part.index)))
+ assert(local.preferredLocations(part) == Seq(s"host-${part.index}"))
+
+ val context = TaskContext.empty()
+ try {
+ val iterator = local.iterator(part, context).asInstanceOf[CometNativeShuffleInputIterator]
+ assert(iterator.partitionIndex == part.index)
+ assert(iterator.inputObjects.toSeq == Seq(null))
+ assert(iterator.planDataByKey("scan-0").sameElements(planData("scan-0")(part.index)))
+ } finally {
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+
test("spill reporting is registered before native shuffle input producers") {
Seq(None, Some(new IllegalStateException("failed native shuffle"))).foreach { failure =>
val writerDisk = new SQLMetric("writerDisk")
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org