fix: recover native Celeborn shuffle from oversized rows

[pingzh]

Which issue does this PR close?

Closes #5527.

Rationale for this change

A large row can exhaust native Celeborn's byte-admission budget even below the configured frame limit, then fail identically on every retry. Align the defaults and let oversized rows complete through local Comet shuffle while preserving ordinary Spark fetch-failure recovery.

What changes are included in this PR?

• Raise default executor admission from 256 MiB to 512 MiB to accommodate the 64 MiB frame limit and encoding workspace. Preserve typed size-limit errors through Rust, DataFusion, and JNI, with diagnostics identifying the limits and required reservation.
• Choose storage before exposing shuffle dependencies to readers. A size-limit failure recomputes every map locally with independent input RDD, shuffle, and stage identities. Remote failures and delayed results cannot abort or corrupt the replacement.
• Start independent materializations concurrently, including non-AQE and nested shuffle branches. Compose upstream futures without blocking preparation workers or starving completion callbacks. Wait before publishing reader dependencies, preserve session artifacts and job properties, and keep cancellation effective during upstream waits and fallback.
• Give each destination separate output counters and publish only the selected destination's statistics. Completed or late remote maps cannot inflate AQE row counts or data sizes after fallback.
• Preserve upstream determinism and normal fetch-failure recovery after partial result completion. Require external shuffle service or shuffle tracking with dynamic allocation, preserve fallback files at executor shutdown, and clean them on explicit shuffle removal.

How are these changes tested?

• 301 tests passed across 13 JVM suites on Spark 3.5.9 and another 301 passed on Spark 4.1.3. The Spark 3.5 run rebuilt from clean targets; both runs include the new starvation regression and the existing recovery, cancellation, concurrency, statistics, and native shuffle coverage.
• A bounded-pool regression runs six dependents of a shared upstream shuffle on two worker threads. It checks worker availability while upstream work is blocked, successful fallback for every shuffle, and exact query results. Restoring the blocking implementation makes this regression time out.
• Regression coverage forces a remote stage failure while the local replacement is active, proves sibling map tasks overlap, exercises nested fallback and cancellation, and verifies exact output statistics after a completed remote map and simulated late accumulator updates. The forced recovery regression reproduces the abort when the shared input RDD is restored.
• Existing coverage includes delayed remote completions, cancellation during replacement registration, and a local fetch failure after a nonempty result partition completes. Tests use a fake Celeborn client with real Spark scheduling, native encoding, and local shuffle I/O.
• Previously validated native changes: 45 RSS tests and 2 JNI tests; 38/50/63 MiB rows under the new default budget; a 70 MiB local shuffle roundtrip; native build, Clippy, and Rust formatting.
• Spotless, Scalastyle, and Linux/macOS suite-registration checks passed.

[sunchao]

Reviewed head 6f2465a2 against base 4219fc79 across five independent scopes. I found one recovery bug and two performance issues worth addressing.

1. [P1] A remote failure can abort the local replacement.
   The replacement dependency reuses the remote stage’s exact input RDD. Spark propagates stage failures using RDD identity, so separate shuffle and stage IDs do not provide sufficient isolation.

   I reproduced this in the PR’s existing tests: the remote stage’s size-limit failure killed the replacement stage. A separate Spark probe confirmed that a fresh sibling input RDD avoids this failure. Create an independent scheduling RDD over the original upstream inputs, or fully retire the remote job before starting its replacement. A narrow wrapper around the old RDD is insufficient.

2. [P2] Non-AQE execution serializes independent shuffle branches.
   The new RDD constructor call waits for materialization to finish. Joins and unions construct their inputs sequentially, so one sibling shuffle now finishes before another starts, even when neither needs fallback.

   In a Spark scheduling probe, two one-second map tasks overlapped for approximately one second with ordinary scheduling and had zero overlap with eager materialization. Submit independent materializations before awaiting their selected dependencies, while preserving the fence against downstream reads.

3. [P2] Fallback double-counts runtime statistics used by AQE.
   The replacement reuses the original metrics. Successful remote maps have already contributed to them; recomputing every map locally adds those contributions again.

   An exact-head Spark/native probe returned 2 rows but reported 3 runtime rows, after one successful remote map and two local maps. Inflated size statistics can also affect AQE join selection. Give each destination separate output statistics and publish the selected destination’s values. Resetting shared counters remains vulnerable to late remote completions.

The core design is reasonable: typed errors and a driver-owned, immutable storage decision are useful abstractions. The materialization wrapper has real cancellation and lifecycle responsibilities. The main improvements should make attempt ownership explicit and move blocking orchestration out of the RDD constructor. I found no meaningful new native encoding/copy overhead. The 512 MiB admission default raises potential memory usage without allocating that amount eagerly.

Validation: native build, 34 RSS tests, and 2 JNI tests passed. The focused Spark 4.1 run passed 9/11 tests, with both failures reproducing the recovery race. CI also has failing fallback tests on Spark 3.5, 4.0, and 4.2. JVM tests used the PR’s fake Celeborn client with real Spark/native execution. Full cluster behavior and whole-query performance were not benchmarked.

[sunchao]

Follow-up on f54d00bd against 4219fc79: the sibling scheduling RDD, asynchronous materialization with deferred dependency selection, and per-destination output counters address the three prior findings at source level. I found no remaining P1/P2 in this increment.

This review used the maintained Spark 3.5/4.0 contracts. I inspected the new regression tests but did not rerun them locally. At 07:51 UTC, CI had 21 successful checks, 30 running, 7 skipped and no failures. The new fallback, concurrency and statistics regressions had not yet been confirmed by completed CI, so this is not a claim of completed runtime validation.

[sunchao]

Re-reviewed f54d00bd against 4219fc79. The previous three findings are fixed. I found one new P2 issue.

[P2] Completion callbacks can starve behind blocked dependency preparation. In CometCelebornShuffleMaterialization.scala:115–125, dependency preparation blocks on upstream materialization using ExecutionContext.global, while upstream completion callbacks require that same pool. Enough waiting dependencies exhaust its capacity: Spark finishes the upstream job, but Comet cannot publish completion, leaving downstream work stuck.

I independently reproduced this using current Comet classes, native upstream execution, two visible driver CPUs, and Scala’s default pool settings:

Pending downstream dependencies	Global threads	Upstream materialization	Downstream jobs admitted
240	241	Completed	240
270	257	Stuck after Spark job succeeded	0

Prefer asynchronous composition of upstream completion, or ensure completion callbacks have execution capacity independent of blocked preparation workers. Increasing the thread limit only moves the threshold. Reproduction and logs (local files).

For the requested dimensions:

• Performance: The issue above causes substantial driver-thread growth and eventually stalls progress. I found no other material overhead in the native success path.
• Design: The sibling RDD correctly isolates replacement scheduling, and deferred waiting restores independent branch startup. The shared executor still needs attention.
• Abstraction and complexity: Per-destination output metrics and the materialization state machine have clear responsibilities. I found no unnecessary abstraction worth removing.

Validation: 24/24 focused Spark tests passed locally. Current CI has 65 successful checks and 9 skipped; its Spark 4.1 shuffle suite passed 463/463 tests. Local tests reused the checksum-verified native artifact because native sources were unchanged. The starvation reproduction uses local Spark and a fake Celeborn client; it does not establish distributed behavior or whole-query throughput.

[sunchao]

Correctness

Re-reviewed f54d00bd against 4219fc79 after the new discussion, comparing the scheduler, future and dependency-selection contracts with Spark 3.5 and 4.0. The three earlier findings remain addressed at source level. The mechanism in the existing [P2] callback-starvation finding is confirmed. I found no additional P1/P2, but the remaining P2 prevents a new approval.

Validation and CI

A component harness compiled the unmodified exact-head materialization class with real Scala 2.12.18, 2.13.16 and 2.13.17 libraries on JDK 17. All six expected-outcome scenarios passed. The harness substitutes Spark's RDD/context/action interfaces, so this confirms the callback mechanism, not a new Spark/native integration or distributed reproduction. I did not rerun the previously reported 24 local Spark tests or independently recount the 463 CI tests. The complete current-head refresh at 16:31 UTC showed 65 successful checks and nine skipped.

Performance

With two visible CPUs and default Scala pool settings, 240 waiting dependencies completed with 241 global threads. At 270 dependencies, all 257 global workers were waiting, the underlying upstream action had completed, and no downstream job was admitted. The same result occurred with each tested Scala version.

This confirms the thread growth and loss of progress described in the existing P2. It is not a whole-query throughput measurement. I found no additional material performance concern in this follow-up.

Design

The sibling scheduling RDD and deferred dependency selection preserve failure isolation and independent startup below pool saturation. Upstream completion still needs execution capacity independent of the workers waiting for it, as discussed in the existing comment. Merely increasing a finite thread limit moves the failure threshold.

Abstraction & complexity

The destination-specific counters separate published output statistics from late remote updates, and the materialization state checks give cancellation and late completion clear boundaries. Those abstractions have concrete responsibilities. I found no additional abstraction change to request beyond resolving the already-linked executor dependency.

[sunchao]

Correctness

Re-reviewed f8d1aa08 against 4219fc79, focusing on the scheduling change since f54d00bd. The existing [P2] callback-starvation finding is addressed. Upstream materializations are composed as futures before dependency preparation and job admission. I found no new or remaining P1/P2 in this increment.

The changed path agrees with the maintained Spark 3.5 and 4.0 RDD, scheduler and future contracts. Caching fixed partition metadata avoids the RDD-lock wait during union dependency discovery. Cancellation still fences late admission, upstream failures propagate, and discovery and submission restore captured job properties and session context. The earlier sibling-RDD isolation, independent startup and destination-specific statistics fixes remain intact.

Validation and CI

All 18 expected-outcome component scenarios passed using the unchanged old/current materialization classes and real Scala 2.12.18, 2.13.16 and 2.13.17 futures on JDK 17. These include bounded-pool progress, upstream failure, cancellation, fallback and property restoration. Spark RDD/context/action interfaces were substituted, so these are not new Spark/native integration results. I inspected the new regression but did not independently verify the reported 301 tests on each Spark version.

At 17:45 UTC, current-head CI had 22 successful checks, seven skipped, four running and 39 queued, with no failures. The integration matrix remains incomplete.

Performance

With two visible CPUs and default Scala pool settings, the prior class stalled at 270 dependents with 257 global workers and no downstream admissions. The current class completed all 270 dependents using two global workers on every tested Scala version. Six dependents also progressed on a fixed two-worker executor, which remained available while upstream completion was withheld.

This confirms removal of the reported worker growth and starvation, not whole-query throughput. The extra discovery pass and future collection operate per materialization. Native encoding and copying are unchanged, and I found no additional material overhead requiring a change.

Design

Composing upstream completion keeps each materialization responsible for its own ancestors and preserves the fence before Spark sees the selected dependency graph. The final preparation and atomic admission checks retain the existing cancellation and fallback boundaries. The fixed-pool test exercises the dependency between waiting work and completion capacity directly, without relying on a larger production thread limit.

Abstraction & complexity

The execution-context accessor is a narrow test seam with the same production default. The iterative discovery helper deduplicates RDDs and upstream futures, while the existing state machine continues to own storage selection and cancellation. These responsibilities are clear, and I found no unnecessary abstraction to request removing in this follow-up.