From b666339d13e99ee74492a67bc32f761cc03a7b64 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 19 Aug 2026 06:09:14 -0600
Subject: [PATCH 1/8] feat: add spark.comet.explain.planOnly.enabled

Builds the Comet plan Comet would have executed and logs it to the driver
log, then discards it and lets Spark run the query unchanged. Both
CometScanRule and CometExecRule short-circuit when the config is set;
CometExecRule uses thread-local bypass flags on both rules to force normal
behavior while it constructs the preview for the report. The WARN is
deduped per SQL execution ID so AQE emits one report per query.

Closes #5335.
---
 .../latest/understanding-comet-plans.md       | 21 +++++
 .../scala/org/apache/comet/CometConf.scala    | 12 +++
 .../apache/comet/rules/CometExecRule.scala    | 74 ++++++++++++++++++
 .../apache/comet/rules/CometScanRule.scala    | 22 ++++++
 .../comet/rules/CometExecRuleSuite.scala      | 76 +++++++++++++++++++
 5 files changed, 205 insertions(+)

diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md
index a600615b4b..a12c3f9f60 100644
--- a/docs/source/user-guide/latest/understanding-comet-plans.md
+++ b/docs/source/user-guide/latest/understanding-comet-plans.md
@@ -208,6 +208,27 @@ operators were arranged after Comet's serialization). See the
 [Metrics Guide](metrics.md) for details on the DataFusion metrics that appear
 in this output.
 
+### `spark.comet.explain.planOnly.enabled`
+
+When enabled, Comet runs its full conversion pass on every query and logs the
+resulting Comet plan and coverage summary to the driver log, then reverts to
+executing the plan on Spark instead of offloading anything to native. Use this
+to evaluate how much of a workload Comet would accelerate without changing the
+execution.
+
+The log line is prefixed with `[Comet plan-only]` and includes the same
+annotated plan and summary as `spark.comet.explain.format=verbose` produces
+against a normal Comet plan. Under AQE the report is emitted once per SQL
+execution.
+
+The estimate reflects Scala-side conversion only. The native plan is never
+handed to DataFusion, so anything that would have failed in DataFusion's
+`create_plan` still counts as accelerated. Treat the percentage as an upper
+bound.
+
+The config requires `spark.comet.exec.enabled=true`. With Comet exec disabled
+the rule that emits the report does not run.
+
 ## Programmatic Access to Fallback Reasons
 
 The configs above route fallback reasons to logs or the SQL UI. If you want
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index f40d0cbdd9..94f8fafd86 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -688,6 +688,18 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_EXPLAIN_PLAN_ONLY_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.explain.planOnly.enabled")
+      .category(CATEGORY_EXEC_EXPLAIN)
+      .doc("When enabled, Comet builds the Comet plan it would have executed and logs it to " +
+        "the driver log, then discards it and lets Spark execute the query. Use this to " +
+        "evaluate how much of a workload Comet would accelerate without changing execution. " +
+        "The estimate is Scala-side only; native planning failures are not surfaced, so the " +
+        "acceleration percentage can be optimistic. Requires `spark.comet.exec.enabled=true`. " +
+        "Disabled by default.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_STRICT_FALLBACK_REASONS: ConfigEntry[Boolean] =
     conf("spark.comet.explain.fallback.strict.enabled")
       .category(CATEGORY_TESTING)
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index cbf6c6a1e8..5d7aa95e13 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -115,6 +115,47 @@ object CometExecRule {
    */
   val SKIP_COMET_BROADCAST_TAG: org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit] =
     org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometBroadcast")
+
+  /**
+   * Bounded set of SQL execution IDs for which plan-only mode has already logged its report. Used
+   * to dedupe the WARN under AQE, where `CometExecRule` fires once against the full initial plan
+   * and again per stage.
+   */
+  private val PLAN_ONLY_REPORTED_LIMIT = 1024
+  private val planOnlyReportedIds: java.util.LinkedHashSet[String] =
+    new java.util.LinkedHashSet[String]()
+
+  /**
+   * Thread-local re-entry guard: set to true while `reportPlanOnlyCoverage` is running the rule
+   * recursively to build the preview plan. Prevents the outer plan-only branch in `_apply` from
+   * firing on the nested call, so the recursive invocation executes the normal transform.
+   */
+  private[rules] val planOnlyPreviewInProgress: ThreadLocal[Boolean] =
+    new ThreadLocal[Boolean] {
+      override def initialValue(): Boolean = false
+    }
+
+  private[comet] def markPlanOnlyReported(executionId: Option[String]): Boolean = {
+    executionId match {
+      case None => true
+      case Some(id) =>
+        planOnlyReportedIds.synchronized {
+          val added = planOnlyReportedIds.add(id)
+          if (planOnlyReportedIds.size > PLAN_ONLY_REPORTED_LIMIT) {
+            val it = planOnlyReportedIds.iterator()
+            it.next()
+            it.remove()
+          }
+          added
+        }
+    }
+  }
+
+  private[comet] def clearPlanOnlyReported(): Unit = {
+    planOnlyReportedIds.synchronized {
+      planOnlyReportedIds.clear()
+    }
+  }
 }
 
 /**
@@ -573,6 +614,26 @@ case class CometExecRule(session: SparkSession)
     newPlan
   }
 
+  /**
+   * Build the Comet plan we would have executed and log it. Called from `_apply` in plan-only
+   * mode; the built plan is discarded. `CometScanRule` and this rule both skip work when
+   * `spark.comet.explain.planOnly.enabled` is set, so we use their bypass flags to force normal
+   * behavior while we compute the preview.
+   */
+  private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
+    val preview = CometScanRule.withForceApply {
+      CometExecRule.planOnlyPreviewInProgress.set(true)
+      try {
+        _apply(CometScanRule(session).apply(plan))
+      } finally {
+        CometExecRule.planOnlyPreviewInProgress.set(false)
+      }
+    }
+    logWarning(
+      s"[Comet plan-only]\n" +
+        new ExtendedExplainInfo().generateExtendedInfo(preview))
+  }
+
   private def _apply(plan: SparkPlan): SparkPlan = {
     // We shouldn't transform Spark query plan if Comet is not loaded.
     if (!isCometLoaded(conf)) return plan
@@ -605,6 +666,19 @@ case class CometExecRule(session: SparkSession)
       // during the bottom-up conversion. Tags persist through AQE stage creation.
       tagUnsafePartialAggregates(planWithJoinRewritten)
 
+      // Plan-only mode: `CometScanRule` skipped work, so `plan` here is still pure Spark. Build
+      // a preview by rerunning both rules against a scan-wrapped copy and log the coverage
+      // report, then return `plan` unchanged so nothing is offloaded to native.
+      if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() &&
+        !CometExecRule.planOnlyPreviewInProgress.get()) {
+        val executionId = Option(
+          session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
+        if (CometExecRule.markPlanOnlyReported(executionId)) {
+          reportPlanOnlyCoverage(plan)
+        }
+        return plan
+      }
+
       var newPlan = transform(planWithJoinRewritten)
 
       // if the plan cannot be run fully natively then explain why (when appropriate
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
index 2c15bb5e4e..df02250e6f 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
@@ -64,6 +64,12 @@ case class CometScanRule(session: SparkSession)
   private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
 
   override def apply(plan: SparkPlan): SparkPlan = {
+    // Plan-only mode: leave the plan untouched. CometExecRule invokes this rule with
+    // `CometScanRule.withForceApply` when it needs a preview plan to compute the coverage
+    // report; the thread-local flag bypasses this skip for that call only.
+    if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() && !CometScanRule.forceApply.get()) {
+      return plan
+    }
     val newPlan = _apply(plan)
     if (showTransformations && !newPlan.fastEquals(plan)) {
       logInfo(s"""
@@ -939,6 +945,22 @@ case class CometScanTypeChecker() extends DataTypeSupport with CometTypeShim {
 
 object CometScanRule extends Logging {
 
+  /**
+   * Thread-local flag telling `CometScanRule.apply` to run its normal wrapping logic even when
+   * `spark.comet.explain.planOnly.enabled` is set. Used by `CometExecRule` when it builds a
+   * preview plan to report on; see `CometExecRule.reportPlanOnlyCoverage`.
+   */
+  private[rules] val forceApply: ThreadLocal[Boolean] = new ThreadLocal[Boolean] {
+    override def initialValue(): Boolean = false
+  }
+
+  private[rules] def withForceApply[T](f: => T): T = {
+    val prev = forceApply.get()
+    forceApply.set(true)
+    try f
+    finally forceApply.set(prev)
+  }
+
   // Per-scheme memo of `NativeBase.isObjectStoreSchemeSupported`. The answer depends only on the
   // URL scheme, so we cache by scheme and never re-cross the JNI boundary for a repeated scheme.
   private val schemeSupportCache =
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index 5444a89fa3..3291e44958 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.QueryStageExec
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec}
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
 
 import org.apache.comet.{CometConf, CometExplainInfo}
@@ -797,4 +798,79 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
+  private def runPlanOnlyQuery(sql: String): SparkPlan = {
+    CometExecRule.clearPlanOnlyReported()
+    val df = spark.sql(sql)
+    val rows = df.collect()
+    // Sanity: results identical to plain Spark run with the config off.
+    withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "false") {
+      checkAnswer(df, spark.sql(sql).collect().toSeq)
+    }
+    assert(rows.nonEmpty)
+    df.queryExecution.executedPlan
+  }
+
+  private def assertNoComet(plan: SparkPlan): Unit = {
+    val cometNodes = stripAQEPlan(plan).collect { case p: CometPlan => p }
+    assert(
+      cometNodes.isEmpty,
+      s"plan-only mode must not offload; found Comet operators: $cometNodes")
+  }
+
+  for (aqeEnabled <- Seq(true, false)) {
+    test(s"plan-only mode: V1 scan, AQE=$aqeEnabled") {
+      withSQLConf(
+        SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+        withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+          assertNoComet(runPlanOnlyQuery("SELECT _2, count(*) FROM tbl GROUP BY _2"))
+        }
+      }
+    }
+
+    test(s"plan-only mode: V2 scan, AQE=$aqeEnabled") {
+      withSQLConf(
+        SQLConf.USE_V1_SOURCE_LIST.key -> "",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+        withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+          assertNoComet(runPlanOnlyQuery("SELECT _2, count(*) FROM tbl GROUP BY _2"))
+        }
+      }
+    }
+  }
+
+  test("plan-only mode: scalar subquery is also reverted") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        val plan = runPlanOnlyQuery("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)")
+        assertNoComet(plan)
+      }
+    }
+  }
+
+  test("plan-only mode: same query with the config off runs on Comet") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "false") {
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        val plan =
+          spark.sql("SELECT _2, count(*) FROM tbl GROUP BY _2").queryExecution.executedPlan
+        val cometNodes = stripAQEPlan(plan).collect { case p: CometPlan => p }
+        assert(cometNodes.nonEmpty, "expected Comet operators when plan-only mode is disabled")
+      }
+    }
+  }
+
 }

From 10e6bd607b443db5ade6670effe2bcde53aebab5 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 19 Aug 2026 06:26:51 -0600
Subject: [PATCH 2/8] refactor: simplify plan-only wiring
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Collapse two ThreadLocals into one shared `planOnlyPreviewInProgress`
  (dropped `CometScanRule.forceApply`/`withForceApply`).
- Move the plan-only branch to the top of the exec-enabled block so
  `normalizePlan`/`RewriteJoin`/`tagUnsafePartialAggregates` are not run
  and thrown away.
- Replace hand-rolled `LinkedHashSet` LRU with a synchronized
  `LinkedHashMap` + `removeEldestEntry`, matching
  `IcebergPlanDataInjector.commonCache`. Drops the `clearPlanOnlyReported`
  test helper.
- Collapse V1/V2 × AQE test variants into one loop; fold `assertNoComet`
  into `runPlanOnlyAndAssertReverted` and drop redundant `collect`s.
---
 .../apache/comet/rules/CometExecRule.scala    | 96 +++++++++----------
 .../apache/comet/rules/CometScanRule.scala    | 24 +----
 .../comet/rules/CometExecRuleSuite.scala      | 82 +++++++---------
 3 files changed, 80 insertions(+), 122 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 5d7aa95e13..36d84b435a 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -117,43 +117,39 @@ object CometExecRule {
     org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometBroadcast")
 
   /**
-   * Bounded set of SQL execution IDs for which plan-only mode has already logged its report. Used
-   * to dedupe the WARN under AQE, where `CometExecRule` fires once against the full initial plan
-   * and again per stage.
+   * SQL execution IDs for which plan-only mode has already logged its report, as a bounded LRU.
+   * Dedupes the WARN under AQE, where `CometExecRule` fires once against the full initial plan
+   * and again per stage. Same synchronized-`LinkedHashMap` LRU pattern used by
+   * `IcebergPlanDataInjector.commonCache`.
    */
   private val PLAN_ONLY_REPORTED_LIMIT = 1024
-  private val planOnlyReportedIds: java.util.LinkedHashSet[String] =
-    new java.util.LinkedHashSet[String]()
+  private val planOnlyReportedIds: java.util.Map[String, java.lang.Boolean] =
+    java.util.Collections.synchronizedMap(
+      new java.util.LinkedHashMap[String, java.lang.Boolean](16, 0.75f, false) {
+        override def removeEldestEntry(
+            eldest: java.util.Map.Entry[String, java.lang.Boolean]): Boolean =
+          size() > PLAN_ONLY_REPORTED_LIMIT
+      })
 
   /**
-   * Thread-local re-entry guard: set to true while `reportPlanOnlyCoverage` is running the rule
-   * recursively to build the preview plan. Prevents the outer plan-only branch in `_apply` from
-   * firing on the nested call, so the recursive invocation executes the normal transform.
+   * Thread-local re-entry guard: set while `reportPlanOnlyCoverage` is building the preview so
+   * both `CometScanRule` and `CometExecRule` run their normal transforms on the nested pass
+   * instead of short-circuiting.
    */
   private[rules] val planOnlyPreviewInProgress: ThreadLocal[Boolean] =
-    new ThreadLocal[Boolean] {
-      override def initialValue(): Boolean = false
-    }
+    ThreadLocal.withInitial(() => java.lang.Boolean.FALSE)
+
+  private[rules] def withPreview[T](f: => T): T = {
+    val prev = planOnlyPreviewInProgress.get()
+    planOnlyPreviewInProgress.set(true)
+    try f
+    finally planOnlyPreviewInProgress.set(prev)
+  }
 
   private[comet] def markPlanOnlyReported(executionId: Option[String]): Boolean = {
     executionId match {
       case None => true
-      case Some(id) =>
-        planOnlyReportedIds.synchronized {
-          val added = planOnlyReportedIds.add(id)
-          if (planOnlyReportedIds.size > PLAN_ONLY_REPORTED_LIMIT) {
-            val it = planOnlyReportedIds.iterator()
-            it.next()
-            it.remove()
-          }
-          added
-        }
-    }
-  }
-
-  private[comet] def clearPlanOnlyReported(): Unit = {
-    planOnlyReportedIds.synchronized {
-      planOnlyReportedIds.clear()
+      case Some(id) => planOnlyReportedIds.put(id, java.lang.Boolean.TRUE) == null
     }
   }
 }
@@ -616,22 +612,14 @@ case class CometExecRule(session: SparkSession)
 
   /**
    * Build the Comet plan we would have executed and log it. Called from `_apply` in plan-only
-   * mode; the built plan is discarded. `CometScanRule` and this rule both skip work when
-   * `spark.comet.explain.planOnly.enabled` is set, so we use their bypass flags to force normal
-   * behavior while we compute the preview.
+   * mode; the built plan is discarded. Both rules short-circuit under plan-only mode, so we set
+   * `planOnlyPreviewInProgress` to force the nested calls to run their normal transforms.
    */
   private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
-    val preview = CometScanRule.withForceApply {
-      CometExecRule.planOnlyPreviewInProgress.set(true)
-      try {
-        _apply(CometScanRule(session).apply(plan))
-      } finally {
-        CometExecRule.planOnlyPreviewInProgress.set(false)
-      }
+    val preview = CometExecRule.withPreview {
+      _apply(CometScanRule(session).apply(plan))
     }
-    logWarning(
-      s"[Comet plan-only]\n" +
-        new ExtendedExplainInfo().generateExtendedInfo(preview))
+    logWarning(s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}")
   }
 
   private def _apply(plan: SparkPlan): SparkPlan = {
@@ -650,6 +638,21 @@ case class CometExecRule(session: SparkSession)
         plan
       }
     } else {
+      // Plan-only mode: build the Comet plan Comet would have executed, log it, and return
+      // the original plan unchanged. `CometScanRule` also short-circuits in this mode, so
+      // `plan` is still pure Spark; `reportPlanOnlyCoverage` rebuilds a scan-wrapped copy for
+      // the preview. Placed before `normalizePlan`/`RewriteJoin`/`tagUnsafePartialAggregates`
+      // so their work is not wasted on the discarded outer pass.
+      if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() &&
+        !CometExecRule.planOnlyPreviewInProgress.get()) {
+        val executionId = Option(
+          session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
+        if (CometExecRule.markPlanOnlyReported(executionId)) {
+          reportPlanOnlyCoverage(plan)
+        }
+        return plan
+      }
+
       val normalizedPlan = normalizePlan(plan)
 
       val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) {
@@ -666,19 +669,6 @@ case class CometExecRule(session: SparkSession)
       // during the bottom-up conversion. Tags persist through AQE stage creation.
       tagUnsafePartialAggregates(planWithJoinRewritten)
 
-      // Plan-only mode: `CometScanRule` skipped work, so `plan` here is still pure Spark. Build
-      // a preview by rerunning both rules against a scan-wrapped copy and log the coverage
-      // report, then return `plan` unchanged so nothing is offloaded to native.
-      if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() &&
-        !CometExecRule.planOnlyPreviewInProgress.get()) {
-        val executionId = Option(
-          session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
-        if (CometExecRule.markPlanOnlyReported(executionId)) {
-          reportPlanOnlyCoverage(plan)
-        }
-        return plan
-      }
-
       var newPlan = transform(planWithJoinRewritten)
 
       // if the plan cannot be run fully natively then explain why (when appropriate
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
index df02250e6f..e7baa19cc3 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
@@ -64,10 +64,10 @@ case class CometScanRule(session: SparkSession)
   private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
 
   override def apply(plan: SparkPlan): SparkPlan = {
-    // Plan-only mode: leave the plan untouched. CometExecRule invokes this rule with
-    // `CometScanRule.withForceApply` when it needs a preview plan to compute the coverage
-    // report; the thread-local flag bypasses this skip for that call only.
-    if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() && !CometScanRule.forceApply.get()) {
+    // Plan-only mode: leave the plan untouched. `CometExecRule.reportPlanOnlyCoverage` sets
+    // `planOnlyPreviewInProgress` when it needs the wrapping to run for the preview plan.
+    if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() &&
+      !CometExecRule.planOnlyPreviewInProgress.get()) {
       return plan
     }
     val newPlan = _apply(plan)
@@ -945,22 +945,6 @@ case class CometScanTypeChecker() extends DataTypeSupport with CometTypeShim {
 
 object CometScanRule extends Logging {
 
-  /**
-   * Thread-local flag telling `CometScanRule.apply` to run its normal wrapping logic even when
-   * `spark.comet.explain.planOnly.enabled` is set. Used by `CometExecRule` when it builds a
-   * preview plan to report on; see `CometExecRule.reportPlanOnlyCoverage`.
-   */
-  private[rules] val forceApply: ThreadLocal[Boolean] = new ThreadLocal[Boolean] {
-    override def initialValue(): Boolean = false
-  }
-
-  private[rules] def withForceApply[T](f: => T): T = {
-    val prev = forceApply.get()
-    forceApply.set(true)
-    try f
-    finally forceApply.set(prev)
-  }
-
   // Per-scheme memo of `NativeBase.isObjectStoreSchemeSupported`. The answer depends only on the
   // URL scheme, so we cache by scheme and never re-cross the JNI boundary for a repeated scheme.
   private val schemeSupportCache =
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index 3291e44958..97296db046 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -798,63 +798,47 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
-  private def runPlanOnlyQuery(sql: String): SparkPlan = {
-    CometExecRule.clearPlanOnlyReported()
-    val df = spark.sql(sql)
-    val rows = df.collect()
-    // Sanity: results identical to plain Spark run with the config off.
-    withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "false") {
-      checkAnswer(df, spark.sql(sql).collect().toSeq)
+  /**
+   * Run `sql` with plan-only mode enabled and assert nothing was offloaded to native. `useV1`
+   * toggles between `USE_V1_SOURCE_LIST=parquet` (V1 `CometScanExec` path) and
+   * `USE_V1_SOURCE_LIST=""` (V2 `CometBatchScanExec` path).
+   */
+  private def runPlanOnlyAndAssertReverted(
+      sql: String,
+      useV1: Boolean = true,
+      aqe: Boolean = true): Unit = {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> (if (useV1) "parquet" else ""),
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+      val executed = spark.sql(sql).queryExecution.executedPlan
+      val cometNodes = stripAQEPlan(executed).collect { case p: CometPlan => p }
+      assert(
+        cometNodes.isEmpty,
+        s"plan-only mode must not offload; found Comet operators: $cometNodes")
     }
-    assert(rows.nonEmpty)
-    df.queryExecution.executedPlan
-  }
-
-  private def assertNoComet(plan: SparkPlan): Unit = {
-    val cometNodes = stripAQEPlan(plan).collect { case p: CometPlan => p }
-    assert(
-      cometNodes.isEmpty,
-      s"plan-only mode must not offload; found Comet operators: $cometNodes")
   }
 
-  for (aqeEnabled <- Seq(true, false)) {
-    test(s"plan-only mode: V1 scan, AQE=$aqeEnabled") {
-      withSQLConf(
-        SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
-        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
-        CometConf.COMET_ENABLED.key -> "true",
-        CometConf.COMET_EXEC_ENABLED.key -> "true",
-        CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
-        withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
-          assertNoComet(runPlanOnlyQuery("SELECT _2, count(*) FROM tbl GROUP BY _2"))
-        }
-      }
-    }
-
-    test(s"plan-only mode: V2 scan, AQE=$aqeEnabled") {
-      withSQLConf(
-        SQLConf.USE_V1_SOURCE_LIST.key -> "",
-        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
-        CometConf.COMET_ENABLED.key -> "true",
-        CometConf.COMET_EXEC_ENABLED.key -> "true",
-        CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
-        withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
-          assertNoComet(runPlanOnlyQuery("SELECT _2, count(*) FROM tbl GROUP BY _2"))
-        }
+  for {
+    useV1 <- Seq(true, false)
+    aqe <- Seq(true, false)
+  } {
+    val label = s"${if (useV1) "V1" else "V2"} scan, AQE=$aqe"
+    test(s"plan-only mode: $label") {
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        runPlanOnlyAndAssertReverted(
+          "SELECT _2, count(*) FROM tbl GROUP BY _2",
+          useV1 = useV1,
+          aqe = aqe)
       }
     }
   }
 
   test("plan-only mode: scalar subquery is also reverted") {
-    withSQLConf(
-      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
-      CometConf.COMET_ENABLED.key -> "true",
-      CometConf.COMET_EXEC_ENABLED.key -> "true",
-      CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
-      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
-        val plan = runPlanOnlyQuery("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)")
-        assertNoComet(plan)
-      }
+    withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+      runPlanOnlyAndAssertReverted("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)")
     }
   }
 

From eb514de1ff76785b81a4940d6a2adb3ea149ab07 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 19 Aug 2026 06:37:12 -0600
Subject: [PATCH 3/8] refactor: replace plan-only ThreadLocal with a forPreview
 parameter

`_apply` on both rules now takes/uses an explicit `forPreview` flag rather
than reading a thread-local. `reportPlanOnlyCoverage` calls the private
`_apply` methods directly, skipping the `Rule.apply` short-circuits, so
the recursive preview reads the parameter instead of ambient state.

Widens `CometScanRule._apply` to `private[rules]` so `CometExecRule` can
invoke it. Drops `planOnlyPreviewInProgress` and `withPreview`.
---
 .../apache/comet/rules/CometExecRule.scala    | 30 ++++---------------
 .../apache/comet/rules/CometScanRule.scala    | 10 +++----
 2 files changed, 11 insertions(+), 29 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 36d84b435a..9ee204ae16 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -131,21 +131,6 @@ object CometExecRule {
           size() > PLAN_ONLY_REPORTED_LIMIT
       })
 
-  /**
-   * Thread-local re-entry guard: set while `reportPlanOnlyCoverage` is building the preview so
-   * both `CometScanRule` and `CometExecRule` run their normal transforms on the nested pass
-   * instead of short-circuiting.
-   */
-  private[rules] val planOnlyPreviewInProgress: ThreadLocal[Boolean] =
-    ThreadLocal.withInitial(() => java.lang.Boolean.FALSE)
-
-  private[rules] def withPreview[T](f: => T): T = {
-    val prev = planOnlyPreviewInProgress.get()
-    planOnlyPreviewInProgress.set(true)
-    try f
-    finally planOnlyPreviewInProgress.set(prev)
-  }
-
   private[comet] def markPlanOnlyReported(executionId: Option[String]): Boolean = {
     executionId match {
       case None => true
@@ -600,7 +585,7 @@ case class CometExecRule(session: SparkSession)
   }
 
   override def apply(plan: SparkPlan): SparkPlan = {
-    val newPlan = _apply(plan)
+    val newPlan = _apply(plan, forPreview = false)
     if (showTransformations && !newPlan.fastEquals(plan)) {
       logInfo(s"""
            |=== Applying Rule $ruleName ===
@@ -612,17 +597,15 @@ case class CometExecRule(session: SparkSession)
 
   /**
    * Build the Comet plan we would have executed and log it. Called from `_apply` in plan-only
-   * mode; the built plan is discarded. Both rules short-circuit under plan-only mode, so we set
-   * `planOnlyPreviewInProgress` to force the nested calls to run their normal transforms.
+   * mode; the built plan is discarded. Passes `forPreview = true` through the nested calls so
+   * both rules run their normal transforms instead of short-circuiting.
    */
   private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
-    val preview = CometExecRule.withPreview {
-      _apply(CometScanRule(session).apply(plan))
-    }
+    val preview = _apply(CometScanRule(session)._apply(plan), forPreview = true)
     logWarning(s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}")
   }
 
-  private def _apply(plan: SparkPlan): SparkPlan = {
+  private def _apply(plan: SparkPlan, forPreview: Boolean): SparkPlan = {
     // We shouldn't transform Spark query plan if Comet is not loaded.
     if (!isCometLoaded(conf)) return plan
 
@@ -643,8 +626,7 @@ case class CometExecRule(session: SparkSession)
       // `plan` is still pure Spark; `reportPlanOnlyCoverage` rebuilds a scan-wrapped copy for
       // the preview. Placed before `normalizePlan`/`RewriteJoin`/`tagUnsafePartialAggregates`
       // so their work is not wasted on the discarded outer pass.
-      if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() &&
-        !CometExecRule.planOnlyPreviewInProgress.get()) {
+      if (!forPreview && CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) {
         val executionId = Option(
           session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
         if (CometExecRule.markPlanOnlyReported(executionId)) {
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
index e7baa19cc3..e852ddc1e4 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
@@ -64,10 +64,10 @@ case class CometScanRule(session: SparkSession)
   private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
 
   override def apply(plan: SparkPlan): SparkPlan = {
-    // Plan-only mode: leave the plan untouched. `CometExecRule.reportPlanOnlyCoverage` sets
-    // `planOnlyPreviewInProgress` when it needs the wrapping to run for the preview plan.
-    if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get() &&
-      !CometExecRule.planOnlyPreviewInProgress.get()) {
+    // Plan-only mode: leave the plan untouched. `CometExecRule.reportPlanOnlyCoverage` calls
+    // `_apply` directly to bypass this short-circuit when it needs the wrapping for the
+    // preview plan.
+    if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) {
       return plan
     }
     val newPlan = _apply(plan)
@@ -80,7 +80,7 @@ case class CometScanRule(session: SparkSession)
     newPlan
   }
 
-  private def _apply(plan: SparkPlan): SparkPlan = {
+  private[rules] def _apply(plan: SparkPlan): SparkPlan = {
     if (!isCometLoaded(conf)) return plan
 
     // Comet does not support structured streaming. The parallel guard in

From 2f94707f8fc087b2a38928c9845e7dd0a395b2e4 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 19 Aug 2026 11:14:39 -0600
Subject: [PATCH 4/8] Address review: report the outer query and include
 post-columnar rules

Plan-only reporting keyed its once-per-query slot on the SQL execution ID
alone. Spark prepares scalar and DPP subqueries as their own top-level plans
before the outer plan reaches the conversion rules, so a nested subquery took
the slot and the outer query - the plan being evaluated - was never reported.
Reports are now keyed on the execution ID and the plan, so the outer query
always gets one, while AQE's per-stage and per-re-optimization applications of
the rule are recognised as re-plans and stay quiet.

The preview also stopped at operator conversion, before Spark inserts columnar
transitions and Comet runs its post-columnar rules. It therefore counted
operators that RevertNativeForTransitionHeavyStages removes and counted
transitions before they existed. The preview now inserts transitions and
applies both post-columnar rules, using a new applyToAllStages entry point so
that every shuffle boundary is visited rather than only the topmost stage.

Tests capture the logged report and assert that the outer query appears under
both AQE modes, and that the reported coverage matches the coverage of the
plan Comet really executes when stage reversion fires.
---
 .../latest/understanding-comet-plans.md       |  20 +++-
 .../comet/CometSparkSessionExtensions.scala   |   6 +-
 .../apache/comet/rules/CometExecRule.scala    | 109 +++++++++++++++---
 ...RevertNativeForTransitionHeavyStages.scala |  13 +++
 .../comet/rules/CometExecRuleSuite.scala      | 104 ++++++++++++++++-
 5 files changed, 229 insertions(+), 23 deletions(-)

diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md
index a12c3f9f60..363d5d1f78 100644
--- a/docs/source/user-guide/latest/understanding-comet-plans.md
+++ b/docs/source/user-guide/latest/understanding-comet-plans.md
@@ -218,14 +218,30 @@ execution.
 
 The log line is prefixed with `[Comet plan-only]` and includes the same
 annotated plan and summary as `spark.comet.explain.format=verbose` produces
-against a normal Comet plan. Under AQE the report is emitted once per SQL
-execution.
+against a normal Comet plan. The preview goes through the whole Comet planning
+sequence, not just operator conversion: Spark's columnar transitions are
+inserted and Comet's post-columnar rules
+(`RevertNativeForTransitionHeavyStages`, `EliminateRedundantTransitions`) are
+applied, so a stage that Comet would have reverted to Spark for having too many
+transitions is reported as reverted.
+
+Spark prepares some plans on their own, ahead of the query that contains them —
+scalar subqueries and dynamic partition pruning subqueries, for instance — so a
+query gets one report per independently planned plan: one for the outer query,
+plus one per such subquery. Repeat applications of the same plan are not
+reported again: under AQE, neither the per-stage applications nor the
+applications that follow each adaptive re-optimization add reports.
 
 The estimate reflects Scala-side conversion only. The native plan is never
 handed to DataFusion, so anything that would have failed in DataFusion's
 `create_plan` still counts as accelerated. Treat the percentage as an upper
 bound.
 
+Under AQE there is a second reason to treat the report as an estimate: it
+describes the plan as it stands before any adaptive re-planning, and the
+post-columnar rules are applied to that whole plan at once rather than to each
+stage as it is created. Coverage of the plan AQE finally executes can differ.
+
 The config requires `spark.comet.exec.enabled=true`. With Comet exec disabled
 the rule that emits the report does not run.
 
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 4dc3619433..22a450d0a0 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -97,7 +97,9 @@ class CometSparkSessionExtensions
     // No-op on Spark 3.5+; see CometSpark34AqeDppFallbackRule's class docstring.
     injectPreSpark35QueryStagePrepRuleShim(extensions, CometSpark34AqeDppFallbackRule)
     extensions.injectQueryStagePrepRule { session => CometScanRule(session) }
-    extensions.injectQueryStagePrepRule { session => CometExecRule(session) }
+    extensions.injectQueryStagePrepRule { session =>
+      CometExecRule(session, queryStagePrep = true)
+    }
     injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters)
     injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery)
     extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) }
@@ -111,6 +113,8 @@ class CometSparkSessionExtensions
     override def preColumnarTransitions: Rule[SparkPlan] = CometExecRule(session)
 
     override def postColumnarTransitions: Rule[SparkPlan] = {
+      // Keep in sync with `CometExecRule.reportPlanOnlyCoverage`, which replays these rules over
+      // the plan it previews so that plan-only reports describe the plan that would have run.
       val rules =
         Seq(RevertNativeForTransitionHeavyStages(session), EliminateRedundantTransitions(session))
       plan => rules.foldLeft(plan) { case (p, rule) => rule(p) }
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 9ee204ae16..75db4b40e6 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -32,7 +32,7 @@ import org.apache.spark.sql.comet._
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution._
-import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
 import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec}
 import org.apache.spark.sql.execution.datasources.WriteFilesExec
@@ -117,32 +117,89 @@ object CometExecRule {
     org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometBroadcast")
 
   /**
-   * SQL execution IDs for which plan-only mode has already logged its report, as a bounded LRU.
-   * Dedupes the WARN under AQE, where `CometExecRule` fires once against the full initial plan
-   * and again per stage. Same synchronized-`LinkedHashMap` LRU pattern used by
-   * `IcebergPlanDataInjector.commonCache`.
+   * A bounded set of keys, used for plan-only reporting state. Evicts in LRU order once `limit`
+   * keys are held, so a long-lived driver retains a fixed amount of reporting state. Same
+   * synchronized-`LinkedHashMap` pattern used by `IcebergPlanDataInjector.commonCache`.
    */
+  private class BoundedKeySet(limit: Int) {
+    private val keys: java.util.Map[String, java.lang.Boolean] =
+      java.util.Collections.synchronizedMap(
+        new java.util.LinkedHashMap[String, java.lang.Boolean](16, 0.75f, true) {
+          override def removeEldestEntry(
+              eldest: java.util.Map.Entry[String, java.lang.Boolean]): Boolean = size() > limit
+        })
+
+    /** Adds `key`, returning true if it was not already present. */
+    def add(key: String): Boolean = keys.put(key, java.lang.Boolean.TRUE) == null
+
+    def contains(key: String): Boolean = keys.containsKey(key)
+  }
+
   private val PLAN_ONLY_REPORTED_LIMIT = 1024
-  private val planOnlyReportedIds: java.util.Map[String, java.lang.Boolean] =
-    java.util.Collections.synchronizedMap(
-      new java.util.LinkedHashMap[String, java.lang.Boolean](16, 0.75f, false) {
-        override def removeEldestEntry(
-            eldest: java.util.Map.Entry[String, java.lang.Boolean]): Boolean =
-          size() > PLAN_ONLY_REPORTED_LIMIT
-      })
-
-  private[comet] def markPlanOnlyReported(executionId: Option[String]): Boolean = {
+
+  /** `executionId:planFingerprint` keys that plan-only mode has already reported. */
+  private val planOnlyReportedPlans = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
+  /** Execution IDs whose plan-only report came from the query-stage-prep rule. */
+  private val planOnlyPrepReportedIds = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
+  /**
+   * Whether plan-only mode should report `plan`, recording that it did so.
+   *
+   * Spark applies this rule many times during one SQL execution, and only some of those
+   * applications correspond to a plan the user is asking about:
+   *
+   *   - Each scalar subquery and DPP subquery is prepared as its own top-level plan, and that
+   *     happens *before* the outer plan reaches the conversion rules. Keying the report on the
+   *     execution ID alone therefore let a nested subquery consume the slot and suppressed the
+   *     outer plan, which is the plan being evaluated. Keying on the execution ID *and* the plan
+   *     gives the outer plan its own report and each separately prepared subquery theirs.
+   *   - Under AQE the rule also runs once per query stage (as a columnar rule) and again on every
+   *     re-optimization (as a query-stage-prep rule). Those are re-planning of a plan already
+   *     reported, so a `plan` containing query stages is skipped, and once the query-stage-prep
+   *     rule has reported an execution the columnar applications for it stay quiet.
+   *
+   * @param queryStagePrep
+   *   whether the calling rule instance is registered as a query-stage-prep rule.
+   */
+  private[comet] def shouldReportPlanOnly(
+      executionId: Option[String],
+      plan: SparkPlan,
+      queryStagePrep: Boolean): Boolean = {
     executionId match {
-      case None => true
-      case Some(id) => planOnlyReportedIds.put(id, java.lang.Boolean.TRUE) == null
+      case None =>
+        // No execution ID means the plan is being built outside an action (`df.explain`, or
+        // reading `queryExecution.executedPlan` directly). AQE creates no stages in that case, so
+        // there is nothing to dedupe against.
+        true
+      case Some(id) =>
+        if (plan.exists(_.isInstanceOf[QueryStageExec])) {
+          // An AQE stage plan or a re-optimized plan: a re-plan of what we already reported.
+          false
+        } else {
+          if (queryStagePrep) {
+            planOnlyPrepReportedIds.add(id)
+          } else if (planOnlyPrepReportedIds.contains(id)) {
+            return false
+          }
+          // The plan's structural hash identifies it: node tags are not part of it, so the same
+          // plan applied twice keys the same and is reported once.
+          planOnlyReportedPlans.add(s"$id:${plan.hashCode()}")
+        }
     }
   }
 }
 
 /**
  * Spark physical optimizer rule for replacing Spark operators with Comet operators.
+ *
+ * @param queryStagePrep
+ *   true for the instance registered with `injectQueryStagePrepRule`, which under AQE sees the
+ *   whole initial plan, and false for the one registered as a columnar rule, which under AQE sees
+ *   one query stage at a time. Only plan-only reporting reads this; see
+ *   [[CometExecRule.shouldReportPlanOnly]].
  */
-case class CometExecRule(session: SparkSession)
+case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
     extends Rule[SparkPlan]
     with ShimSubqueryBroadcast {
 
@@ -599,9 +656,23 @@ case class CometExecRule(session: SparkSession)
    * Build the Comet plan we would have executed and log it. Called from `_apply` in plan-only
    * mode; the built plan is discarded. Passes `forPreview = true` through the nested calls so
    * both rules run their normal transforms instead of short-circuiting.
+   *
+   * Conversion is only the first half of Comet planning. Normally Spark then inserts the columnar
+   * transitions and runs Comet's post-columnar rules (see
+   * `CometSparkSessionExtensions.CometExecColumnar.postColumnarTransitions`), which can revert
+   * whole stages back to Spark and drop redundant transitions. Those steps run here too, so the
+   * report describes the plan that would really have executed and counts the transitions that
+   * would really have been there, rather than the pre-transition conversion result.
+   *
+   * `RevertNativeForTransitionHeavyStages` is applied with `applyToAllStages` because the preview
+   * holds the whole plan at once, whereas under AQE Spark hands that rule one stage at a time.
    */
   private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
-    val preview = _apply(CometScanRule(session)._apply(plan), forPreview = true)
+    val converted = _apply(CometScanRule(session)._apply(plan), forPreview = true)
+    val withTransitions =
+      ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted)
+    val reverted = RevertNativeForTransitionHeavyStages(session).applyToAllStages(withTransitions)
+    val preview = EliminateRedundantTransitions(session).apply(reverted)
     logWarning(s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}")
   }
 
@@ -629,7 +700,7 @@ case class CometExecRule(session: SparkSession)
       if (!forPreview && CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) {
         val executionId = Option(
           session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
-        if (CometExecRule.markPlanOnlyReported(executionId)) {
+        if (CometExecRule.shouldReportPlanOnly(executionId, plan, queryStagePrep)) {
           reportPlanOnlyCoverage(plan)
         }
         return plan
diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
index 1e0cfc79e0..88783c6572 100644
--- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
@@ -76,6 +76,19 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession)
       .getOrElse(withRevertedStages)
   }
 
+  /**
+   * Applies the revert decision to every stage of `plan`, regardless of whether AQE is enabled.
+   *
+   * `apply` picks the AQE branch when AQE is on because Spark hands it a single query stage at a
+   * time there, so only the topmost stage of `plan` is considered. Callers holding a whole plan
+   * that has not been split into stages - the plan-only preview in
+   * `CometExecRule.reportPlanOnlyCoverage` - need every shuffle boundary visited to see the
+   * reversions that the real per-stage applications would make.
+   */
+  private[rules] def applyToAllStages(plan: SparkPlan): SparkPlan = {
+    if (!enabled) plan else applyForNonAQE(plan)
+  }
+
   /**
    * Reverts the stage if C2R count exceeds threshold. Wraps in R2C if exchange needs columnar.
    */
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index 97296db046..f769721efb 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -21,6 +21,7 @@ package org.apache.comet.rules
 
 import scala.util.Random
 
+import org.apache.logging.log4j.Level
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.FunctionIdentifier
 import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo}
@@ -34,7 +35,7 @@ import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleEx
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
 
-import org.apache.comet.{CometConf, CometExplainInfo}
+import org.apache.comet.{CometConf, CometCoverageStats, CometExplainInfo}
 import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus}
 import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator}
 
@@ -857,4 +858,105 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
+  private val PLAN_ONLY_PREFIX = "[Comet plan-only]"
+
+  /** Runs `f` and returns the `[Comet plan-only]` reports that `CometExecRule` logged. */
+  private def capturePlanOnlyReports(f: => Unit): Seq[String] = {
+    val appender = new LogAppender("Comet plan-only reports")
+    withLogAppender(
+      appender,
+      loggerNames = Seq(classOf[CometExecRule].getName),
+      level = Some(Level.WARN)) {
+      f
+    }
+    appender.loggingEvents
+      .map(_.getMessage.getFormattedMessage)
+      .filter(_.startsWith(PLAN_ONLY_PREFIX))
+      .toSeq
+  }
+
+  /** The `Comet accelerated N out of M eligible operators` counts in a plan-only report. */
+  private def coverageOf(report: String): (Int, Int) = {
+    val pattern = """Comet accelerated (\d+) out of (\d+) eligible operators""".r
+    pattern
+      .findFirstMatchIn(report)
+      .map(m => (m.group(1).toInt, m.group(2).toInt))
+      .getOrElse(fail(s"report has no coverage summary:\n$report"))
+  }
+
+  // The outer query is planned after any subquery it contains, so a report slot owned by the
+  // first plan Spark prepares would describe the subquery and never the query being evaluated.
+  for (aqe <- Seq(true, false)) {
+    test(s"plan-only mode: report describes the outer query, not just a subquery (AQE=$aqe)") {
+      withSQLConf(
+        SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+        withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+          val reports = capturePlanOnlyReports {
+            spark.sql("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)").collect()
+          }
+          assert(reports.nonEmpty, "expected a plan-only report")
+          // The outer plan's Filter appears in no subquery plan, so its presence proves the
+          // outer query was reported and not suppressed by the subquery's earlier planning.
+          assert(
+            reports.exists(_.contains("Filter")),
+            s"no report describes the outer query:\n${reports.mkString("\n\n")}")
+          assert(
+            reports.distinct.size == reports.size,
+            s"the same plan was reported more than once:\n${reports.mkString("\n\n")}")
+          // Expected: one report for the subquery plan, one for the outer plan. AQE applies the
+          // rule again per stage and per re-optimization; those must not add reports.
+          assert(
+            reports.size <= 4,
+            s"expected a report per planned plan, got ${reports.size}:\n" +
+              reports.mkString("\n\n"))
+        }
+      }
+    }
+  }
+
+  test("plan-only mode: coverage accounts for post-columnar stage reversion") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      // AQE off so that Spark applies the post-columnar rules to the whole plan exactly once,
+      // which is what the preview does, making the two directly comparable.
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false") {
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        val query = "SELECT _2, count(*), sum(_1) FROM tbl GROUP BY _2"
+
+        // Comet accelerates part of this plan when the stage is left alone, so a preview that
+        // stopped before the post-columnar rules would report a non-zero count below.
+        withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") {
+          val df = sql(query)
+          df.collect()
+          assert(CometCoverageStats.forPlan(df.queryExecution.executedPlan).cometOperators > 0)
+        }
+
+        withSQLConf(
+          CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true",
+          CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") {
+          // What Comet really executes with reversion enabled.
+          val df = sql(query)
+          df.collect()
+          val executed = CometCoverageStats.forPlan(df.queryExecution.executedPlan)
+
+          val reports = withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+            capturePlanOnlyReports(sql(query).collect())
+          }
+          assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}")
+          assert(
+            coverageOf(reports.head) ==
+              (executed.cometOperators, executed.cometOperators + executed.sparkOperators),
+            s"report disagrees with the executed plan ($executed):\n${reports.head}")
+        }
+      }
+    }
+  }
+
 }

From 1683a65a446403f691f3fbad030f2706c533093d Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 26 Aug 2026 15:19:38 -0600
Subject: [PATCH 5/8] Address review: fix Spark 3.x test compile, dedupe
 without an execution ID, preview subqueries

Move the coverage assertions inside their withSQLConf block. Spark 3.4 and 3.5
declare withSQLConf as returning Unit, so binding its result inferred Unit and
broke test compilation on every supported Spark 3.x build.

Replace the execution-ID keyed report slots with a mark on the plan itself.
df.rdd.count() and reading executedPlan without an action can plan, and in the
first case execute AQE stages, with no execution ID installed, and the old no-ID
path reported unconditionally. Catalyst copies tags onto replacement nodes, so
the mark survives the rewrites between one rule application and the next and
identifies an AQE stage or final-plan pass as a re-plan of something already
reported. The execution-scoped hash is kept for the one case a mark cannot
cover: a subquery referenced twice is prepared twice, as separate but identical
plans.

Preview the plan behind each subquery expression before computing coverage.
Extended explain counts the plans owned by a node's expressions, and normal
planning has already converted those by the time the outer plan arrives, so the
outer report counted accelerated subquery operators as Spark. Two things were
needed to make the numbers line up: a prepared subquery arrives with codegen
wrappers and transitions already inserted, which blocks conversion, so those are
stripped first; and ReuseExchangeAndSubquery is replayed at the top level, since
it is the last preparation step and otherwise each copy of a subquery is counted
separately. With AQE off the outer report now equals CometCoverageStats for the
plan Comet really executes.

Under AQE the outer report still describes the pre-adaptive plan, in which
subqueries have not been planned yet and so count as Spark. Called out in the
user guide.
---
 .../latest/understanding-comet-plans.md       |  23 ++-
 .../apache/comet/rules/CometExecRule.scala    | 184 ++++++++++++++----
 .../comet/rules/CometExecRuleSuite.scala      |  84 +++++++-
 3 files changed, 237 insertions(+), 54 deletions(-)

diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md
index 363d5d1f78..c2e09b35f4 100644
--- a/docs/source/user-guide/latest/understanding-comet-plans.md
+++ b/docs/source/user-guide/latest/understanding-comet-plans.md
@@ -228,19 +228,28 @@ transitions is reported as reverted.
 Spark prepares some plans on their own, ahead of the query that contains them —
 scalar subqueries and dynamic partition pruning subqueries, for instance — so a
 query gets one report per independently planned plan: one for the outer query,
-plus one per such subquery. Repeat applications of the same plan are not
-reported again: under AQE, neither the per-stage applications nor the
-applications that follow each adaptive re-optimization add reports.
+plus one per such subquery. The outer report counts its subqueries too, the same
+way normal Comet planning does, so the reports for one query describe
+overlapping sets of operators and their counts should not be added up. Repeat
+applications of the same plan are not reported again: under AQE, neither the
+per-stage applications nor the applications that follow each adaptive
+re-optimization add reports.
 
 The estimate reflects Scala-side conversion only. The native plan is never
 handed to DataFusion, so anything that would have failed in DataFusion's
 `create_plan` still counts as accelerated. Treat the percentage as an upper
 bound.
 
-Under AQE there is a second reason to treat the report as an estimate: it
-describes the plan as it stands before any adaptive re-planning, and the
-post-columnar rules are applied to that whole plan at once rather than to each
-stage as it is created. Coverage of the plan AQE finally executes can differ.
+Under AQE the report is an estimate for a second reason: it describes the plan
+as it stands before any adaptive re-planning, and the post-columnar rules are
+applied to that whole plan at once rather than to each stage as it is created.
+Coverage of the plan AQE finally executes can differ. One case is worth calling
+out, because it moves the number the other way: AQE does not plan a subquery
+into the outer plan until after the report has been produced, so the outer
+report counts a subquery's operators as un-accelerated Spark even where Comet
+would accelerate them. For a subquery-heavy query under AQE, read the
+per-subquery reports rather than the outer percentage, or turn AQE off for the
+evaluation run.
 
 The config requires `spark.comet.exec.enabled=true`. With Comet exec disabled
 the rule that emits the report does not run.
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 75db4b40e6..581bf58eaa 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -43,8 +43,9 @@ import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, V2CommandEx
 import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan
 import org.apache.spark.sql.execution.datasources.v2.json.JsonScan
 import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
-import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange, ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery
 import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
@@ -131,8 +132,6 @@ object CometExecRule {
 
     /** Adds `key`, returning true if it was not already present. */
     def add(key: String): Boolean = keys.put(key, java.lang.Boolean.TRUE) == null
-
-    def contains(key: String): Boolean = keys.containsKey(key)
   }
 
   private val PLAN_ONLY_REPORTED_LIMIT = 1024
@@ -140,24 +139,70 @@ object CometExecRule {
   /** `executionId:planFingerprint` keys that plan-only mode has already reported. */
   private val planOnlyReportedPlans = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
 
-  /** Execution IDs whose plan-only report came from the query-stage-prep rule. */
-  private val planOnlyPrepReportedIds = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+  /**
+   * Set on the root of every plan that plan-only mode has reported. Catalyst copies a node's tags
+   * onto the node that replaces it (`TreeNode.copyTagsFrom`), so the mark survives the rewrites
+   * Spark applies between one application of this rule and the next, which is what lets a later
+   * application recognize a plan it has already described.
+   */
+  private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported")
+
+  /**
+   * Whether `plan` is an application of this rule to a plan already reported, rather than to a
+   * plan the user is asking about. All the state lives on the plan itself, so this holds whether
+   * or not the application carries a SQL execution ID - `df.rdd.count()` and reading
+   * `executedPlan` without an action both plan, and in the first case execute AQE stages, with no
+   * execution ID set.
+   *
+   * Under AQE one query reaches the conversion rules five times, and only the first is the plan
+   * being evaluated:
+   *
+   *   - the initial plan, through the query-stage-prep rule - the one to report;
+   *   - the same plan again as a columnar rule, now wrapped in `AdaptiveSparkPlanExec`;
+   *   - each query stage as it is created, again as a columnar rule, rooted at the `Exchange` the
+   *     stage was cut at;
+   *   - the re-optimized plan after each stage materializes, through the prep rule;
+   *   - the final plan once every stage has materialized, as a columnar rule.
+   *
+   * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the wrapper is matched
+   * by name rather than through its contents. The two re-optimization passes hold
+   * `QueryStageExec` nodes for the stages already materialized. The remaining two - a stage, and
+   * a final plan for a query AQE never had to cut into stages - carry the mark left by the prep
+   * rule.
+   *
+   * @param queryStagePrep
+   *   whether the calling rule instance is registered as a query-stage-prep rule.
+   * @param aqeEnabled
+   *   whether AQE is enabled for the session. A plan rooted at an `Exchange` is a stage only when
+   *   AQE cuts stages; with AQE off it is an ordinary plan (`df.repartition(n)`, say) and must
+   *   still be reported.
+   */
+  private def isReapplication(
+      plan: SparkPlan,
+      queryStagePrep: Boolean,
+      aqeEnabled: Boolean): Boolean = {
+    plan.exists(p =>
+      p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] ||
+        p.getTagValue(PLAN_ONLY_REPORTED).isDefined) ||
+    (aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange])
+  }
 
   /**
-   * Whether plan-only mode should report `plan`, recording that it did so.
+   * Whether plan-only mode should report `plan`, marking it reported if so.
    *
-   * Spark applies this rule many times during one SQL execution, and only some of those
+   * Spark applies this rule many times while executing one query, and only some of those
    * applications correspond to a plan the user is asking about:
    *
    *   - Each scalar subquery and DPP subquery is prepared as its own top-level plan, and that
-   *     happens *before* the outer plan reaches the conversion rules. Keying the report on the
-   *     execution ID alone therefore let a nested subquery consume the slot and suppressed the
-   *     outer plan, which is the plan being evaluated. Keying on the execution ID *and* the plan
-   *     gives the outer plan its own report and each separately prepared subquery theirs.
-   *   - Under AQE the rule also runs once per query stage (as a columnar rule) and again on every
-   *     re-optimization (as a query-stage-prep rule). Those are re-planning of a plan already
-   *     reported, so a `plan` containing query stages is skipped, and once the query-stage-prep
-   *     rule has reported an execution the columnar applications for it stay quiet.
+   *     happens *before* the outer plan reaches the conversion rules. A single report slot per
+   *     SQL execution therefore let a nested subquery consume the slot and suppressed the outer
+   *     plan, which is the plan being evaluated. Marking plans individually gives the outer plan
+   *     its own report and each separately prepared subquery theirs.
+   *   - Under AQE the rule also runs per query stage and again after every re-optimization. Those
+   *     are re-planning of a plan already reported; see [[isReapplication]].
+   *   - A subquery referenced from more than one place in the outer plan is prepared once per
+   *     reference, as a separate but identical plan each time. The mark cannot catch those - they
+   *     share no nodes - so within one SQL execution the plan's structural hash dedupes them.
    *
    * @param queryStagePrep
    *   whether the calling rule instance is registered as a query-stage-prep rule.
@@ -165,27 +210,16 @@ object CometExecRule {
   private[comet] def shouldReportPlanOnly(
       executionId: Option[String],
       plan: SparkPlan,
-      queryStagePrep: Boolean): Boolean = {
-    executionId match {
-      case None =>
-        // No execution ID means the plan is being built outside an action (`df.explain`, or
-        // reading `queryExecution.executedPlan` directly). AQE creates no stages in that case, so
-        // there is nothing to dedupe against.
-        true
-      case Some(id) =>
-        if (plan.exists(_.isInstanceOf[QueryStageExec])) {
-          // An AQE stage plan or a re-optimized plan: a re-plan of what we already reported.
-          false
-        } else {
-          if (queryStagePrep) {
-            planOnlyPrepReportedIds.add(id)
-          } else if (planOnlyPrepReportedIds.contains(id)) {
-            return false
-          }
-          // The plan's structural hash identifies it: node tags are not part of it, so the same
-          // plan applied twice keys the same and is reported once.
-          planOnlyReportedPlans.add(s"$id:${plan.hashCode()}")
-        }
+      queryStagePrep: Boolean,
+      aqeEnabled: Boolean): Boolean = {
+    if (isReapplication(plan, queryStagePrep, aqeEnabled)) {
+      false
+    } else {
+      plan.setTagValue(PLAN_ONLY_REPORTED, ())
+      // Node tags are not part of a plan's structural hash, so the mark set above does not
+      // perturb the key. Without an execution ID there is nothing to scope the state to, and
+      // `isReapplication` has already ruled out the repeat applications AQE makes, so report.
+      executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.hashCode()}"))
     }
   }
 }
@@ -654,8 +688,16 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
 
   /**
    * Build the Comet plan we would have executed and log it. Called from `_apply` in plan-only
-   * mode; the built plan is discarded. Passes `forPreview = true` through the nested calls so
-   * both rules run their normal transforms instead of short-circuiting.
+   * mode; the built plan is discarded.
+   */
+  private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
+    val preview = buildPreview(plan, topLevel = true)
+    logWarning(s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}")
+  }
+
+  /**
+   * The plan Comet would have executed for `plan`. Passes `forPreview = true` through the nested
+   * calls so both conversion rules run their normal transforms instead of short-circuiting.
    *
    * Conversion is only the first half of Comet planning. Normally Spark then inserts the columnar
    * transitions and runs Comet's post-columnar rules (see
@@ -666,14 +708,68 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
    *
    * `RevertNativeForTransitionHeavyStages` is applied with `applyToAllStages` because the preview
    * holds the whole plan at once, whereas under AQE Spark hands that rule one stage at a time.
+   *
+   * @param topLevel
+   *   false when previewing the plan behind a subquery expression. `ReuseExchangeAndSubquery` is
+   *   the last step of Spark's preparation and `QueryExecution.preparations` omits it for a
+   *   subquery, so the preview follows suit.
    */
-  private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
-    val converted = _apply(CometScanRule(session)._apply(plan), forPreview = true)
+  private def buildPreview(plan: SparkPlan, topLevel: Boolean): SparkPlan = {
+    val converted =
+      _apply(CometScanRule(session)._apply(previewSubqueriesOf(plan)), forPreview = true)
     val withTransitions =
       ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted)
     val reverted = RevertNativeForTransitionHeavyStages(session).applyToAllStages(withTransitions)
     val preview = EliminateRedundantTransitions(session).apply(reverted)
-    logWarning(s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}")
+    if (topLevel) ReuseExchangeAndSubquery.apply(preview) else preview
+  }
+
+  /**
+   * `plan` with the plan behind each of its subquery expressions replaced by that plan's own
+   * preview.
+   *
+   * Extended explain walks a node's `innerChildren`, which for a `SparkPlan` are the plans owned
+   * by its expressions, and counts their operators towards the report. Normal planning has
+   * already converted those plans by the time the outer plan reaches this rule - Spark prepares a
+   * scalar subquery through the full preparation sequence, columnar rules included, before
+   * substituting it into the outer plan - so leaving them untouched here would report every
+   * subquery operator as un-accelerated Spark and understate coverage relative to what Comet
+   * really executes.
+   *
+   * Each subquery is also reported in its own right, because Spark prepares it as a top-level
+   * plan of its own; those reports and the counts here therefore describe overlapping sets of
+   * operators.
+   */
+  private def previewSubqueriesOf(plan: SparkPlan): SparkPlan = {
+    plan.transformAllExpressions { case subquery: ExecSubqueryExpression =>
+      subquery.withNewPlan(previewSubquery(subquery.plan))
+    }
+  }
+
+  private def previewSubquery(subquery: BaseSubqueryExec): BaseSubqueryExec = subquery match {
+    // Reuse bookkeeping: the plan to preview is one level further down.
+    case reused: ReusedSubqueryExec => reused.copy(child = previewSubquery(reused.child))
+    case other =>
+      val preview = buildPreview(stripPreparation(other.child), topLevel = false)
+      other.withNewChildren(Seq(preview)).asInstanceOf[BaseSubqueryExec]
+  }
+
+  /**
+   * `plan` with the artifacts of a finished plan preparation removed: whole-stage codegen
+   * wrappers and the columnar transitions Spark inserted.
+   *
+   * A subquery arrives inside the outer plan fully prepared -
+   * `ApplyColumnarRulesAndInsertTransitions` and `CollapseCodegenStages` have both run over it -
+   * whereas the conversion rules only ever see a plan midway through preparation. A
+   * `HashAggregateExec` still wrapped in `WholeStageCodegenExec` is left unconverted, so
+   * previewing the prepared form would report a subquery as falling back that Comet in fact
+   * accelerates. [[buildPreview]] re-inserts the transitions once conversion is done.
+   */
+  private def stripPreparation(plan: SparkPlan): SparkPlan = plan.transformUp {
+    case WholeStageCodegenExec(child) => child
+    case InputAdapter(child) => child
+    case ColumnarToRowExec(child) => child
+    case RowToColumnarExec(child) => child
   }
 
   private def _apply(plan: SparkPlan, forPreview: Boolean): SparkPlan = {
@@ -700,7 +796,11 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
       if (!forPreview && CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) {
         val executionId = Option(
           session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
-        if (CometExecRule.shouldReportPlanOnly(executionId, plan, queryStagePrep)) {
+        if (CometExecRule.shouldReportPlanOnly(
+            executionId,
+            plan,
+            queryStagePrep,
+            conf.adaptiveExecutionEnabled)) {
           reportPlanOnlyCoverage(plan)
         }
         return plan
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index f769721efb..0ebe0e9e4c 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -946,14 +946,88 @@ class CometExecRuleSuite extends CometTestBase {
           df.collect()
           val executed = CometCoverageStats.forPlan(df.queryExecution.executedPlan)
 
-          val reports = withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
-            capturePlanOnlyReports(sql(query).collect())
+          // The assertions stay inside the config block: on Spark 3.4 and 3.5 `withSQLConf` is
+          // declared to return `Unit`, so a value cannot be carried out of one.
+          withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+            val reports = capturePlanOnlyReports(sql(query).collect())
+            assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}")
+            assert(
+              coverageOf(reports.head) ==
+                (executed.cometOperators, executed.cometOperators + executed.sparkOperators),
+              s"report disagrees with the executed plan ($executed):\n${reports.head}")
           }
-          assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}")
+        }
+      }
+    }
+  }
+
+  // `df.rdd.count()` and reading `executedPlan` without running an action can plan - and in the
+  // first case execute AQE stages - without a SQL execution ID installed, so the reporting state
+  // cannot be scoped to one. Each of these must still produce exactly one report per query, not
+  // one per stage and re-optimization.
+  for (aqe <- Seq(true, false)) {
+    test(s"plan-only mode: one report per query without a SQL execution ID (AQE=$aqe)") {
+      withSQLConf(
+        SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+        withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+          val query = "SELECT _2, count(*) FROM tbl GROUP BY _2"
+
+          val viaRdd = capturePlanOnlyReports(spark.sql(query).rdd.count())
+          assert(
+            viaRdd.size == 1,
+            s"expected one report for df.rdd.count(), got ${viaRdd.size}:\n" +
+              viaRdd.mkString("\n\n"))
+
+          val viaExecutedPlan =
+            capturePlanOnlyReports(spark.sql(query).queryExecution.executedPlan)
+          assert(
+            viaExecutedPlan.size == 1,
+            s"expected one report for executedPlan, got ${viaExecutedPlan.size}:\n" +
+              viaExecutedPlan.mkString("\n\n"))
+        }
+      }
+    }
+  }
+
+  // A scalar subquery is planned as a top-level plan of its own and substituted into the outer
+  // plan, and extended explain counts the plans owned by a node's expressions. The outer preview
+  // therefore has to preview its subqueries too, or it reports operators Comet does accelerate as
+  // Spark. AQE is off here so the preview and the executed plan are the same single pass; see the
+  // user guide for why the two can differ under AQE.
+  test("plan-only mode: outer report coverage matches normal planning for a scalar subquery") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true") {
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        val query = "SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)"
+
+        // What Comet really executes, subquery operators included.
+        val df = sql(query)
+        df.collect()
+        val executed = CometCoverageStats.forPlan(df.queryExecution.executedPlan)
+        assert(
+          executed.cometOperators > 0,
+          "test query must be partly accelerated for the comparison to mean anything")
+
+        withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+          val reports = capturePlanOnlyReports(sql(query).collect())
+          // One report for the separately planned subquery, one for the outer query. Only the
+          // outer one describes the Filter.
+          val outer = reports.filter(_.contains("Filter"))
+          assert(
+            outer.size == 1,
+            s"expected exactly one report for the outer query, got ${outer.size}:\n" +
+              reports.mkString("\n\n"))
           assert(
-            coverageOf(reports.head) ==
+            coverageOf(outer.head) ==
               (executed.cometOperators, executed.cometOperators + executed.sparkOperators),
-            s"report disagrees with the executed plan ($executed):\n${reports.head}")
+            s"outer report disagrees with the executed plan ($executed):\n${outer.head}")
         }
       }
     }

From c75a87ce5728eb08cb1b58ecf1b6eb2b489c4738 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 26 Aug 2026 17:41:49 -0600
Subject: [PATCH 6/8] Address review: DPP preparation scope, empty adaptive
 plans, V2 scan coverage

Preview a DPP subquery inside its BroadcastExchangeExec rather than around it.
PlanDynamicPruningFilters prepares the build plan and wraps it in the exchange
afterwards, so the plan that went through the post-columnar rules is the
exchange's child. Handing the exchange to the preview left that child a stage
whose top boundary was the exchange, so RevertNativeForTransitionHeavyStages
never fired and the outer report claimed 4/12 accelerated where the executed
plan has 2/12.

Suppress reports for the rest of an execution once AQE has cut a query stage.
When a shuffle materializes empty, AQE re-plans from the logical plan and can
collapse the tree to an empty relation; that plan shares no nodes with the one
already reported, so the mark does not reach it, and it holds no query stages
either, so the structural checks do not either. Its hash differs too, so both
guards admitted a second 0% report under the same execution ID. Subqueries are
all compiled before the first stage is cut, so this costs no report except for a
subquery AQE only plans once stages are under way, which the user guide now
notes.

Set USE_V1_SOURCE_LIST before creating the Parquet fixture in the plan-only scan
tests. withParquetTable resolves the relation through spark.read and registers
the result as a temp view, so changing the source list afterwards left both the
V1 and the V2 variant planning a FileSourceScanExec. The scan node is now
asserted so this cannot regress unnoticed.

New tests cover an empty adaptive query under AQE on and off, and compare the
outer report for a DPP subquery against CometCoverageStats for the plan Comet
really executes with stage reversion forced on.
---
 .../latest/understanding-comet-plans.md       |   5 +-
 .../apache/comet/rules/CometExecRule.scala    | 113 +++++++++-----
 .../comet/rules/CometExecRuleSuite.scala      | 145 ++++++++++++++++--
 3 files changed, 206 insertions(+), 57 deletions(-)

diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md
index c2e09b35f4..336fc5f6d4 100644
--- a/docs/source/user-guide/latest/understanding-comet-plans.md
+++ b/docs/source/user-guide/latest/understanding-comet-plans.md
@@ -233,7 +233,10 @@ way normal Comet planning does, so the reports for one query describe
 overlapping sets of operators and their counts should not be added up. Repeat
 applications of the same plan are not reported again: under AQE, neither the
 per-stage applications nor the applications that follow each adaptive
-re-optimization add reports.
+re-optimization add reports, and nor does a plan AQE re-plans wholesale after a
+stage materializes empty. One consequence under AQE is that a subquery Spark
+only plans once query stages are under way — a DPP subquery, for instance — has
+no report of its own.
 
 The estimate reflects Scala-side conversion only. The native plan is never
 handed to DataFusion, so anything that would have failed in DataFusion's
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 581bf58eaa..f9625ab36f 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -132,6 +132,8 @@ object CometExecRule {
 
     /** Adds `key`, returning true if it was not already present. */
     def add(key: String): Boolean = keys.put(key, java.lang.Boolean.TRUE) == null
+
+    def contains(key: String): Boolean = keys.containsKey(key)
   }
 
   private val PLAN_ONLY_REPORTED_LIMIT = 1024
@@ -139,6 +141,9 @@ object CometExecRule {
   /** `executionId:planFingerprint` keys that plan-only mode has already reported. */
   private val planOnlyReportedPlans = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
 
+  /** Execution IDs for which AQE has begun cutting the plan into query stages. */
+  private val planOnlyStagedExecutions = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
   /**
    * Set on the root of every plan that plan-only mode has reported. Catalyst copies a node's tags
    * onto the node that replaces it (`TreeNode.copyTagsFrom`), so the mark survives the rewrites
@@ -148,61 +153,69 @@ object CometExecRule {
   private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported")
 
   /**
-   * Whether `plan` is an application of this rule to a plan already reported, rather than to a
-   * plan the user is asking about. All the state lives on the plan itself, so this holds whether
-   * or not the application carries a SQL execution ID - `df.rdd.count()` and reading
-   * `executedPlan` without an action both plan, and in the first case execute AQE stages, with no
-   * execution ID set.
-   *
-   * Under AQE one query reaches the conversion rules five times, and only the first is the plan
-   * being evaluated:
+   * Whether `plan` is the plan of a query stage AQE has just cut, which reaches the columnar rule
+   * rooted at the `Exchange` the cut was made at.
    *
-   *   - the initial plan, through the query-stage-prep rule - the one to report;
-   *   - the same plan again as a columnar rule, now wrapped in `AdaptiveSparkPlanExec`;
-   *   - each query stage as it is created, again as a columnar rule, rooted at the `Exchange` the
-   *     stage was cut at;
-   *   - the re-optimized plan after each stage materializes, through the prep rule;
-   *   - the final plan once every stage has materialized, as a columnar rule.
+   * With AQE off a plan rooted at an `Exchange` is an ordinary plan - `df.repartition(n)`, say -
+   * and must still be reported, hence the `aqeEnabled` guard.
+   */
+  private def isQueryStage(
+      plan: SparkPlan,
+      queryStagePrep: Boolean,
+      aqeEnabled: Boolean): Boolean = {
+    aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange]
+  }
+
+  /**
+   * Whether `plan` is an application of this rule to a plan already reported, rather than to a
+   * plan the user is asking about. The state is on the plan itself, so this holds whether or not
+   * the application carries a SQL execution ID - `df.rdd.count()` and reading `executedPlan`
+   * without an action both plan, and in the first case execute AQE stages, with no execution ID
+   * set.
    *
    * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the wrapper is matched
-   * by name rather than through its contents. The two re-optimization passes hold
-   * `QueryStageExec` nodes for the stages already materialized. The remaining two - a stage, and
-   * a final plan for a query AQE never had to cut into stages - carry the mark left by the prep
-   * rule.
+   * by name rather than through its contents. A re-optimized plan holds `QueryStageExec` nodes
+   * for the stages already materialized. A final plan for a query AQE never had to cut into
+   * stages holds neither, and is recognized by the mark left when it was first reported.
    *
    * @param queryStagePrep
    *   whether the calling rule instance is registered as a query-stage-prep rule.
-   * @param aqeEnabled
-   *   whether AQE is enabled for the session. A plan rooted at an `Exchange` is a stage only when
-   *   AQE cuts stages; with AQE off it is an ordinary plan (`df.repartition(n)`, say) and must
-   *   still be reported.
    */
-  private def isReapplication(
-      plan: SparkPlan,
-      queryStagePrep: Boolean,
-      aqeEnabled: Boolean): Boolean = {
+  private def isReapplication(plan: SparkPlan): Boolean = {
     plan.exists(p =>
       p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] ||
-        p.getTagValue(PLAN_ONLY_REPORTED).isDefined) ||
-    (aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange])
+        p.getTagValue(PLAN_ONLY_REPORTED).isDefined)
   }
 
   /**
    * Whether plan-only mode should report `plan`, marking it reported if so.
    *
    * Spark applies this rule many times while executing one query, and only some of those
-   * applications correspond to a plan the user is asking about:
+   * applications correspond to a plan the user is asking about. Under AQE one query reaches the
+   * rules at least five times:
+   *
+   *   - the initial plan, through the query-stage-prep rule - the one to report;
+   *   - the same plan again as a columnar rule, now wrapped in `AdaptiveSparkPlanExec`;
+   *   - each query stage as it is created, again as a columnar rule;
+   *   - the re-optimized plan after each stage materializes, through the prep rule;
+   *   - the final plan once every stage has materialized, as a columnar rule.
+   *
+   * Three mechanisms sort those out, because no one of them covers every shape:
    *
    *   - Each scalar subquery and DPP subquery is prepared as its own top-level plan, and that
    *     happens *before* the outer plan reaches the conversion rules. A single report slot per
    *     SQL execution therefore let a nested subquery consume the slot and suppressed the outer
    *     plan, which is the plan being evaluated. Marking plans individually gives the outer plan
-   *     its own report and each separately prepared subquery theirs.
-   *   - Under AQE the rule also runs per query stage and again after every re-optimization. Those
-   *     are re-planning of a plan already reported; see [[isReapplication]].
+   *     its own report and each separately prepared subquery theirs; see [[isReapplication]].
+   *   - A mark cannot survive AQE replacing the physical tree wholesale, which is what happens
+   *     when a stage materializes empty and the plan collapses to an empty relation: the new plan
+   *     shares no nodes with the one reported and holds no query stages either. Once AQE has cut
+   *     a stage for an execution, though, everything that arrives afterwards is AQE re-planning
+   *     something already reported, and subqueries are all compiled before the first stage is
+   *     cut, so suppressing the rest of the execution costs no report.
    *   - A subquery referenced from more than one place in the outer plan is prepared once per
-   *     reference, as a separate but identical plan each time. The mark cannot catch those - they
-   *     share no nodes - so within one SQL execution the plan's structural hash dedupes them.
+   *     reference, as a separate but identical plan each time. Neither of the above catches
+   *     those, so within one SQL execution the plan's structural hash dedupes them.
    *
    * @param queryStagePrep
    *   whether the calling rule instance is registered as a query-stage-prep rule.
@@ -212,13 +225,19 @@ object CometExecRule {
       plan: SparkPlan,
       queryStagePrep: Boolean,
       aqeEnabled: Boolean): Boolean = {
-    if (isReapplication(plan, queryStagePrep, aqeEnabled)) {
+    if (isQueryStage(plan, queryStagePrep, aqeEnabled)) {
+      // Record that AQE has started executing this query before dropping the stage itself.
+      executionId.foreach(planOnlyStagedExecutions.add)
+      false
+    } else if (isReapplication(plan)) {
+      false
+    } else if (executionId.exists(planOnlyStagedExecutions.contains)) {
       false
     } else {
       plan.setTagValue(PLAN_ONLY_REPORTED, ())
       // Node tags are not part of a plan's structural hash, so the mark set above does not
-      // perturb the key. Without an execution ID there is nothing to scope the state to, and
-      // `isReapplication` has already ruled out the repeat applications AQE makes, so report.
+      // perturb the key. Without an execution ID there is nothing to scope the state to, and the
+      // checks above have already ruled out the repeat applications AQE makes, so report.
       executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.hashCode()}"))
     }
   }
@@ -750,8 +769,24 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
     // Reuse bookkeeping: the plan to preview is one level further down.
     case reused: ReusedSubqueryExec => reused.copy(child = previewSubquery(reused.child))
     case other =>
-      val preview = buildPreview(stripPreparation(other.child), topLevel = false)
-      other.withNewChildren(Seq(preview)).asInstanceOf[BaseSubqueryExec]
+      other.withNewChildren(Seq(previewPreparedPlan(other.child))).asInstanceOf[BaseSubqueryExec]
+  }
+
+  /**
+   * Preview `plan`, which Spark prepared as a plan in its own right.
+   *
+   * A DPP subquery is the exception to that framing: `PlanDynamicPruningFilters` prepares the
+   * build plan and only then wraps it in a `BroadcastExchangeExec`, so the plan that went through
+   * the post-columnar rules - and the stage `RevertNativeForTransitionHeavyStages` judged - is
+   * the exchange's child, not the exchange. Previewing the exchange instead leaves its child a
+   * stage bounded at the top by the exchange, which stops the reversion firing, and the report
+   * then counts operators as accelerated that the executed plan runs on Spark. Descend through
+   * the wrapper and put it back, so the preview keeps the boundary Spark's preparation used.
+   */
+  private def previewPreparedPlan(plan: SparkPlan): SparkPlan = plan match {
+    case exchange: BroadcastExchangeExec =>
+      exchange.withNewChildren(Seq(previewPreparedPlan(exchange.child)))
+    case other => buildPreview(stripPreparation(other), topLevel = false)
   }
 
   /**
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index 0ebe0e9e4c..611c7c0e37 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.QueryStageExec
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec}
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
@@ -800,9 +801,14 @@ class CometExecRuleSuite extends CometTestBase {
   }
 
   /**
-   * Run `sql` with plan-only mode enabled and assert nothing was offloaded to native. `useV1`
-   * toggles between `USE_V1_SOURCE_LIST=parquet` (V1 `CometScanExec` path) and
-   * `USE_V1_SOURCE_LIST=""` (V2 `CometBatchScanExec` path).
+   * With plan-only mode enabled, plan `sql` over a fresh Parquet table and assert nothing was
+   * offloaded to native. `useV1` toggles between `USE_V1_SOURCE_LIST=parquet` (the V1
+   * `FileSourceScanExec` path) and `USE_V1_SOURCE_LIST=""` (the V2 `BatchScanExec` path).
+   *
+   * The source list has to be set before the table is created. `withParquetTable` resolves the
+   * relation through `spark.read` and registers the result as a temp view, so the choice of V1 or
+   * V2 is baked in at that point; changing the config afterwards leaves both variants planning a
+   * `FileSourceScanExec`. The scan node is asserted below so that this cannot regress unnoticed.
    */
   private def runPlanOnlyAndAssertReverted(
       sql: String,
@@ -814,11 +820,22 @@ class CometExecRuleSuite extends CometTestBase {
       CometConf.COMET_ENABLED.key -> "true",
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
-      val executed = spark.sql(sql).queryExecution.executedPlan
-      val cometNodes = stripAQEPlan(executed).collect { case p: CometPlan => p }
-      assert(
-        cometNodes.isEmpty,
-        s"plan-only mode must not offload; found Comet operators: $cometNodes")
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        val executed = stripAQEPlan(spark.sql(sql).queryExecution.executedPlan)
+        val cometNodes = executed.collect { case p: CometPlan => p }
+        assert(
+          cometNodes.isEmpty,
+          s"plan-only mode must not offload; found Comet operators: $cometNodes")
+        if (useV1) {
+          assert(
+            executed.exists(_.isInstanceOf[FileSourceScanExec]),
+            s"expected the V1 scan path, got:\n$executed")
+        } else {
+          assert(
+            executed.exists(_.isInstanceOf[BatchScanExec]),
+            s"expected the V2 scan path, got:\n$executed")
+        }
+      }
     }
   }
 
@@ -828,19 +845,15 @@ class CometExecRuleSuite extends CometTestBase {
   } {
     val label = s"${if (useV1) "V1" else "V2"} scan, AQE=$aqe"
     test(s"plan-only mode: $label") {
-      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
-        runPlanOnlyAndAssertReverted(
-          "SELECT _2, count(*) FROM tbl GROUP BY _2",
-          useV1 = useV1,
-          aqe = aqe)
-      }
+      runPlanOnlyAndAssertReverted(
+        "SELECT _2, count(*) FROM tbl GROUP BY _2",
+        useV1 = useV1,
+        aqe = aqe)
     }
   }
 
   test("plan-only mode: scalar subquery is also reverted") {
-    withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
-      runPlanOnlyAndAssertReverted("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)")
-    }
+    runPlanOnlyAndAssertReverted("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)")
   }
 
   test("plan-only mode: same query with the config off runs on Comet") {
@@ -1033,4 +1046,102 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
+  // When a shuffle materializes empty, AQE re-plans from the logical plan and can collapse the
+  // whole tree to an empty relation. That plan shares no nodes with the one already reported and
+  // holds no query stages either, so neither the mark nor the stage check recognizes it.
+  for (aqe <- Seq(true, false)) {
+    test(s"plan-only mode: one report when the plan becomes empty (AQE=$aqe)") {
+      withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+        val reports = capturePlanOnlyReports {
+          spark
+            .sql("SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2")
+            .collect()
+        }
+        assert(
+          reports.size == 1,
+          s"expected one report, got ${reports.size}:\n${reports.mkString("\n\n")}")
+        // The one report must describe the query, not the empty relation AQE replaced it with.
+        assert(
+          reports.head.contains("HashAggregate"),
+          s"the report does not describe the query:\n${reports.head}")
+      }
+    }
+  }
+
+  /** Registers a `fact` table partitioned on the join key plus a small `dim` table. */
+  private def withDppTables(f: => Unit): Unit = {
+    withTempDir { dir =>
+      withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") {
+        val sess = spark
+        import sess.implicits._
+        (0 until 400)
+          .map(i => (i, i % 10, s"f$i"))
+          .toDF("fact_id", "fact_key", "fact_str")
+          .write
+          .partitionBy("fact_key")
+          .parquet(s"${dir.getAbsolutePath}/fact")
+        (0 until 10)
+          .map(i => (i, i, s"d$i"))
+          .toDF("dim_id", "dim_key", "dim_str")
+          .write
+          .parquet(s"${dir.getAbsolutePath}/dim")
+      }
+      spark.read.parquet(s"${dir.getAbsolutePath}/fact").createOrReplaceTempView("fact")
+      spark.read.parquet(s"${dir.getAbsolutePath}/dim").createOrReplaceTempView("dim")
+      withTempView("fact", "dim")(f)
+    }
+  }
+
+  // `PlanDynamicPruningFilters` prepares the DPP build plan and only wraps it in a
+  // `BroadcastExchangeExec` afterwards, so the stage `RevertNativeForTransitionHeavyStages` judged
+  // is the exchange's child. A preview that hands it the exchange instead leaves that child's top
+  // boundary open, the reversion never fires, and the report claims acceleration the executed plan
+  // does not have. Reversion is forced on here so a missed reversion changes the number.
+  test("plan-only mode: outer report coverage matches normal planning for a DPP subquery") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      // AQE off: the preview describes the pre-adaptive plan, so only the non-adaptive plan is
+      // comparable. Comet also cannot currently run this query with AQE on and reversion enabled.
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+      SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0",
+      CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false") {
+      withDppTables {
+        val query = "SELECT f.fact_id, f.fact_str, d.dim_str FROM fact f " +
+          "JOIN dim d ON f.fact_key = d.dim_key WHERE d.dim_id < 10"
+
+        val df = sql(query)
+        df.collect()
+        val plan = df.queryExecution.executedPlan
+        // `exists` walks children only, and a DPP subquery hangs off the scan's expressions.
+        assert(
+          plan.collectWithSubqueries { case p: SubqueryBroadcastExec => p }.nonEmpty,
+          s"test query must produce a DPP subquery:\n$plan")
+        val executed = CometCoverageStats.forPlan(plan)
+
+        withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+          val reports = capturePlanOnlyReports(sql(query).collect())
+          // One report for the separately prepared DPP build plan, one for the outer query. Only
+          // the outer one describes the join.
+          val outer = reports.filter(_.contains("BroadcastHashJoin"))
+          assert(
+            outer.size == 1,
+            s"expected exactly one report for the outer query, got ${outer.size}:\n" +
+              reports.mkString("\n\n"))
+          assert(
+            coverageOf(outer.head) ==
+              (executed.cometOperators, executed.cometOperators + executed.sparkOperators),
+            s"outer report disagrees with the executed plan ($executed):\n${outer.head}")
+        }
+      }
+    }
+  }
+
 }

From 53db476c63236100aa93e801226dc25253c67ca9 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 27 Aug 2026 11:26:21 -0600
Subject: [PATCH 7/8] fix: suppress AQE re-plan reports without an execution ID
 and dedupe reused subqueries

Mark the logical plan a reported physical plan was built from, alongside the
existing mark on the physical tree. When a stage materializes empty AQE
re-plans from the logical plan and can collapse the whole tree to an empty
relation, leaving a physical tree that shares no node with the one reported and
holds no query stage, so only the logical mark reaches it. The mark needs no
SQL execution ID, which the previous stage-tracking guard did, so the duplicate
0% report is now suppressed on the RDD path PySpark's `df.rdd` takes as well.

That also removes the reason the stage-tracking guard existed, and with it the
suppression of every report that followed the first stage cut, so a subquery
AQE only plans once stages are under way - a DPP subquery - gets its report
back.

Key the per-execution dedupe on the canonicalized plan. A subquery referenced
twice is prepared once per reference as plans differing only in expression IDs,
which Spark then collapses to one `ReusedSubqueryExec`; the raw structural hash
saw two plans and reported the same work twice.
---
 .../latest/understanding-comet-plans.md       |  4 +-
 .../apache/comet/rules/CometExecRule.scala    | 68 ++++++-------
 .../comet/rules/CometExecRuleSuite.scala      | 96 ++++++++++++++++++-
 3 files changed, 127 insertions(+), 41 deletions(-)

diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md
index 336fc5f6d4..8a2fce6b60 100644
--- a/docs/source/user-guide/latest/understanding-comet-plans.md
+++ b/docs/source/user-guide/latest/understanding-comet-plans.md
@@ -234,9 +234,7 @@ overlapping sets of operators and their counts should not be added up. Repeat
 applications of the same plan are not reported again: under AQE, neither the
 per-stage applications nor the applications that follow each adaptive
 re-optimization add reports, and nor does a plan AQE re-plans wholesale after a
-stage materializes empty. One consequence under AQE is that a subquery Spark
-only plans once query stages are under way — a DPP subquery, for instance — has
-no report of its own.
+stage materializes empty.
 
 The estimate reflects Scala-side conversion only. The native plan is never
 handed to DataFusion, so anything that would have failed in DataFusion's
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index f9625ab36f..75526fd18a 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -141,9 +141,6 @@ object CometExecRule {
   /** `executionId:planFingerprint` keys that plan-only mode has already reported. */
   private val planOnlyReportedPlans = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
 
-  /** Execution IDs for which AQE has begun cutting the plan into query stages. */
-  private val planOnlyStagedExecutions = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
-
   /**
    * Set on the root of every plan that plan-only mode has reported. Catalyst copies a node's tags
    * onto the node that replaces it (`TreeNode.copyTagsFrom`), so the mark survives the rewrites
@@ -152,6 +149,18 @@ object CometExecRule {
    */
   private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported")
 
+  /**
+   * The same mark, placed on the logical plan the reported physical plan was linked to.
+   *
+   * AQE can replace the physical tree wholesale - a stage that materializes empty collapses the
+   * whole plan to an empty relation - leaving a tree that shares no node with the one already
+   * reported and holds no query stage either, so [[PLAN_ONLY_REPORTED]] cannot reach it. Every
+   * such tree is planned from a logical plan derived from the one already reported, and Catalyst
+   * copies tags onto replacement nodes there too, so the logical mark does reach it.
+   */
+  private val PLAN_ONLY_REPORTED_LOGICAL: TreeNodeTag[Unit] =
+    TreeNodeTag[Unit]("comet.planOnlyReportedLogical")
+
   /**
    * Whether `plan` is the plan of a query stage AQE has just cut, which reaches the columnar rule
    * rooted at the `Exchange` the cut was made at.
@@ -176,15 +185,15 @@ object CometExecRule {
    * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the wrapper is matched
    * by name rather than through its contents. A re-optimized plan holds `QueryStageExec` nodes
    * for the stages already materialized. A final plan for a query AQE never had to cut into
-   * stages holds neither, and is recognized by the mark left when it was first reported.
-   *
-   * @param queryStagePrep
-   *   whether the calling rule instance is registered as a query-stage-prep rule.
+   * stages holds neither, and is recognized by the mark left when it was first reported. A plan
+   * AQE rebuilt from scratch carries none of those, and is recognized by the mark on the logical
+   * plan it was built from.
    */
   private def isReapplication(plan: SparkPlan): Boolean = {
     plan.exists(p =>
       p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] ||
-        p.getTagValue(PLAN_ONLY_REPORTED).isDefined)
+        p.getTagValue(PLAN_ONLY_REPORTED).isDefined) ||
+    plan.logicalLink.exists(_.getTagValue(PLAN_ONLY_REPORTED_LOGICAL).isDefined)
   }
 
   /**
@@ -200,22 +209,20 @@ object CometExecRule {
    *   - the re-optimized plan after each stage materializes, through the prep rule;
    *   - the final plan once every stage has materialized, as a columnar rule.
    *
-   * Three mechanisms sort those out, because no one of them covers every shape:
+   * Two mechanisms sort those out, because neither on its own covers every shape:
    *
-   *   - Each scalar subquery and DPP subquery is prepared as its own top-level plan, and that
-   *     happens *before* the outer plan reaches the conversion rules. A single report slot per
-   *     SQL execution therefore let a nested subquery consume the slot and suppressed the outer
-   *     plan, which is the plan being evaluated. Marking plans individually gives the outer plan
-   *     its own report and each separately prepared subquery theirs; see [[isReapplication]].
-   *   - A mark cannot survive AQE replacing the physical tree wholesale, which is what happens
-   *     when a stage materializes empty and the plan collapses to an empty relation: the new plan
-   *     shares no nodes with the one reported and holds no query stages either. Once AQE has cut
-   *     a stage for an execution, though, everything that arrives afterwards is AQE re-planning
-   *     something already reported, and subqueries are all compiled before the first stage is
-   *     cut, so suppressing the rest of the execution costs no report.
+   *   - A mark on the plan already reported, both on the physical tree and on the logical plan it
+   *     was built from, so that a later application recognizes it however AQE rewrote it; see
+   *     [[isReapplication]]. Marking plans individually, rather than holding one report slot per
+   *     SQL execution, is what gives the outer query a report of its own: each scalar subquery
+   *     and DPP subquery is prepared as a top-level plan in its own right, and for most of them
+   *     that happens *before* the outer plan reaches the conversion rules, so a single slot would
+   *     be consumed by a subquery and the plan being evaluated would never be described.
    *   - A subquery referenced from more than one place in the outer plan is prepared once per
-   *     reference, as a separate but identical plan each time. Neither of the above catches
-   *     those, so within one SQL execution the plan's structural hash dedupes them.
+   *     reference, as a separate plan each time, differing only in expression IDs; Spark then
+   *     collapses them to one `ReusedSubqueryExec`. No mark connects those, so within one SQL
+   *     execution the canonicalized plan's hash dedupes them. Canonicalization is what makes the
+   *     expression IDs drop out; the raw structural hash sees two different plans.
    *
    * @param queryStagePrep
    *   whether the calling rule instance is registered as a query-stage-prep rule.
@@ -225,20 +232,15 @@ object CometExecRule {
       plan: SparkPlan,
       queryStagePrep: Boolean,
       aqeEnabled: Boolean): Boolean = {
-    if (isQueryStage(plan, queryStagePrep, aqeEnabled)) {
-      // Record that AQE has started executing this query before dropping the stage itself.
-      executionId.foreach(planOnlyStagedExecutions.add)
-      false
-    } else if (isReapplication(plan)) {
-      false
-    } else if (executionId.exists(planOnlyStagedExecutions.contains)) {
+    if (isQueryStage(plan, queryStagePrep, aqeEnabled) || isReapplication(plan)) {
       false
     } else {
       plan.setTagValue(PLAN_ONLY_REPORTED, ())
-      // Node tags are not part of a plan's structural hash, so the mark set above does not
-      // perturb the key. Without an execution ID there is nothing to scope the state to, and the
-      // checks above have already ruled out the repeat applications AQE makes, so report.
-      executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.hashCode()}"))
+      plan.logicalLink.foreach(_.setTagValue(PLAN_ONLY_REPORTED_LOGICAL, ()))
+      // Node tags are not part of a plan's canonical form, so the marks set above do not perturb
+      // the key. Without an execution ID there is nothing to scope the state to, and the checks
+      // above have already ruled out the repeat applications AQE makes, so report.
+      executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.canonicalized.hashCode()}"))
     }
   }
 }
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index 611c7c0e37..4ba243e831 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -1049,17 +1049,25 @@ class CometExecRuleSuite extends CometTestBase {
   // When a shuffle materializes empty, AQE re-plans from the logical plan and can collapse the
   // whole tree to an empty relation. That plan shares no nodes with the one already reported and
   // holds no query stages either, so neither the mark nor the stage check recognizes it.
-  for (aqe <- Seq(true, false)) {
-    test(s"plan-only mode: one report when the plan becomes empty (AQE=$aqe)") {
+  for {
+    aqe <- Seq(true, false)
+    // `collect()` installs a SQL execution ID; going straight to the RDD - the path PySpark's
+    // `df.rdd` takes through `Dataset.javaToPython` - does not, so the suppression cannot be
+    // scoped to one.
+    (action, runIt) <- Seq[(String, org.apache.spark.sql.DataFrame => Unit)](
+      "collect" -> (df => df.collect()),
+      "toRdd.count" -> (df => df.queryExecution.toRdd.count()))
+  } {
+    test(s"plan-only mode: one report when the plan becomes empty (AQE=$aqe, $action)") {
       withSQLConf(
         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
         CometConf.COMET_ENABLED.key -> "true",
         CometConf.COMET_EXEC_ENABLED.key -> "true",
         CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
         val reports = capturePlanOnlyReports {
-          spark
-            .sql("SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2")
-            .collect()
+          runIt(
+            spark.sql(
+              "SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2"))
         }
         assert(
           reports.size == 1,
@@ -1072,6 +1080,36 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
+  // The same nested scalar subquery projected twice is prepared once per reference, as separate
+  // but structurally identical plans that differ only in expression IDs. Spark reuses one of them
+  // through `ReusedSubqueryExec`, so reporting both describes the same work twice.
+  test("plan-only mode: a subquery referenced twice is reported once") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        val reports = capturePlanOnlyReports {
+          spark
+            .sql("""SELECT _1,
+                   |  (SELECT max(_2) FROM tbl
+                   |   WHERE _1 > (SELECT min(_2) FROM tbl)) AS a,
+                   |  (SELECT max(_2) FROM tbl
+                   |   WHERE _1 > (SELECT min(_2) FROM tbl)) AS b
+                   |FROM tbl""".stripMargin)
+            .collect()
+        }
+        // One for the innermost `min` subquery, one for the `max` subquery, one for the outer
+        // query. The second reference to the `max` subquery must not add a fourth.
+        assert(
+          reports.size == 3,
+          s"expected three reports, got ${reports.size}:\n${reports.mkString("\n\n")}")
+      }
+    }
+  }
+
   /** Registers a `fact` table partitioned on the join key plus a small `dim` table. */
   private def withDppTables(f: => Unit): Unit = {
     withTempDir { dir =>
@@ -1144,4 +1182,52 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
+  // Every plan Spark prepares in its own right gets its own report, including one AQE only plans
+  // once query stages are under way. A DPP build plan is the case in point: it is prepared by
+  // `PlanAdaptiveDynamicPruningFilters`, a stage optimizer rule, so it arrives after the outer
+  // query has been reported and after the first stage has been cut.
+  test("plan-only mode: a DPP subquery planned mid-execution is reported under AQE") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+      SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+      withDppTables {
+        val query = "SELECT f.fact_id, f.fact_str, d.dim_str FROM fact f " +
+          "JOIN dim d ON f.fact_key = d.dim_key WHERE d.dim_id < 10"
+        val reports = capturePlanOnlyReports(sql(query).collect())
+        assert(
+          reports.size == 2,
+          s"expected a report for the outer query and one for the DPP build plan, got " +
+            s"${reports.size}:\n${reports.mkString("\n\n")}")
+        assert(
+          reports.count(_.contains("BroadcastHashJoin")) == 1,
+          s"exactly one report should describe the outer query:\n${reports.mkString("\n\n")}")
+      }
+    }
+  }
+
+  // A query AQE cuts into several stages reaches the rules once per stage and once per
+  // re-optimization on top of the initial planning. None of those may add a report.
+  test("plan-only mode: one report for a multi-stage query under AQE") {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+      // Force a shuffled join so the plan has more than one shuffle boundary.
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+      withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+        val query = "SELECT a._2, count(*) FROM tbl a JOIN tbl b ON a._1 = b._2 " +
+          "GROUP BY a._2 ORDER BY 1"
+        val reports = capturePlanOnlyReports(sql(query).collect())
+        assert(
+          reports.size == 1,
+          s"expected one report, got ${reports.size}:\n${reports.mkString("\n\n")}")
+      }
+    }
+  }
 }

From 0896e36ed798ca3eda213e1fca4d9be6d47bd157 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 27 Aug 2026 14:00:28 -0600
Subject: [PATCH 8/8] 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.
---
 .../apache/comet/rules/CometExecRule.scala    | 70 ++++++++++++-------
 .../comet/rules/CometExecRuleSuite.scala      | 29 ++++++--
 2 files changed, 65 insertions(+), 34 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 75526fd18a..797c3d4eb5 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -149,18 +149,6 @@ object CometExecRule {
    */
   private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported")
 
-  /**
-   * The same mark, placed on the logical plan the reported physical plan was linked to.
-   *
-   * AQE can replace the physical tree wholesale - a stage that materializes empty collapses the
-   * whole plan to an empty relation - leaving a tree that shares no node with the one already
-   * reported and holds no query stage either, so [[PLAN_ONLY_REPORTED]] cannot reach it. Every
-   * such tree is planned from a logical plan derived from the one already reported, and Catalyst
-   * copies tags onto replacement nodes there too, so the logical mark does reach it.
-   */
-  private val PLAN_ONLY_REPORTED_LOGICAL: TreeNodeTag[Unit] =
-    TreeNodeTag[Unit]("comet.planOnlyReportedLogical")
-
   /**
    * Whether `plan` is the plan of a query stage AQE has just cut, which reaches the columnar rule
    * rooted at the `Exchange` the cut was made at.
@@ -185,15 +173,34 @@ object CometExecRule {
    * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the wrapper is matched
    * by name rather than through its contents. A re-optimized plan holds `QueryStageExec` nodes
    * for the stages already materialized. A final plan for a query AQE never had to cut into
-   * stages holds neither, and is recognized by the mark left when it was first reported. A plan
-   * AQE rebuilt from scratch carries none of those, and is recognized by the mark on the logical
-   * plan it was built from.
+   * stages holds neither, and is recognized by the mark left when it was first reported.
    */
   private def isReapplication(plan: SparkPlan): Boolean = {
     plan.exists(p =>
       p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] ||
-        p.getTagValue(PLAN_ONLY_REPORTED).isDefined) ||
-    plan.logicalLink.exists(_.getTagValue(PLAN_ONLY_REPORTED_LOGICAL).isDefined)
+        p.getTagValue(PLAN_ONLY_REPORTED).isDefined)
+  }
+
+  /**
+   * Whether `plan` is a plan AQE re-optimized down to nothing, which is the one re-optimization
+   * result [[isReapplication]] cannot recognize.
+   *
+   * A re-optimized plan normally holds a `QueryStageExec` for each stage that has materialized.
+   * The exception is a stage that materialized empty: `AQEPropagateEmptyRelation` then replaces
+   * the stages, and everything above them, with an empty relation, so the plan Spark hands back
+   * shares no node with the plan already reported and holds no query stage either. Empty is the
+   * only shape that eliminates every stage, and Catalyst records it as a plan whose row count is
+   * known to be zero, so the logical link answers the question directly.
+   *
+   * The check is confined to the query-stage-prep rule, the only one AQE hands a re-optimized
+   * plan to. A genuinely empty query - `WHERE false`, an empty local relation - is planned once,
+   * reaches the columnar rule instead, and is still reported.
+   */
+  private def isAdaptiveReplanToNothing(
+      plan: SparkPlan,
+      queryStagePrep: Boolean,
+      aqeEnabled: Boolean): Boolean = {
+    aqeEnabled && queryStagePrep && plan.logicalLink.exists(_.maxRows.contains(0L))
   }
 
   /**
@@ -209,21 +216,27 @@ object CometExecRule {
    *   - the re-optimized plan after each stage materializes, through the prep rule;
    *   - the final plan once every stage has materialized, as a columnar rule.
    *
-   * Two mechanisms sort those out, because neither on its own covers every shape:
+   * Three mechanisms sort those out, because no one of them covers every shape:
    *
-   *   - A mark on the plan already reported, both on the physical tree and on the logical plan it
-   *     was built from, so that a later application recognizes it however AQE rewrote it; see
-   *     [[isReapplication]]. Marking plans individually, rather than holding one report slot per
-   *     SQL execution, is what gives the outer query a report of its own: each scalar subquery
-   *     and DPP subquery is prepared as a top-level plan in its own right, and for most of them
-   *     that happens *before* the outer plan reaches the conversion rules, so a single slot would
-   *     be consumed by a subquery and the plan being evaluated would never be described.
+   *   - A mark on the plan already reported, so that a later application recognizes it however
+   *     Spark rewrote it in between; see [[isReapplication]]. Marking plans individually, rather
+   *     than holding one report slot per SQL execution, is what gives the outer query a report of
+   *     its own: each scalar subquery and DPP subquery is prepared as a top-level plan in its own
+   *     right, and for most of them that happens *before* the outer plan reaches the conversion
+   *     rules, so a single slot would be consumed by a subquery and the plan being evaluated
+   *     would never be described.
+   *   - One re-optimization result carries no trace of the plan it replaced, and is recognized by
+   *     what it is rather than by a mark; see [[isAdaptiveReplanToNothing]].
    *   - A subquery referenced from more than one place in the outer plan is prepared once per
    *     reference, as a separate plan each time, differing only in expression IDs; Spark then
    *     collapses them to one `ReusedSubqueryExec`. No mark connects those, so within one SQL
    *     execution the canonicalized plan's hash dedupes them. Canonicalization is what makes the
    *     expression IDs drop out; the raw structural hash sees two different plans.
    *
+   * None of this is scoped to a SQL execution ID, which `df.rdd.count()` and reading
+   * `executedPlan` without an action both plan - and in the first case execute AQE stages -
+   * without.
+   *
    * @param queryStagePrep
    *   whether the calling rule instance is registered as a query-stage-prep rule.
    */
@@ -235,11 +248,14 @@ object CometExecRule {
     if (isQueryStage(plan, queryStagePrep, aqeEnabled) || isReapplication(plan)) {
       false
     } else {
+      // Mark before deciding. A plan AQE re-optimized to nothing is not worth a report, but the
+      // final plan Spark builds from it reaches the columnar rule next and has to be recognized
+      // as a plan already dealt with.
       plan.setTagValue(PLAN_ONLY_REPORTED, ())
-      plan.logicalLink.foreach(_.setTagValue(PLAN_ONLY_REPORTED_LOGICAL, ()))
-      // Node tags are not part of a plan's canonical form, so the marks set above do not perturb
+      // Node tags are not part of a plan's canonical form, so the mark set above does not perturb
       // the key. Without an execution ID there is nothing to scope the state to, and the checks
       // above have already ruled out the repeat applications AQE makes, so report.
+      !isAdaptiveReplanToNothing(plan, queryStagePrep, aqeEnabled) &&
       executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.canonicalized.hashCode()}"))
     }
   }
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index 4ba243e831..c513a94e26 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -1057,24 +1057,39 @@ class CometExecRuleSuite extends CometTestBase {
     (action, runIt) <- Seq[(String, org.apache.spark.sql.DataFrame => Unit)](
       "collect" -> (df => df.collect()),
       "toRdd.count" -> (df => df.queryExecution.toRdd.count()))
+    // The second shape is the one where the plan that AQE reports back cannot be recognized by a
+    // mark of any kind: `RemoveRedundantSorts` drops the outer sort before this rule sees the
+    // plan, so the mark lands on the join below it, and empty propagation replaces the join and
+    // then the sort, leaving a root that inherited the unmarked sort's tags.
+    (shape, query, marker) <- Seq(
+      (
+        "aggregate",
+        "SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2",
+        "HashAggregate"),
+      (
+        "join under a removed root sort",
+        """SELECT a.id FROM range(0, 20, 1, 2) a
+          |JOIN range(0, 20, 1, 2) b ON a.id % 7 = b.id % 7
+          |WHERE a.id < 0
+          |SORT BY a.id % 7""".stripMargin,
+        "SortMergeJoin"))
   } {
-    test(s"plan-only mode: one report when the plan becomes empty (AQE=$aqe, $action)") {
+    test(s"plan-only mode: one report when the plan becomes empty ($shape, AQE=$aqe, $action)") {
       withSQLConf(
         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+        // Force a shuffled join so the join shape has a stage that can materialize empty.
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+        SQLConf.SHUFFLE_PARTITIONS.key -> "2",
         CometConf.COMET_ENABLED.key -> "true",
         CometConf.COMET_EXEC_ENABLED.key -> "true",
         CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
-        val reports = capturePlanOnlyReports {
-          runIt(
-            spark.sql(
-              "SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2"))
-        }
+        val reports = capturePlanOnlyReports(runIt(spark.sql(query)))
         assert(
           reports.size == 1,
           s"expected one report, got ${reports.size}:\n${reports.mkString("\n\n")}")
         // The one report must describe the query, not the empty relation AQE replaced it with.
         assert(
-          reports.head.contains("HashAggregate"),
+          reports.head.contains(marker),
           s"the report does not describe the query:\n${reports.head}")
       }
     }