fix: recover native Celeborn shuffle from oversized rows

.github/workflows/pr_build_linux.yml

@@ -374,6 +374,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

.github/workflows/pr_build_macos.yml

@@ -147,6 +147,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

common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java

@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet;
+
+/** 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);
+  }
+}

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-->
 

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.
 

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() {