feat: enable native-only Celeborn shuffle planning (8/n)

[pingzh]

Which issue does this PR close?

Part of #5352. This is the eighth foundational PR and does not close the issue.

Previous PRs:

• feat: add RSS partition writer and task-owned JNI callback (1/n) #5473
• feat: add partition writer destinations to native shuffle plans (2/n) #5476
• feat: add destination-aware native shuffle execution (3/n) #5481
• feat: bind task-owned RSS callbacks to native shuffle plans (4/n) #5491
• feat: add Celeborn shuffle manager and partition pusher (5/n) #5501
• feat: complete Celeborn map-side shuffle push lifecycle (6/n) #5513
• feat: add raw Celeborn native shuffle reader (7/n) #5531

Rationale for this change

The preceding PRs added the Celeborn-backed native shuffle writer, map-side lifecycle, and raw reader. This PR enables the planner to select that path through CometCelebornShuffleManager with explicit native opt-in.

Unsupported exchanges must retain ordinary Spark/Celeborn shuffle without selecting Comet's JVM columnar shuffle. Fallback must also keep aggregate producers and consumers on compatible intermediate buffer formats.

What changes are included in this PR?

• Recognize the composite manager and require Comet execution, Comet shuffle, and explicit spark.comet.shuffle.mode=native. Default auto and jvm retain ordinary Spark/Celeborn shuffle.
• Snapshot effective application configuration through Celeborn's configuration API, preserving defaults, aliases, and fallback-policy precedence. Reject native planning for I/O encryption, disabled stage reruns, unsupported configuration APIs, and applicable local-fallback policies or partition thresholds. Session overrides cannot bypass these application-level restrictions.
• Require a native child and supported native partitioning/schema; preserve sticky AQE fallback and the existing local Comet shuffle behavior. Apply the shared guards to collect-limit and top-K shuffle producers as well.
• Keep incompatible aggregate buffers on Spark when an exchange or intermediate merge cannot run natively. Restore affected native ancestors, preserve completed/stage boundaries, and retain safe mixed execution such as non-decimal AVG.
• Make the shuffle-mode setting public, document configuration and limitations, and register the new planner suite in Linux and macOS CI.

How are these changes tested?

New planner and manager tests cover the opt-in matrix, runtime versus session configuration, supported and unsupported exchanges, AQE, special shuffle producers, fallback-policy boundaries, optional API failures, and immutable configuration snapshots.

Executed Spark-reference comparisons cover fallback queries, collect-list/set buffer compatibility, and mixed DISTINCT/percentile/FIRST/LAST pipelines. Planner-only checks cover repeated rule application, native-ancestor rollback, and completed aggregate boundaries. The local Comet shuffle path has a matching aggregate regression.

Local validation:

• Final full-reactor Maven run on Spark 4.1.3 / Scala 2.13 / JDK 21: 116 tests passed, including the corrected FIRST/LAST diagnostic assertion.
• Broader 13-suite shuffle, JNI, configuration, and aggregate run: 379/380 passed before the final diagnostic correction; its sole failure passed in the final rerun.
• Rust core, shuffle, and JNI library tests: 312 passed, with four existing HDFS tests ignored.
• Spotless, ScalaStyle, cargo fmt --all --check, and workspace Clippy with warnings denied passed.
• 63 configuration compatibility checks passed against published Celeborn 0.6.3 and 0.7.0 clients. Scala 2.12 parsing and the corresponding reflected API signatures were also checked.

Native plan-selection tests use a delegated test manager without executing remote native dependencies. No live Celeborn cluster or complete legacy Spark build matrix was run locally.

[sunchao]

Reviewed head a6e665563d2e9938f77ffa79258c6e8da2275932 against base/merge-base 98cd8c967995906500b12a709c25724b4cfba634. No introduced or materially worsened P1/P2 found.

Validation included 904 focused JVM tests across Spark 3.4/3.5/4.0/4.1/4.2, 312 native tests (four existing HDFS tests ignored), exact-base controls, and consumed native/JNI query, configuration, corruption, and cleanup checks with published Celeborn 0.6.3/0.7.0 clients.

Stock-client execution used Spark 3.5 with controlled loopback workers and lifecycle metadata; no deployed-cluster or performance validation. CI is still in progress.

[sunchao]

cc @andygrove @wForget

[ziting-openai]

Reviewed head a6e665563d2e9938f77ffa79258c6e8da2275932. No P1 or significant P2 findings after tracing native eligibility, application-versus-session configuration, ordinary shuffle fallback, aggregate buffer compatibility, AQE boundaries, and limit/top-K producers.

Static source and test review only; I did not run tests or a live Celeborn cluster.

Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.

[sunchao]

[P2] Preserve aggregate buffers when a Celeborn result stage reverts

The new composite-manager admission exposes a late fallback that the earlier aggregate-buffer checks do not cover. With native Celeborn shuffle, AQE off, and spark.comet.exec.transitionRevert.enabled=true / spark.comet.exec.transitionRevert.maxTransitions=0, an ordinary two-stage percentile(v, CAST(0.5 AS DOUBLE)) over supported partitioned native input can keep its native Partial and accepted exchange while the row-result stage reverts only Final.

With whole-stage codegen disabled and the default JVM C2R, the bridge projects the native list-of-doubles state, but Spark Final reads it as its binary value/count-map buffer. For example, a 256-double partial becomes a 2,088-byte unsafe array whose leading count is decoded as a 65,536-byte entry length; readFully cannot satisfy that request. This is a concrete missing regression case for the late result-stage fallback, rather than just an ArrayType/BinaryType declaration mismatch.

The post-rule stops at the exchange boundary; the earlier prediction/refused-exchange restoration is not run again. Please preserve a compatible producer/consumer pair or decline this late reversion.

This finding is source-derived on the examined Spark 3.5 path; no runtime reproduction was executed for this review. The reverter is pre-existing, but BASE rejects this composite-manager configuration and the new admission makes it reachable. Transition reversion is off by default.

[andygrove]

canAggregateBeConverted documents itself as a conservative check and every other branch returns a bare false, but the new FIRST/LAST branch calls withFallbackReason on the aggregate. That makes a predicate which runs over the pre-conversion plan mutate it.

As far as I can tell the tag cannot change any conversion decision. The only readers of hasFallbackReason in this file are reportUnexplainedFallback and the generic-message path in the case op => fallthrough, and both use it to decide whether to add "is not supported" or throw under COMET_STRICT_FALLBACK_REASONS, never to refuse a conversion. And when the same node reaches CometBaseAggregate.doConvert, the string written there at operators.scala:1785 is character for character the one written here, so it dedupes into the same FALLBACK_REASONS set and the explain output is identical either way.

Would dropping the withFallbackReason call and just returning false work here? The one thing it would give up is pre-satisfying the strict-reason check for that node, which doConvert already covers on its own.

[ziting-openai]

Reviewed head d5559d9812bd4cdc3e1751a168af2d979470eb84. No P1s; two significant P2s remain, documented inline: the late-fallback guard skips an inner aggregate in nested plans, and the new regression does not compile on the Spark 3.4/3.5 Scala 2.12 profiles.

Approved under the requested threshold of zero P1s and fewer than three P2s. This does not mark either P2 resolved or indicate green CI: the two legacy-profile jobs currently fail test compilation with found: Unit; required: SparkPlan; other checks remain in progress. No merge is being requested.

Static source and test inspection plus current-head CI log inspection only; I did not run tests or a live Celeborn cluster.

Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.

[ziting-openai]

[P2] Keep traversing after a Final that does not itself cross a boundary

This branch returns false without visiting the aggregate's children. With native Celeborn shuffle, AQE off, and spark.comet.exec.transitionRevert.enabled=true / maxTransitions=0, a nested query such as SELECT k, percentile(p, CAST(0.5 AS DOUBLE)) FROM (SELECT k, percentile(v, CAST(0.5 AS DOUBLE)) AS p FROM t GROUP BY k) q GROUP BY k on supported partitioned input can have Final_outer -> Partial_outer -> Final_inner -> exchange -> Partial_inner: the outer aggregation keeps the existing partitioning on k, so it needs no additional exchange. At Final_outer, both local checks are false because reachesBoundaryBeforeAggregate stops at Partial_outer; Final_inner is never examined. revertToSpark then restores the inner consumer to Spark while retaining the native exchange/producer, so Spark's percentile merge reads Comet's array state as its binary buffer and can fail. Please also recurse into children when these local checks are false, and add a nested-aggregate regression with whole-stage codegen disabled. This is source-traced, not an executed reproduction.

Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.

[ziting-openai]

[P2] Avoid using withSQLConf's result as a SparkPlan on legacy Spark

reverted is inferred as Unit on the Spark 3.4/3.5 Scala 2.12 profiles, so collect(reverted) at line 464 does not compile. Both current-head CI jobs fail in scala-test-compile-first with found: Unit; required: org.apache.spark.sql.execution.SparkPlan (Spark 3.4 log, Spark 3.5 log). Please declare the result as a SparkPlan outside withSQLConf and assign it inside the block, so the regression works with the older helper signature. Apply the same adjustment to the new producer test in RevertNativeForTransitionHeavyStagesSuite, which also captures this helper's return value.

Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.

[ziting-openai]

Reviewed head 43e1fcdc3393a31ebd81e5b990b6721ce6e9fdaf. The two affected tests now assign a typed SparkPlan inside withSQLConf, addressing the legacy-helper return-value issue in code. Current-head Spark 3.4/3.5 CI was still running during review, so this is not a claim of a green matrix.

No P1s; one existing P2 remains: the nested-aggregate fallback guard still stops before inspecting an inner buffer consumer. Production code is unchanged from the previous review, so that finding is not resolved. No duplicate inline comment posted.

Approved under the requested threshold of zero P1s and fewer than three P2s. Static source/test inspection only; I did not run tests or a live Celeborn cluster.

Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.

[sunchao]

Re-reviewed updated head 43e1fcdc3393a31ebd81e5b990b6721ce6e9fdaf. The earlier Unit-return issue and the specific FIRST/LAST diagnostic-placement issue are fixed in source. [P2] The existing nested aggregate fallback issue remains: the transition guard can stop at an outer aggregate and miss a deeper consumer of native partial state. This remains a source-derived finding; no query reproduction was run.

[ziting-openai]

Reviewed head 906d09bfc952254b276a7810e87c2cf6ad825eaf. The visitor now continues into aggregate children when its local producer/consumer checks are false, addressing the nested-aggregate P2. The added nested-percentile regression compares with Spark and requires one exchange plus both native Partial/Final pairs.

No remaining P1 or significant P2 findings in this re-review. The earlier legacy Spark test-helper fix is retained. Current-head CI is still running; this review does not claim that the new regression or full matrix has passed. Static source/test inspection only; I did not run tests or a live Celeborn cluster.

Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.

[sunchao]

Re-reviewed 906d09bfc952254b276a7810e87c2cf6ad825eaf against base/merge-base e0fa3ccb16d4c9b3a03a38b09af6ed0fb242e027.

Independently verified the nested aggregate P2 is fixed in consumed Spark/JNI execution: scalar and grouped percentile results now match Spark, with native producer/consumer pairs preserved. Matched prior-head controls still reproduce the failures. AQE, deeper nesting, null/empty inputs, boundary and compatible-SUM controls also pass. The legacy compilation fix remains intact.

Fresh validation: native core build; 360 focused JVM tests across Spark 3.4/3.5/4.1; 430 current-head SQL comparisons matching Spark, plus 55 exact-base and 55 prior-head causal controls. No remaining introduced or materially worsened P1/P2 found.

SQL validation used published Celeborn 0.6.3/0.7.0 TCP clients and production native/JNI writer-reader code with task-owned loopback lifecycle/worker fixtures. No deployed-cluster or performance validation. CI is still running.

[andygrove]

LGTM. Thanks @pingzh. One nit: I posted a review and although I can see that the feedback was addressed, it would be helpful to get a response on the comment so I know when to do another review.

[sunchao]

Merged, thanks @pingzh for the PR, @andygrove and @ziting-openai for the review!