@@ -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 }
0 commit comments