yew1eb commented on PR #3853: URL: https://github.com/apache/celeborn/pull/3853#issuecomment-6311349312
@venkata91 Thanks for the detailed review and repros — all addressed in ed7b8099 (branch rebased onto latest main; the force-push compare shows the exact delta). **Failed non-AQE queries.** Went with (a): raw exchanges are skipped entirely when `executionFailure` is set, leaving those shuffles to GC cleanup — same as Spark's own query-end cleanup. Your `bad.union(good)` repro is now a test: a job-counting listener asserts no job is launched on the listener-bus thread after the query fails. **AQE collection.** Switched to (b): materialized `ShuffleQueryStageExec`s are collected from `AdaptiveExecutionContext.stageCache` (public in 3.0–4.2), which covers scalar subqueries, stages dropped from the final plan, and materialized children of unmaterialized stages — close to Spark 4's `context.shuffleIds`, but available on 3.x. Tree-encountered shuffle stages are no longer descended into, and the `inAdaptivePlan` reset is gone. Tests added for all three cases (your subquery repro, empty join side, failure in the second map stage). **stageRerun check.** Moved to the `SparkShuffleManager` constructor (driver only), so the SparkContext fails to start; the fail-fast test now intercepts around `getOrCreate()`. **Nits.** `hasMapping` dropped from `isAppShuffleRegistered`; config renamed to `celeborn.client.spark.sql.queryEndShuffleCleanup.enabled` (matching the PR description) with the limitations documented (SQL/Dataset only, non-vanilla exchanges like Gluten unsupported, failed/cancelled stages fall back to GC cleanup). Suite passes 14/14 locally on both Spark 3.5 and 4.0. -- 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]
