yew1eb commented on PR #3853: URL: https://github.com/apache/celeborn/pull/3853#issuecomment-6058852415
@venkata91 Thanks for the thorough review and the repros — all findings confirmed and addressed in the latest commit. Summary of the changes: **Blocker 1 — Spark 4 + AQE no-op.** Fixed. The extractor now matches generic `QueryStageExec` and descends into `stage.plan`, so Spark 4's `ResultQueryStageExec` leaf is reached; `TableCacheQueryStageExec` (plan = `InMemoryTableScanExec`) contains no shuffles and is unaffected. The suite now passes 10/10 locally on both Spark 3.5 and Spark 4.0. **Blocker 2 — silent empty results with `stageRerun.enabled=false`.** Init now fails fast when cleanup is enabled but stage rerun is not, following the existing `fetch.cleanFailedShuffle` precedent. Your repro scenarios are covered by new tests: sequential and concurrent re-execution of the same DataFrame, both returning correct results with stage rerun on. The `UNKNOWN_APP_SHUFFLE_ID` race you mentioned is left as a follow-up since neither of us could trigger it. **Lazy `shuffleDependency`.** Extraction is now side-effect free on the listener thread: - Shuffle ids are read only from materialized `ShuffleQueryStageExec` (dependency already computed); raw exchanges inside an `AdaptiveSparkPlanExec` subtree are never touched. - For failed non-AQE queries, exchanges with `RangePartitioning` output are skipped (reading them could launch a sampling job); the `isAppShuffleRegistered` filter drops anything never actually registered. - Two subtleties surfaced while testing the multi-join case: AQE nests earlier stages inside later stages' plan subtrees, so traversal must continue below stage/exchange boundaries; and identical inputs (self-joins) canonicalize multiple exchanges into one. - Scalar subqueries are covered via `subqueriesAll`. Non-vanilla `ShuffleExchangeLike` implementations (e.g. Gluten) remain intentionally unsupported to avoid coupling to third-party exchange internals. **Planning failures / nits.** Failed queries are detected via `SparkListenerSQLExecutionEnd.executionFailure`; the residual re-planning risk of `qe.executedPlan` is contained by the `NonFatal` catch (no public `isExecutedPlanAvailable` exists on Spark 3.5/4.0–4.2). `LifecycleManager` is now constructor-injected, and `Throwable` was narrowed to `NonFatal`. **Config key.** Aligned to `celeborn.client.spark.sql.queryEndShuffleCleanup.enabled`; docs regenerated via `ConfigurationSuite`. `spark-connect-server` added to the doc and PR description. **Tests.** Added: failed AQE query, sequential/concurrent DataFrame re-execution, and the stage-rerun fail-fast case; removed the redundant `spark.stop()` calls. **Re: scope and SPARK-47764.** Confirmed it shipped in 4.0.0 scoped to Spark Connect (`shuffleCleanupMode` defaults to `DoNotCleanup`); the classic and ThriftServer variants arrived only in 4.1 and 4.2 respectively, all opt-in. So this feature still covers ground Spark doesn't: all of Spark 3.x and classic/Thrift on 4.0. Both implementations share the same `shuffleDriverComponents.removeShuffle` entry point, so they coexist idempotently if a user enables both. That said, narrowing this listener's scope on Spark 4.1+/4.2+ in favor of the built-in cleanup could be a reasonable follow-up once those versions are the common denominator. Happy to share the debug logs for the multi-join nesting case if useful. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
