dongjoon-hyun commented on code in PR #55839:
URL: https://github.com/apache/spark/pull/55839#discussion_r3815937714
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -397,6 +401,8 @@ case class AdaptiveSparkPlanExec(
currentPhysicalPlan.treeString,
newPhysicalPlan.treeString).mkString("\n")
logOnLevel(log"Plan changed:\n${MDC(QUERY_PLAN, plans)}")
cleanUpTempTags(newPhysicalPlan)
+ obsoleteCancelledStageIds ++=
+ cancelObsoleteStages(newPhysicalPlan, stagesToReplace)
Review Comment:
`cancelObsoleteStages` only ever sees `stagesToReplace`, and
`stagesToReplace` is reset three lines below at `:408` every time a plan is
adopted. So this only covers stages created since the *last* adoption.
`replaceWithQueryStagesInLogicalPlan` wraps a stage into a
`LogicalQueryStage` regardless of whether it has materialized, so a stage that
is still materializing can be folded into `currentLogicalPlan` by one adoption
and then be invisible to the next. Concretely:
- Round 1: stages A (empty, fast) and B (slow, still materializing) are
created. A completes, a re-plan is adopted for an unrelated reason. B now lives
in `currentLogicalPlan` as a `LogicalQueryStage`, and `stagesToReplace` is
cleared.
- Round 2: stage C materializes, re-planning proves B's branch empty and
drops it from `newPhysicalPlan`.
- `cancelObsoleteStages(newPhysicalPlan, stagesToReplace)` only sees `[C]`.
B is neither cancelled nor added to `obsoleteCancelledStageIds`, so it keeps
running; and if it later fails (task error, executor loss) `errors.append(ex)`
at `:361` fires and `cleanUpAndThrowException` aborts the query over a stage
that is no longer in any plan.
All five new tests exercise a single adoption, so this isn't covered. If the
intent is "cancel every exchange stage the adopted plan no longer references",
the candidate set needs to come from `currentPhysicalPlan` before the swap
rather than from `stagesToReplace` alone.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -420,6 +426,45 @@ case class AdaptiveSparkPlanExec(
.get.asInstanceOf[T]
}
+ private def cancelObsoleteStages(
+ newPhysicalPlan: SparkPlan,
+ stagesToReplace: Seq[QueryStageExec]): Seq[Int] = {
+ val newStages = newPhysicalPlan.collect {
+ case stage: QueryStageExec => stage
+ }
+ val obsoleteStages = stagesToReplace.collect {
+ case stage: ExchangeQueryStageExec
+ if !newStages.exists(newStage =>
+ newStage.id == stage.id ||
newStage.resultOption.eq(stage.resultOption)) => stage
+ }
+ obsoleteStages.flatMap { stage =>
+ context.withStageLifecycleLock {
Review Comment:
Holding `stageLifecycleLock` across `stage.cancel(...)` keeps a query-global
lock over a call that can block for a long time.
`ShuffleQueryStageExec.doCancel` -> `cancelShuffleJob` is
`this.synchronized` on the exchange node (`ShuffleExchangeExec.scala:129`). The
same monitor is held by `triggerFuture` while it runs
`mapOutputStatisticsFuture` (`ShuffleExchangeExec.scala:92-105`), which
evaluates `inputRDD.getNumPartitions` and `shuffleDependency`
(`ShuffleExchangeExec.scala:227-233`) -- for a file-scan child that is
partition planning and file listing, which can take seconds.
While the AQE main thread waits on that monitor, every other thread sharing
this `AdaptiveExecutionContext` -- each subquery's `AdaptiveSparkPlanExec` --
is blocked in `createNonResultQueryStages` at `:672` / `:699` waiting for
`stageLifecycleLock`. Nothing held a shared lock across job teardown before
this PR.
The race the lock closes is between the `stageCache` lookup /
`markSharedStageResult` and the cache eviction, so keeping the `isMaterialized`
/ `isSharedStageResult` checks and `removeStageFromCache` inside the lock while
moving `stage.cancel(...)` outside it would preserve the invariant without
serialising stage creation on job cancellation. That would also take the
quadratic `stageCache` scan @cloud-fan flagged at `:437` out of the critical
section.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
