peter-toth commented on code in PR #55839:
URL: https://github.com/apache/spark/pull/55839#discussion_r3702602173
##########
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:
**Finding 1.** The description (and the JIRA) describe a narrower change
than this diff.
- The wrapper list under "What changes were proposed" names
`AQEShuffleReadExec`, `SortExec`, `ProjectExec`, `ColumnarToRowExec`, but the
rule also gained a `BaseLimitExec` branch
(`AQEPropagateEmptyRelation.scala:82`).
- "cancellation failures from those obsolete stages are ignored instead of
being treated as query failures" reads as though AQE already cancelled obsolete
stages. It doesn't. On the merge base, `git grep '\.cancel(' --
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/` returns
exactly one hit, in `cleanUpAndThrowException`. This method is what makes AQE
proactively kill exchange stages that re-planning dropped, and that is a
user-visible change: jobs get cancelled mid-flight and
`SparkListenerJobEnd(JobFailed)` events show up where none did before -- your
own test asserts one.
- The `stageLifecycleLock` that now serialises exchange reuse against
cancellation, and the `markSharedStageResult` / `isSharedStageResult` tracking,
aren't mentioned at all.
- "How was this patch tested" lists 2 tests; the PR adds 5.
Since the cancellation is the part with a real behaviour change, please
spell it out under "Does this PR introduce _any_ user-facing change?" so it's
findable from the commit log.
##########
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:
**Finding 2.** Neither end-to-end test fails on the unfixed rule. Reproduced
in a worktree at `5a4b32a`:
```
git checkout cce435599ed3 --
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEPropagateEmptyRelation.scala
build/sbt 'sql/testOnly *AdaptiveQueryExecSuite -- -z "empty materialized
stage short-circuits AQE through sort wrappers" -z "empty filtered global
aggregate stage preserves its single output row"'
```
Both pass, `JobFailed`/"cancel" assertion included.
The reason is that there is no sort wrapper in play. `EnsureRequirements`
builds its `SortExec` with the plain constructor, so it carries neither
`LOGICAL_PLAN_TAG` nor `LOGICAL_PLAN_INHERITED_TAG` and `SparkPlan.logicalLink`
is `None` for it (`SparkPlan.scala:112`).
`replaceWithQueryStagesInLogicalPlan`'s `collectFirst` therefore resolves
`physicalNode` to the `ShuffleQueryStageExec` itself,
`LogicalQueryStage.isDirectStage` is `true`, and the pre-existing `case stage:
QueryStageExec if stage.isMaterialized` branch is what collapses the join -- on
base exactly as here. The test name promises the new path and measures the old
one.
(The second test is a legitimate guard against the new `case aggregate:
BaseAggregateExec => getEstimatedRowCount(aggregate.child)` fall-through, which
is why it also passes on base -- base never had that fall-through. No complaint
there, it just doesn't cover the new capability either.)
The `nonEmpty` side of the global-aggregate case is the cheapest thing to
pin, since it's the one place where the new `Some(1)` changes an outcome:
`PropagateEmptyRelationBase` only consults `nonEmpty(p.right)` for a
`LeftSemi`/`LeftAnti` join with **no** join condition. On base
`getEstimatedRowCount` returns `None` for `LogicalQueryStage(_,
HashAggregateExec(no grouping) -> stage)` and the join stays; here it returns
`Some(1)` and the join is rewritten to its left side. Something like:
```scala
val agg = spark.range(1).where("id < 0").agg(count("*").as("c"))
val df = testData.join(agg, Seq.empty[String], "left_semi")
checkAnswer(df, testData.collect().toSeq)
assert(collect(stripAQEPlan(df.queryExecution.executedPlan)) { case j:
BaseJoinExec => j }.isEmpty)
```
I haven't run that exact query, so please confirm the join really arrives
with `condition = None` and that the assertion fails with the rule at base
before committing it -- the SQL `EXISTS` spelling is the obvious alternative
but `RewriteNonCorrelatedExists` may rewrite it out from under you.
##########
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:
**Finding 6.** None of the new stage-lifecycle machinery is documented, and
this is the corner of AQE where the invariants are hardest to re-derive.
`withStageLifecycleLock`, `markSharedStageResult`, `isSharedStageResult`,
`cancelObsoleteStages` and `removeStageFromCache` have no scaladoc at all, so
the safety argument lives only in the reviewer's head. Three parts of it are
worth writing down, because I had to reconstruct each one:
- why `!stage.isMaterialized` is sufficient for a stage that is nested
inside a *later* stage's `plan` and therefore invisible to
`newPhysicalPlan.collect` (it can only be nested there once materialized, since
`createNonResultQueryStages` creates a stage only when
`allChildStagesMaterialized`);
- why `isSharedStageResult` is needed on top of the `resultOption.eq`
retention check (it covers reuse by a *different* `AdaptiveSparkPlanExec`
sharing this `context` -- a subquery -- whose reuse instance isn't in this plan
at all);
- why the lock has to span the `stageCache` lookup and `reuseQueryStage`
together (`markSharedStageResult` must be visible before a concurrent
cancellation reads it).
One consequence also deserves stating outright: `sharedStageResults` is
never cleared, so once an exchange has been reused anywhere in the query,
neither it nor any of its reuse instances can ever be cancelled as obsolete
again. That's a fine trade, but it means the optimization silently stops
applying to reuse-heavy plans, which a future reader shouldn't have to discover
by experiment.
##########
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:
**Finding 8.** Worth considering landing this as two PRs, because these are
two separate features rather than one change.
Neither half depends on the other:
- `AQEPropagateEmptyRelation` references nothing the other half adds -- no
`cancelObsoleteStages`, no `withStageLifecycleLock`, no
`markSharedStageResult`. Its only new import is `execution.{BaseLimitExec,
ColumnarToRowExec, ProjectExec, SortExec, SparkPlan}`.
- `cancelObsoleteStages` references neither the rule nor `EmptyRelation`.
It's driven purely by `stagesToReplace` minus what survives in
`newPhysicalPlan`, i.e. by *any* re-plan that drops a branch -- a join flipping
to a broadcast, a skew-join rewrite, an `EmptyRelation` from the pre-existing
direct-stage propagation. It doesn't need this rule in order to have work to do.
That second point isn't only a reading of the code: when I reverted
`AQEPropagateEmptyRelation.scala` to the merge base and re-ran the new tests
(finding 2), the cancellation still fired -- `TaskKilled (Stage cancelled:
[SPARK_JOB_CANCELLED] ... The query stage is no longer referenced by the
current adaptive plan.)` across two jobs. The cancellation half is fully
exercised by AQE behaviour that already exists on master.
So this is two features sharing one JIRA, and the bundling has already cost
something concrete. Each needs its own "what changed / why / how tested" --
finding 1 is largely the description of one feature being stretched over two.
And the rule half's gaps (findings 2 and 3) are easy to miss beside the
cancellation work: the three approvals so far all describe the wrapper
recursion as verified for safety, which it is, without touching whether it runs.
##########
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:
**Finding 10.** `markSharedStageResult` and `isSharedStageResult` are fully
public, while the sibling `withStageLifecycleLock` a few lines up is
`private[adaptive]`. Every caller, `AdaptiveQueryExecSuite` included, is in
`org.apache.spark.sql.execution.adaptive`, so `private[adaptive]` works for
both and keeps them off `AdaptiveExecutionContext`'s public 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 =>
+ 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:
**Finding 4.** `removeStageFromCache(stage)` runs before
`stage.cancel(...)`, so a cancellation that throws leaves the stage in the
worst of both worlds: evicted from `stageCache` but still materializing, and
absent from `obsoleteCancelledStageIds`.
Failure scenario: `ShuffleQueryStageExec.doCancel` -> `cancelShuffleJob` ->
`futureAction.get().foreach(_.cancel(reason))` throws (or
`BroadcastQueryStageExec` -> `cancelJobsWithTag` throws). We log and return
`None`, and the stage keeps running. Its cache entry is gone, so the next
identical exchange in the query builds a second stage for the same shuffle
instead of reusing this one. Then, when the half-cancelled stage reports
`StageFailure`, `obsoleteCancelledStageIds.contains(stage.id)` is false at
`:361`, so `errors.append(ex)` -> `cleanUpAndThrowException` fails the whole
query over a stage that is no longer in the plan -- exactly the failure this PR
set out to remove.
@viirya read this as safe on the grounds that a failed cancel isn't silently
swallowed. That's accurate; my point is that for an obsolete stage, not
swallowing it is the bug. Note this does not undo @shrirangmhalgi's fix: stages
*skipped* because they are materialized or shared must still return `None` and
must not be suppressed, and they do below.
```scala
if (!stage.isMaterialized &&
!context.isSharedStageResult(stage.resultOption)) {
try {
stage.cancel("The query stage is no longer referenced by the current
adaptive plan.")
} catch {
case NonFatal(t) =>
logError(s"Exception in cancelling obsolete query stage:
${stage.treeString}", t)
}
// The stage is obsolete either way: drop it from the cache so it can't be
reused, and
// suppress its materialization failure, which is no longer relevant to
this query.
removeStageFromCache(stage)
Some(stage.id)
} else {
None
}
```
`obsolete stage cancellation preserves shared reused exchange stages` pins
the current behaviour (`assert(failedCancelledIds.isEmpty)`), so it would need
updating alongside.
##########
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:
**Finding 3.** Three of the five wrapper branches never execute. I wrapped
each new case in a marker printing the branch name and its return value, then
ran the whole suite (`build/sbt 'sql/testOnly *AdaptiveQueryExecSuite'`, 134
tests, all green):
| branch | hits |
| --- | --- |
| `SortExec` | 24 (12x `Some(10)`, 4x `Some(80)`, 8x `None`) |
| `BaseAggregateExec`, empty grouping | 34 (17x `Some(2)`, 12x `Some(1)`, 5x
`None`) |
| `BaseLimitExec` | 10 -- all ten from `limit wrappers propagate only
provably empty query stages`, none from a real query |
| `AQEShuffleReadExec` | 0 |
| `ProjectExec` | 0 |
| `ColumnarToRowExec` | 0 |
So the shape the description leads with cannot be the one that fires --
`AQEShuffleReadExec` is never reached. That's consistent with how
`physicalNode` is picked (finding 2): a physical wrapper only becomes
`LogicalQueryStage.physicalPlan` if it carries the same logical link as the
stage below it, and `AQEShuffleReadExec` (built by `OptimizeSkewedJoin` with
the plain constructor, no link) doesn't, while `ColumnarToRowExec` is inserted
by `postStageCreationRules` and so lands *inside* a stage's own plan or inside
`ResultQueryStageExec.plan`, after which there is no more re-planning. And
wherever a physical wrapper does sit above a stage, the corresponding *logical*
operator sits above the `LogicalQueryStage`, where
`PropagateEmptyRelationBase.commonApplyFunc` already propagates emptiness
through `Project` / `Filter` / `Sort` / `GlobalLimit` / `LocalLimit` / `Offset`
/ `RepartitionOperation`.
Worth noting that even `SortExec` never returned `Some(0)` in that run -- it
only ever fed `nonEmpty`, so the one branch that does execute has no observed
effect on `isEmpty`.
Either way out is fine by me: show a query that reaches each branch and
assert the outcome in a test that fails on base, or drop the branches that
can't be reached and keep the aggregate case, which does fire 34 times and is
the part with a real semantic delta over base.
This is where I part company with the approvals: @yadavay-amzn and @viirya
both verified that the four wrappers *preserve* emptiness, which I agree with
-- a pass-through can't be unsafe. Reachability is the separate question, and
it's the one that decides whether the branches earn their place.
##########
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:
**Finding 7.** This test turns on a race. `Thread.sleep(200)` x 200 rows
over 20 partitions has to keep the right stage unfinished until the left stage
materializes, the re-plan runs and `cancelObsoleteStages` reaches it; if the
right side wins, no job is cancelled and the `JobFailed` assertion at
`:398-404` fails. It also costs seconds of wall clock in an otherwise
sub-second suite (2.7s and 4.7s across my two runs), and
`e.getMessage.toLowerCase(Locale.ROOT).contains("cancel")` couples it to the
wording of `SPARK_JOB_CANCELLED`.
Two cheaper options: have the right side block on a `CountDownLatch` the
test releases after observing the cancellation, instead of a fixed sleep; or
drop the `JobFailed` assertion from this test entirely and leave cancellation
to `obsolete stage cancellation preserves shared reused exchange stages`, which
already covers it deterministically.
##########
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:
**Finding 5.** A cancelled obsolete shuffle stage never registers its
`shuffleId`, so whatever map output it had already written stays on disk for
the rest of the application.
`context.shuffleIds` is populated only in the success branch of the
materialization callback (`:329-334`). `SQLExecution.extractShuffleIds` prefers
that set over scanning the plan whenever an `AdaptiveSparkPlanExec` is present
(`SQLExecution.scala:111-121`), and it is what drives `shuffleCleanupMode` at
the end of the execution (`SQLExecution.scala:242-264`): `RemoveShuffleFiles`
-> `shuffleDriverComponents.removeShuffle`, `SkipMigration` ->
`addShuffleToSkip`. Before this PR the obsolete stage ran to completion, hit
the success branch and got cleaned up with the rest; now it is cancelled, so
the map tasks that finished before the cancellation took effect are only
reclaimed when `ContextCleaner` GCs the RDD.
The obvious fix is to mirror the success-path registration for the stages we
cancel:
```scala
stage.plan.collect { case s: ShuffleExchangeLike =>
context.shuffleIds.put(s.shuffleId, true) }
```
One wrinkle to handle: `ShuffleExchangeExec.shuffleId` is
`shuffleDependency.shuffleId` off a `lazy val`, so touching it on a stage whose
job never started would *create* a shuffle dependency rather than describe one.
Worth gating on the stage having actually been submitted.
##########
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:
**Finding 9.** `LogicalQueryStage.computeStats()` already does what this
method does, global-aggregate special case included
(`LogicalQueryStage.scala:57-71`):
```scala
val physicalStats = physicalPlan.collectFirst {
case a: BaseAggregateExec if a.groupingExpressions.isEmpty =>
a.collectFirst { case s: QueryStageExec => s.computeStats()
}.flatten.map { stat =>
if (stat.rowCount.contains(0)) stat.copy(rowCount = Some(1)) else stat
}
case s: QueryStageExec => s.computeStats()
}.flatten
```
So the tree now holds two implementations of "find the row count under this
`LogicalQueryStage`'s physical plan, and count a global aggregate over an empty
stage as one row", with different traversal rules: `computeStats` uses
`collectFirst` and therefore looks through *any* node -- its own `TODO this is
not accurate when there is other physical nodes above QueryStageExec` says as
much, and a `FilterExec` or a join above the stage would be walked straight
through -- while this method uses an allowlist. They can disagree, and a fix to
one won't reach the other.
Worth considering: lift the allowlist walk into a single helper next to
`LogicalQueryStage` and have both `computeStats` and this rule call it. That
also tightens `computeStats`, since the allowlist is the more conservative of
the two. If you'd rather not touch `computeStats` here, a comment on each
pointing at the other would at least make the coupling visible.
--
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]
