feat: add spark.comet.explain.planOnly.enabled

docs/source/user-guide/latest/understanding-comet-plans.md

@@ -208,6 +208,31 @@ 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 plans every query as it normally would, logs the resulting Comet plan and
+coverage summary to the driver log with the prefix `[Comet plan-only]`, and then executes the
+original plan on Spark. Use it to estimate how much of a workload Comet would accelerate
+without changing execution. The report is the same annotated plan that
+`spark.comet.explain.format=verbose` produces, and reflects Comet's post-columnar rules, so a
+stage Comet would revert to Spark for having too many transitions is reported as reverted.
+Requires `spark.comet.exec.enabled=true`.
+
+Keep the following in mind when reading the reports:
+
+- Spark plans scalar and dynamic partition pruning subqueries separately from the query that
+  contains them, so each gets its own report alongside the outer query's. The outer report
+  also counts its subqueries, so do not add the reports together.
+- Only the JVM side of planning runs. Anything that would fail when DataFusion builds the
+  native plan still counts as accelerated, so treat the percentage as an upper bound.
+- Comet's split Iceberg V2 write (`spark.comet.write.iceberg.splitOperator.enabled`) is
+  declined in plan-only mode, so such writes run on, and are reported as, Spark.
+- Under AQE the report describes the plan before any adaptive re-planning, so coverage of the
+  plan that finally executes can differ. In particular, AQE plans subqueries into the outer
+  query only after the report is produced, so the outer report counts their operators as
+  Spark. For subquery-heavy queries, read the per-subquery reports or disable AQE for the
+  evaluation run.
+
 ## Programmatic Access to Fallback Reasons
 
 The configs above route fallback reasons to logs or the SQL UI. If you want

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -785,6 +785,17 @@ 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 logs the plan it would have executed, with a coverage " +
+          "summary, to the driver log and then lets Spark execute the query unchanged. Native " +
+          "planning failures are not detected, so the coverage can be optimistic. Requires " +
+          "`spark.comet.exec.enabled=true`.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_STRICT_FALLBACK_REASONS: ConfigEntry[Boolean] =
     conf("spark.comet.explain.fallback.strict.enabled")
       .category(CATEGORY_TESTING)

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -34,7 +34,7 @@ import org.apache.spark.sql.internal.SQLConf
 
 import org.apache.comet.CometConf._
 import org.apache.comet.iceberg.IcebergWriteStrategy
-import org.apache.comet.rules.{CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages}
+import org.apache.comet.rules.{CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule}
 import org.apache.comet.shims.ShimCometSparkSessionExtensions
 
 /**
@@ -96,7 +96,9 @@ class CometSparkSessionExtensions
     // Registered before CometRule so tags are in place when conversion runs.
     // No-op on Spark 3.5+; see CometSpark34AqeDppFallbackRule's class docstring.
     injectPreSpark35QueryStagePrepRuleShim(extensions, CometSpark34AqeDppFallbackRule)
-    extensions.injectQueryStagePrepRule { session => CometRule(session) }
+    extensions.injectQueryStagePrepRule { session =>
+      CometRule(session, queryStagePrep = true)
+    }
     injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters)
     injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery)
     extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) }
@@ -106,8 +108,7 @@ class CometSparkSessionExtensions
     override def preColumnarTransitions: Rule[SparkPlan] = CometRule(session)
 
     override def postColumnarTransitions: Rule[SparkPlan] = {
-      val rules =
-        Seq(RevertNativeForTransitionHeavyStages(session), EliminateRedundantTransitions(session))
+      val rules = CometRule.postColumnarRules(session)
       plan => rules.foldLeft(plan) { case (p, rule) => rule(p) }
     }
   }

spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala

@@ -38,6 +38,10 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy {
     if (!CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.get(session.sessionState.conf)) {
       return Nil
     }
+    // Planner strategies run before CometRule, so plan-only mode needs its own guard here.
+    if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(session.sessionState.conf)) {
+      return Nil
+    }
 
     plan match {
       case ad: AppendData =>

spark/src/main/scala/org/apache/comet/rules/CometRule.scala

@@ -19,9 +19,83 @@
 
 package org.apache.comet.rules
 
+import scala.util.control.NonFatal
+
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, ExecSubqueryExpression, InputAdapter, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, SQLExecution, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange}
+import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+import org.apache.comet.CometSparkSessionExtensions.isCometLoaded
+import org.apache.comet.shims.ShimCometStreaming
+
+object CometRule {
+
+  /** Comet's post-columnar rules, shared by `CometColumnar` and the plan-only preview. */
+  def postColumnarRules(session: SparkSession, wholePlan: Boolean = false): Seq[Rule[SparkPlan]] =
+    Seq(
+      RevertNativeForTransitionHeavyStages(session, wholePlan),
+      EliminateRedundantTransitions(session))
+
+  /**
+   * `executionId:canonicalPlanHash` keys already reported, LRU-bounded for long-lived drivers.
+   */
+  private val planOnlyReportedPlans: java.util.Set[String] =
+    java.util.Collections.newSetFromMap(
+      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() > 1024
+        }))
+
+  /**
+   * Marks the root of a reported plan. Catalyst copies tags onto replacement nodes, so later
+   * applications of the rule can recognize the plan however Spark rewrote it in between.
+   */
+  private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported")
+
+  /**
+   * Whether plan-only mode should report `plan`, marking it reported if so.
+   *
+   * Under AQE the rule sees one query several times: the initial plan (prep rule, the one to
+   * report), the same plan wrapped in `AdaptiveSparkPlanExec`, each query stage, each
+   * re-optimization and the final plan. Only the first is reported. Each scalar or DPP subquery
+   * is prepared as a top-level plan of its own and gets its own report. None of this can rely on
+   * a SQL execution ID, since `df.rdd.count()` and `executedPlan` plan without one.
+   */
+  private[comet] def shouldReportPlanOnly(
+      executionId: Option[String],
+      plan: SparkPlan,
+      queryStagePrep: Boolean,
+      aqeEnabled: Boolean): Boolean = {
+    // Under AQE the columnar rule sees each new query stage rooted at its Exchange.
+    val isQueryStage = aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange]
+    // Already reported: a re-optimized plan holds query stages, and a final plan carries the
+    // mark. `AdaptiveSparkPlanExec` is a leaf to `exists`, so it is matched directly.
+    val isReapplication = plan.exists(p =>
+      p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] ||
+        p.getTagValue(PLAN_ONLY_REPORTED).isDefined)
+    if (isQueryStage || isReapplication) {
+      false
+    } else {
+      // Mark even when not reporting, so the final plan built from an empty re-plan is recognized.
+      plan.setTagValue(PLAN_ONLY_REPORTED, ())
+      // AQE re-plans to an empty relation when a stage materializes empty, sharing no nodes or
+      // stages with the reported plan. Only the prep rule sees re-plans, so a genuinely empty
+      // query still reaches the columnar rule and is reported.
+      val replannedToNothing =
+        aqeEnabled && queryStagePrep && plan.logicalLink.exists(_.maxRows.contains(0L))
+      // A subquery referenced twice is prepared twice, differing only in expression IDs, so dedupe
+      // on the canonical plan within an execution. Tags are not part of the canonical form.
+      !replannedToNothing &&
+      executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.canonicalized.hashCode()}"))
+    }
+  }
+}
 
 /**
  * Comet's plan conversion pass: scan conversion followed by operator conversion.
@@ -36,11 +110,109 @@ import org.apache.spark.sql.execution.SparkPlan
  * delegates to their `apply`. Spark's own plan change log sees one rule: query-stage preparation
  * logs this pass as `org.apache.comet.rules.CometRule`, which is the name
  * `spark.sql.planChangeLog.rules` has to match.
+ *
+ * @param queryStagePrep
+ *   true for the `injectQueryStagePrepRule` instance, which sees the whole initial plan under
+ *   AQE. Only plan-only reporting reads it.
  */
-case class CometRule(session: SparkSession) extends Rule[SparkPlan] {
+case class CometRule(session: SparkSession, queryStagePrep: Boolean = false)
+    extends Rule[SparkPlan] {
 
   private val scanRule = CometScanRule(session)
   private val execRule = CometExecRule(session)
 
-  override def apply(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan))
+  override def apply(plan: SparkPlan): SparkPlan = {
+    if (planOnlyApplies(plan)) {
+      reportPlanOnlyCoverage(plan)
+      plan
+    } else {
+      convert(plan)
+    }
+  }
+
+  private def convert(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan))
+
+  /** Mirrors the conversion rules' own guards; plan-only is scoped to exec being enabled. */
+  private def planOnlyApplies(plan: SparkPlan): Boolean =
+    CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) &&
+      isCometLoaded(conf) &&
+      !ShimCometStreaming.isStreamingPlan(plan) &&
+      CometConf.COMET_EXEC_ENABLED.get(conf)
+
+  /** Logs the Comet plan for `plan` unless already reported. Never fails the query. */
+  private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
+    try {
+      val executionId = Option(
+        session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
+      if (CometRule.shouldReportPlanOnly(
+          executionId,
+          plan,
+          queryStagePrep,
+          conf.adaptiveExecutionEnabled)) {
+        val preview = buildPreview(plan, topLevel = true)
+        logWarning(
+          s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}")
+      }
+    } catch {
+      case NonFatal(e) =>
+        logWarning("[Comet plan-only] could not build a coverage report for this query", e)
+    }
+  }
+
+  /**
+   * The plan Comet would execute for `plan`: conversion, columnar transitions, then the
+   * post-columnar rules, with stage reversion visiting every stage since this is the whole plan.
+   *
+   * @param topLevel
+   *   false for subquery plans, which Spark prepares without `ReuseExchangeAndSubquery`.
+   */
+  private def buildPreview(plan: SparkPlan, topLevel: Boolean): SparkPlan = {
+    val converted = convert(previewSubqueriesOf(plan))
+    val withTransitions =
+      ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted)
+    val preview = CometRule
+      .postColumnarRules(session, wholePlan = true)
+      .foldLeft(withTransitions) { case (p, rule) => rule(p) }
+    if (topLevel) ReuseExchangeAndSubquery.apply(preview) else preview
+  }
+
+  /**
+   * `plan` with each subquery's plan replaced by its preview. Extended explain counts subquery
+   * operators, and normal planning has already converted them, so leaving them would understate
+   * coverage.
+   */
+  private def previewSubqueriesOf(plan: SparkPlan): SparkPlan = {
+    plan.transformAllExpressions { case subquery: ExecSubqueryExpression =>
+      subquery.withNewPlan(previewSubquery(subquery.plan))
+    }
+  }
+
+  private def previewSubquery(subquery: BaseSubqueryExec): BaseSubqueryExec = subquery match {
+    case reused: ReusedSubqueryExec => reused.copy(child = previewSubquery(reused.child))
+    case other =>
+      other.withNewChildren(Seq(previewPreparedPlan(other.child))).asInstanceOf[BaseSubqueryExec]
+  }
+
+  /**
+   * `PlanDynamicPruningFilters` prepares a DPP build plan before wrapping it in a
+   * `BroadcastExchangeExec`, so preview the exchange's child on its own to keep the stage
+   * boundary that stage reversion saw.
+   */
+  private def previewPreparedPlan(plan: SparkPlan): SparkPlan = plan match {
+    case exchange: BroadcastExchangeExec =>
+      exchange.withNewChildren(Seq(previewPreparedPlan(exchange.child)))
+    case other => buildPreview(stripPreparation(other), topLevel = false)
+  }
+
+  /**
+   * Removes codegen wrappers and transitions from an already-prepared subquery plan, since the
+   * conversion rules expect a plan from before those are inserted. [[buildPreview]] re-inserts
+   * the transitions.
+   */
+  private def stripPreparation(plan: SparkPlan): SparkPlan = plan.transformUp {
+    case WholeStageCodegenExec(child) => child
+    case InputAdapter(child) => child
+    case ColumnarToRowExec(child) => child
+    case RowToColumnarExec(child) => child
+  }
 }

spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala

@@ -36,8 +36,12 @@ import org.apache.comet.serde.QueryPlanSerde
  * Reverts a query stage to Spark row-based execution when it has too many columnar-to-row (C2R)
  * transitions. Each C2R indicates Comet could not keep execution columnar and had to fall back.
  * With columnar shuffle enabled, each C2R implies a corresponding R2C round-trip.
+ *
+ * @param wholePlan
+ *   visit every stage even under AQE, where Spark normally hands this rule one stage at a time.
+ *   Set by the plan-only preview, which holds the whole plan.
  */
-case class RevertNativeForTransitionHeavyStages(session: SparkSession)
+case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan: Boolean = false)
     extends Rule[SparkPlan]
     with Logging {
 
@@ -47,7 +51,7 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession)
   override def apply(plan: SparkPlan): SparkPlan = {
     if (!enabled) return plan
 
-    if (session.sessionState.conf.adaptiveExecutionEnabled) {
+    if (session.sessionState.conf.adaptiveExecutionEnabled && !wholePlan) {
       applyForAQE(plan)
     } else {
       applyForNonAQE(plan)

spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala

@@ -79,6 +79,26 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes
     }
   }
 
+  test("plan-only mode leaves the write with Spark") {
+    withDetectionCatalog { dir =>
+      createTable(dir, "plan_only", partitionSpec = "")
+      // IcebergWriteStrategy runs before CometRule, so it needs its own plan-only guard.
+      // withSQLConf returns Unit on Spark 3.4/3.5, hence the var.
+      var plan: SparkPlan = null
+      withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+        plan = captureWritePlan("plan_only", allowWriteFailure = false) {
+          spark.sql(s"INSERT INTO $catalog.$ns.plan_only VALUES (1, 'us', 1.0)")
+        }
+      }
+      assert(
+        findWriteExec(plan).isEmpty,
+        s"plan-only mode must not split the write into Comet's two-operator shape:\n$plan")
+      assert(
+        !containsCometWriteExec(plan),
+        s"plan-only mode must not offload the write to Comet:\n$plan")
+    }
+  }
+
   test("SparkWrite reflection helpers all resolve on the current Iceberg runtime") {
     withDetectionCatalog { dir =>
       createTable(dir, "refl_probe", partitionSpec = "")