sunchao commented on code in PR #55839:
URL: https://github.com/apache/spark/pull/55839#discussion_r3816463036
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEPropagateEmptyRelation.scala:
##########
@@ -52,19 +53,46 @@ object AQEPropagateEmptyRelation extends
PropagateEmptyRelationBase {
// - positive value means an estimated row count which can be
over-estimated
// - none means the plan has not materialized or the plan can not be
estimated
private def getEstimatedRowCount(plan: LogicalPlan): Option[BigInt] = plan
match {
- case LogicalQueryStage(_, stage: QueryStageExec) if stage.isMaterialized =>
+ case LogicalQueryStage(_, physicalPlan) =>
+ getEstimatedRowCount(physicalPlan)
+
+ case _: EmptyRelation => Some(0)
+
+ case _ => None
+ }
+
+ private def getEstimatedRowCount(plan: SparkPlan): Option[BigInt] = plan
match {
Review Comment:
I kept the correctness-sensitive helper separate and documented why
`LogicalQueryStage.computeStats()` cannot be reused directly here.
`computeStats()` walks through arbitrary physical operators and may fall back
to logical estimates; that is acceptable for planning statistics, but using it
to prove emptiness/nonemptiness could incorrectly cross a filter or join and
change query results. This rule instead accepts runtime statistics only from a
materialized stage through an explicit allowlist and applies the
global-aggregate one-row correction locally. Consolidating the helpers would
first require changing the broader statistics contract, which is outside this
patch.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -986,6 +1042,23 @@ case class AdaptiveExecutionContext(session:
SparkSession, qe: QueryExecution) {
val stageCache: TrieMap[SparkPlan, ExchangeQueryStageExec] =
new TrieMap[SparkPlan, ExchangeQueryStageExec]()
+ private val stageLifecycleLock = new Object
+
+ private[adaptive] def withStageLifecycleLock[T](body: => T): T = {
+ stageLifecycleLock.synchronized(body)
+ }
+
+ private val sharedStageResults =
+ new ConcurrentHashMap[AtomicReference[Option[Any]], Boolean]()
+
+ def markSharedStageResult(resultOption: AtomicReference[Option[Any]]): Unit
= {
+ sharedStageResults.put(resultOption, true)
+ }
+
+ def isSharedStageResult(resultOption: AtomicReference[Option[Any]]): Boolean
= {
Review Comment:
Fixed. Both `markSharedStageResult` and `isSharedStageResult` are now
`private[adaptive]`, matching the lifecycle-lock helper and keeping this
bookkeeping off the public `AdaptiveExecutionContext` surface.
##########
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 =>
Review Comment:
Fixed. The adopted physical plan is scanned once to build the retained
stage-ID and result-identity sets, and obsolete candidates are checked against
those sets without a nested `exists`. Cache removal uses the exchange's
canonicalized key rather than scanning the entire `stageCache` for every
obsolete stage. The stage-lifecycle critical section is also kept clear of a
blocking shuffle-submission/cancellation wait.
##########
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:
Agreed. `stagesToReplace` only describes the current adoption window, so it
cannot be the sole source of obsolete-stage candidates. Candidate discovery now
unions and deduplicates stages from the previous physical plan and the latest
replacement batch before comparing them with the newly adopted plan. The
existing materialized/shared-stage guards remain in place; an obsolete
broadcast or identifiable cleanup-protected shuffle that cannot safely be
cancelled is still tracked so a later ordinary failure cannot abort the query
once the stage is confirmed private and unreferenced. Opaque delegated shuffles
remain excluded because their hidden delegate shuffle IDs cannot safely be
registered. The deterministic regression constructs an old physical plan
containing an unfinished stage absent from the latest replacement batch, checks
candidate deduplication, verifies that the old stage is cancelled and its later
ordinary failure is ignored, and confirms that the newly retained stage stays
cached.
##########
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:
Agreed that a slow shuffle monitor must not stall every adaptive subquery
sharing the context. Cancellation now claims a stage-result-scoped
`CompletableFuture` reservation; exchange-cache lookup, reuse, and
shared-result marking stay under short global-lock sections, while
shuffle-monitor acquisition and `stage.cancel()` occur outside that lock. Both
cache lookup paths wait for same-result completion outside the global lock and
retry, while unrelated exchanges continue. Deterministic regressions cover a
blocked generic cancellation, a genuinely held shuffle monitor, same-result
reuse, unrelated exchange progress, and a failed cancellation whose waiting
consumer reuses the original cached stage. Atomic shuffle-submission
eligibility checks, `RemoveShuffleFiles` protections, fatal-error propagation,
guarded cleanup of identifiable failed obsolete stages, and conservative
propagation of opaque delegated-shuffle failures remain intact.
--
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]
