sunchao commented on code in PR #55839:
URL: https://github.com/apache/spark/pull/55839#discussion_r3816456749
##########
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(
Review Comment:
Agreed. The published PR description still needs to be updated. I prepared a
replacement that distinguishes the two behaviors: conservative
runtime-row-count propagation through the supported sort, limit, and aggregate
paths, and lifecycle handling of private obsolete exchange stages. It documents
proactive shuffle cancellation and user-visible listener events, the
broadcast/delegated-shuffle/cleanup-mode exclusions, guarded handling of
eventual failures from identifiable obsolete stages, the stricter treatment of
opaque delegated shuffles, cross-plan stage tracking, exchange-reuse
synchronization, fatal-error propagation, and regression coverage.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -363,6 +368,314 @@ class AdaptiveQueryExecSuite
}
}
+ test("empty materialized stage short-circuits AQE through sort wrappers") {
Review Comment:
Good catch. I replaced the ineffective end-to-end assertion with the
conditionless left-semi join over a global aggregate that you suggested. The
regression verifies that the original plan contains the join, that the
aggregate remains a one-row result over empty input, and that AQE eliminates
the join using the new `Some(1)` information. Its direct rule assertion depends
on the global-aggregate cardinality handling rather than only checking an
unchanged final result. The final combined patch passed the complete 139-test
suite.
##########
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 {
+ case stage: QueryStageExec if stage.isMaterialized =>
stage.getRuntimeStatistics.rowCount
- case LogicalQueryStage(_, agg: BaseAggregateExec) if
agg.groupingExpressions.nonEmpty &&
- agg.child.isInstanceOf[QueryStageExec] =>
- val stage = agg.child.asInstanceOf[QueryStageExec]
- if (stage.isMaterialized) {
- stage.getRuntimeStatistics.rowCount
+ case read: AQEShuffleReadExec =>
Review Comment:
Agreed. I removed the `AQEShuffleReadExec`, `ProjectExec`, and
`ColumnarToRowExec` branches instead of treating emptiness preservation as
evidence of reachability. The remaining allowlist covers `SortExec`,
`BaseLimitExec`, and `BaseAggregateExec`. The wrapper regression obtains a
logically linked sort/shuffle shape from a real planned `ORDER BY ... LIMIT ...
DISTINCT` query and directly exercises sort and limit propagation using that
shape; the limit-specific assertion is a direct rule test, not a claim that a
naturally executed query necessarily chooses the limit branch. The
global-aggregate regression observes the new nonempty inference directly.
##########
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 {
+ if (!stage.isMaterialized &&
!context.isSharedStageResult(stage.resultOption)) {
+ removeStageFromCache(stage)
Review Comment:
Fixed while retaining the distinction between retained/shared stages,
successfully cancelled stages, identifiable obsolete stages whose cancellation
fails or is unsafe, and opaque delegated shuffles. Cancellation is attempted
before the exchange-cache entry is removed, so a thrown cancellation leaves the
existing stage available for reuse. Failed attempts, obsolete broadcasts, and
known submitted cleanup-protected shuffles are tracked separately from actual
successful cancellations. Their later nonfatal materialization failure is
ignored only if the stage is still obsolete and private; at that point any
visible submitted shuffle ID is registered and the failed cache entry is
evicted. Fatal failures, failures from retained/shared stages, and failures
from opaque delegated shuffles continue to propagate; the opaque delegated
stage also remains cached. The regressions check cancellation IDs,
failed-cancellation cache retention, eventual eligible-stage eviction, and
conservative deleg
ated-stage handling.
##########
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 {
+ if (!stage.isMaterialized &&
!context.isSharedStageResult(stage.resultOption)) {
+ removeStageFromCache(stage)
+ try {
+ stage.cancel("The query stage is no longer referenced by the
current adaptive plan.")
Review Comment:
Fixed. A submitted shuffle ID is registered before cancellation is
attempted, including when cancellation throws. Submission is checked before
reading `shuffleId`, so an unsubmitted exchange does not initialize its lazy
shuffle dependency. Submitted shuffles are not proactively cancelled under
`RemoveShuffleFiles`, because Spark acknowledges cancellation before executor
map writers have necessarily stopped; otherwise end-of-query cleanup could
delete files while those writers are still active. Such an identifiable
obsolete shuffle is still tracked: if it later fails with an ordinary error
while still private and unreferenced, its submitted shuffle ID is registered,
its failed cache entry is removed, and the irrelevant failure is ignored
without cancelling the underlying action. Wrapped fatal errors and failures
after reuse still propagate. An opaque delegated wrapper with an empty
`futureAction` is handled more conservatively: it is not cancelled, its failure
is not suppressed, an
d its cache entry is retained because the delegate's hidden shuffle ID cannot
safely be registered.
##########
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
= {
Review Comment:
Added documentation for the stage-lifecycle lock, shared-result bookkeeping,
cancellation eligibility, guarded eventual failure handling for identifiable
uncancellable obsolete stages, cache removal, and cleanup registration. Opaque
delegated shuffles remain excluded from failure suppression because their
hidden delegate shuffle IDs cannot be safely registered. The comments explain
why a stage nested under a later stage must already be materialized, why
cross-subquery reuse needs shared-result identity tracking beyond the current
plan, why cache lookup/reuse/marking must remain atomic, and why once-shared
results stay protected for the lifetime of the adaptive context. A
stage-result-scoped `CompletableFuture` reservation also lets both
exchange-cache lookup paths wait for that result outside the query-global lock
while unrelated exchanges continue.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -363,6 +368,314 @@ class AdaptiveQueryExecSuite
}
}
+ test("empty materialized stage short-circuits AQE through sort wrappers") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
+ val jobEndEvents = new mutable.ArrayBuffer[SparkListenerJobEnd]
+ val listener = new SparkListener {
+ override def onJobEnd(jobEnd: SparkListenerJobEnd): Unit =
jobEndEvents.synchronized {
+ jobEndEvents += jobEnd
+ }
+ }
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ val left = spark.range(0, 1, 1, 1).where("id <
0").select($"id".as("k"))
+ val right = spark.range(0, 200, 1, 20).as[Long].map { id =>
Review Comment:
Fixed. I removed the sleep-based job-cancellation/listener assertion from
the wrapper test. Rule propagation and exchange-stage lifecycle behavior are
now checked separately using deterministic plan and stage assertions, so the
regression does not depend on relative task completion times or
cancellation-message wording.
##########
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 ++=
Review Comment:
I agree these are independent changes: row-count propagation can stand
alone, and obsolete-stage cancellation also applies to pre-existing adaptive
replans. The implementation and regression coverage now distinguish the two
behaviors, and I have prepared an updated PR description that still needs to be
applied. I am happy to split proactive stage cancellation into a separate
PR/JIRA if you would prefer to review and land the propagation change
independently.
--
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]
