Skip to content

Commit 0896e36

Browse files
committed
fix: recognize an AQE re-plan that collapsed to an empty relation
A mark on the plan already reported cannot reach the plan AQE builds when a stage materializes empty, and marking the logical plan only reached it while the physical root's logical link happened to be the logical root. When a preparation rule drops the root operator first - `RemoveRedundantSorts` removing a sort above a sort merge join - the mark lands on the join below it, empty propagation replaces the join and then the sort, and the replacement root inherits the unmarked sort's tags. Recognize that plan by what it is instead. A re-optimized plan holds a `QueryStageExec` for every stage that materialized; the only way to eliminate all of them is empty propagation, which leaves a plan Catalyst records as producing zero rows, so the logical link answers the question directly. The check is confined to the query-stage-prep rule, so a genuinely empty query is still reported. This subsumes the logical mark, which is dropped.
1 parent 53db476 commit 0896e36

2 files changed

Lines changed: 65 additions & 34 deletions

File tree

‎spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala‎

Lines changed: 43 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -149,18 +149,6 @@ object CometExecRule {
149149
*/
150150
private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported")
151151

152-
/**
153-
* The same mark, placed on the logical plan the reported physical plan was linked to.
154-
*
155-
* AQE can replace the physical tree wholesale - a stage that materializes empty collapses the
156-
* whole plan to an empty relation - leaving a tree that shares no node with the one already
157-
* reported and holds no query stage either, so [[PLAN_ONLY_REPORTED]] cannot reach it. Every
158-
* such tree is planned from a logical plan derived from the one already reported, and Catalyst
159-
* copies tags onto replacement nodes there too, so the logical mark does reach it.
160-
*/
161-
private val PLAN_ONLY_REPORTED_LOGICAL: TreeNodeTag[Unit] =
162-
TreeNodeTag[Unit]("comet.planOnlyReportedLogical")
163-
164152
/**
165153
* Whether `plan` is the plan of a query stage AQE has just cut, which reaches the columnar rule
166154
* rooted at the `Exchange` the cut was made at.
@@ -185,15 +173,34 @@ object CometExecRule {
185173
* `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the wrapper is matched
186174
* by name rather than through its contents. A re-optimized plan holds `QueryStageExec` nodes
187175
* for the stages already materialized. A final plan for a query AQE never had to cut into
188-
* stages holds neither, and is recognized by the mark left when it was first reported. A plan
189-
* AQE rebuilt from scratch carries none of those, and is recognized by the mark on the logical
190-
* plan it was built from.
176+
* stages holds neither, and is recognized by the mark left when it was first reported.
191177
*/
192178
private def isReapplication(plan: SparkPlan): Boolean = {
193179
plan.exists(p =>
194180
p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] ||
195-
p.getTagValue(PLAN_ONLY_REPORTED).isDefined) ||
196-
plan.logicalLink.exists(_.getTagValue(PLAN_ONLY_REPORTED_LOGICAL).isDefined)
181+
p.getTagValue(PLAN_ONLY_REPORTED).isDefined)
182+
}
183+
184+
/**
185+
* Whether `plan` is a plan AQE re-optimized down to nothing, which is the one re-optimization
186+
* result [[isReapplication]] cannot recognize.
187+
*
188+
* A re-optimized plan normally holds a `QueryStageExec` for each stage that has materialized.
189+
* The exception is a stage that materialized empty: `AQEPropagateEmptyRelation` then replaces
190+
* the stages, and everything above them, with an empty relation, so the plan Spark hands back
191+
* shares no node with the plan already reported and holds no query stage either. Empty is the
192+
* only shape that eliminates every stage, and Catalyst records it as a plan whose row count is
193+
* known to be zero, so the logical link answers the question directly.
194+
*
195+
* The check is confined to the query-stage-prep rule, the only one AQE hands a re-optimized
196+
* plan to. A genuinely empty query - `WHERE false`, an empty local relation - is planned once,
197+
* reaches the columnar rule instead, and is still reported.
198+
*/
199+
private def isAdaptiveReplanToNothing(
200+
plan: SparkPlan,
201+
queryStagePrep: Boolean,
202+
aqeEnabled: Boolean): Boolean = {
203+
aqeEnabled && queryStagePrep && plan.logicalLink.exists(_.maxRows.contains(0L))
197204
}
198205

199206
/**
@@ -209,21 +216,27 @@ object CometExecRule {
209216
* - the re-optimized plan after each stage materializes, through the prep rule;
210217
* - the final plan once every stage has materialized, as a columnar rule.
211218
*
212-
* Two mechanisms sort those out, because neither on its own covers every shape:
219+
* Three mechanisms sort those out, because no one of them covers every shape:
213220
*
214-
* - A mark on the plan already reported, both on the physical tree and on the logical plan it
215-
* was built from, so that a later application recognizes it however AQE rewrote it; see
216-
* [[isReapplication]]. Marking plans individually, rather than holding one report slot per
217-
* SQL execution, is what gives the outer query a report of its own: each scalar subquery
218-
* and DPP subquery is prepared as a top-level plan in its own right, and for most of them
219-
* that happens *before* the outer plan reaches the conversion rules, so a single slot would
220-
* be consumed by a subquery and the plan being evaluated would never be described.
221+
* - A mark on the plan already reported, so that a later application recognizes it however
222+
* Spark rewrote it in between; see [[isReapplication]]. Marking plans individually, rather
223+
* than holding one report slot per SQL execution, is what gives the outer query a report of
224+
* its own: each scalar subquery and DPP subquery is prepared as a top-level plan in its own
225+
* right, and for most of them that happens *before* the outer plan reaches the conversion
226+
* rules, so a single slot would be consumed by a subquery and the plan being evaluated
227+
* would never be described.
228+
* - One re-optimization result carries no trace of the plan it replaced, and is recognized by
229+
* what it is rather than by a mark; see [[isAdaptiveReplanToNothing]].
221230
* - A subquery referenced from more than one place in the outer plan is prepared once per
222231
* reference, as a separate plan each time, differing only in expression IDs; Spark then
223232
* collapses them to one `ReusedSubqueryExec`. No mark connects those, so within one SQL
224233
* execution the canonicalized plan's hash dedupes them. Canonicalization is what makes the
225234
* expression IDs drop out; the raw structural hash sees two different plans.
226235
*
236+
* None of this is scoped to a SQL execution ID, which `df.rdd.count()` and reading
237+
* `executedPlan` without an action both plan - and in the first case execute AQE stages -
238+
* without.
239+
*
227240
* @param queryStagePrep
228241
* whether the calling rule instance is registered as a query-stage-prep rule.
229242
*/
@@ -235,11 +248,14 @@ object CometExecRule {
235248
if (isQueryStage(plan, queryStagePrep, aqeEnabled) || isReapplication(plan)) {
236249
false
237250
} else {
251+
// Mark before deciding. A plan AQE re-optimized to nothing is not worth a report, but the
252+
// final plan Spark builds from it reaches the columnar rule next and has to be recognized
253+
// as a plan already dealt with.
238254
plan.setTagValue(PLAN_ONLY_REPORTED, ())
239-
plan.logicalLink.foreach(_.setTagValue(PLAN_ONLY_REPORTED_LOGICAL, ()))
240-
// Node tags are not part of a plan's canonical form, so the marks set above do not perturb
255+
// Node tags are not part of a plan's canonical form, so the mark set above does not perturb
241256
// the key. Without an execution ID there is nothing to scope the state to, and the checks
242257
// above have already ruled out the repeat applications AQE makes, so report.
258+
!isAdaptiveReplanToNothing(plan, queryStagePrep, aqeEnabled) &&
243259
executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.canonicalized.hashCode()}"))
244260
}
245261
}

‎spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1057,24 +1057,39 @@ class CometExecRuleSuite extends CometTestBase {
10571057
(action, runIt) <- Seq[(String, org.apache.spark.sql.DataFrame => Unit)](
10581058
"collect" -> (df => df.collect()),
10591059
"toRdd.count" -> (df => df.queryExecution.toRdd.count()))
1060+
// The second shape is the one where the plan that AQE reports back cannot be recognized by a
1061+
// mark of any kind: `RemoveRedundantSorts` drops the outer sort before this rule sees the
1062+
// plan, so the mark lands on the join below it, and empty propagation replaces the join and
1063+
// then the sort, leaving a root that inherited the unmarked sort's tags.
1064+
(shape, query, marker) <- Seq(
1065+
(
1066+
"aggregate",
1067+
"SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2",
1068+
"HashAggregate"),
1069+
(
1070+
"join under a removed root sort",
1071+
"""SELECT a.id FROM range(0, 20, 1, 2) a
1072+
|JOIN range(0, 20, 1, 2) b ON a.id % 7 = b.id % 7
1073+
|WHERE a.id < 0
1074+
|SORT BY a.id % 7""".stripMargin,
1075+
"SortMergeJoin"))
10601076
} {
1061-
test(s"plan-only mode: one report when the plan becomes empty (AQE=$aqe, $action)") {
1077+
test(s"plan-only mode: one report when the plan becomes empty ($shape, AQE=$aqe, $action)") {
10621078
withSQLConf(
10631079
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
1080+
// Force a shuffled join so the join shape has a stage that can materialize empty.
1081+
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
1082+
SQLConf.SHUFFLE_PARTITIONS.key -> "2",
10641083
CometConf.COMET_ENABLED.key -> "true",
10651084
CometConf.COMET_EXEC_ENABLED.key -> "true",
10661085
CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
1067-
val reports = capturePlanOnlyReports {
1068-
runIt(
1069-
spark.sql(
1070-
"SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2"))
1071-
}
1086+
val reports = capturePlanOnlyReports(runIt(spark.sql(query)))
10721087
assert(
10731088
reports.size == 1,
10741089
s"expected one report, got ${reports.size}:\n${reports.mkString("\n\n")}")
10751090
// The one report must describe the query, not the empty relation AQE replaced it with.
10761091
assert(
1077-
reports.head.contains("HashAggregate"),
1092+
reports.head.contains(marker),
10781093
s"the report does not describe the query:\n${reports.head}")
10791094
}
10801095
}

0 commit comments

Comments
 (0)