From b666339d13e99ee74492a67bc32f761cc03a7b64 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 19 Aug 2026 06:09:14 -0600 Subject: [PATCH 01/14] 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 a600615b4b3..a12c3f9f601 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 f40d0cbdd97..94f8fafd86f 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 cbf6c6a1e89..5d7aa95e132 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 2c15bb5e4e7..df02250e6f1 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 5444a89fa36..3291e449588 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 Date: Wed, 19 Aug 2026 06:26:51 -0600 Subject: [PATCH 02/14] 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 5d7aa95e132..36d84b435a3 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 df02250e6f1..e7baa19cc35 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 3291e449588..97296db046a 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 Date: Wed, 19 Aug 2026 06:37:12 -0600 Subject: [PATCH 03/14] 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 36d84b435a3..9ee204ae16e 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 e7baa19cc35..e852ddc1e48 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 Date: Wed, 19 Aug 2026 11:14:39 -0600 Subject: [PATCH 04/14] 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 a12c3f9f601..363d5d1f787 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 4dc36194331..22a450d0a00 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 9ee204ae16e..75db4b40e63 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 1e0cfc79e0b..88783c65722 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 97296db046a..f769721efbe 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 Date: Wed, 26 Aug 2026 15:19:38 -0600 Subject: [PATCH 05/14] 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 363d5d1f787..c2e09b35f48 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 75db4b40e63..581bf58eaa9 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 f769721efbe..0ebe0e9e4c8 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 Date: Wed, 26 Aug 2026 17:41:49 -0600 Subject: [PATCH 06/14] 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 c2e09b35f48..336fc5f6d47 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 581bf58eaa9..f9625ab36f4 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 0ebe0e9e4c8..611c7c0e378 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 Date: Thu, 27 Aug 2026 11:26:21 -0600 Subject: [PATCH 07/14] 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 336fc5f6d47..8a2fce6b60b 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 f9625ab36f4..75526fd18aa 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 611c7c0e378..4ba243e831a 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 Date: Thu, 27 Aug 2026 14:00:28 -0600 Subject: [PATCH 08/14] 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 75526fd18aa..797c3d4eb59 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 4ba243e831a..c513a94e261 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}") } } From 74fe586c94455b18e757eb51bd2b3a839e0a1583 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 27 Aug 2026 15:26:03 -0600 Subject: [PATCH 09/14] fix: drop a redundant interpolator and guard the plan-only report The interpolator broke both scalafix lint jobs. The guard keeps a preview failure from failing the query, which is the one thing plan-only mode promises not to do. --- .../apache/comet/rules/CometExecRule.scala | 39 ++++++++++++------- .../comet/rules/CometExecRuleSuite.scala | 2 +- 2 files changed, 25 insertions(+), 16 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 797c3d4eb59..1771c052564 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -20,6 +20,7 @@ package org.apache.comet.rules import scala.collection.mutable.ListBuffer +import scala.util.control.NonFatal import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder} @@ -132,8 +133,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 @@ -724,12 +723,30 @@ 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. + * If `plan` is one plan-only mode has not already described, build the Comet plan we would have + * executed and log it. Called from `_apply` in plan-only mode; the built plan is discarded. + * + * Nothing here may fail the query. Plan-only mode exists so that a workload can be assessed + * without taking on risk, and the preview rebuilds and rewrites a plan Spark has already + * prepared, so a plan shape it mishandles has to cost the report rather than the query. */ private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = { - val preview = buildPreview(plan, topLevel = true) - logWarning(s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") + try { + val executionId = Option( + session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)) + if (CometExecRule.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) + } } /** @@ -847,15 +864,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // the preview. Placed before `normalizePlan`/`RewriteJoin`/`tagUnsafePartialAggregates` // so their work is not wasted on the discarded outer pass. 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, - conf.adaptiveExecutionEnabled)) { - reportPlanOnlyCoverage(plan) - } + 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 c513a94e261..cbbae063b4a 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -1215,7 +1215,7 @@ class CometExecRuleSuite extends CometTestBase { 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 " + + "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, From 6ecc2041734a4c6147bb4c3d74b651bbe787d888 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 09:20:29 -0600 Subject: [PATCH 10/14] fix: keep plan-only mode from offloading Iceberg V2 writes `IcebergWriteStrategy` is registered with `injectPlannerStrategy`, not as a `Rule[SparkPlan]`, so it runs during physical planning ahead of both conversion rules and their plan-only short-circuit never reaches it. An Iceberg V2 write was still emitted as Comet's two-operator shape and executed natively, which is exactly what plan-only mode promises not to do. Neither `IcebergCommitExec` nor `IcebergWriteExec` extends `CometPlan`, so the plan-only suite's "no Comet operators" assertion could not see it. Decline the strategy in plan-only mode. The cost is that the report no longer describes the write as accelerated; that is documented alongside the existing native-planner caveat. Everything else registered in `CometSparkSessionExtensions.apply` is already inert once no Comet nodes exist: `CometPlanAdaptiveDynamicPruningFilters` only matches Comet scans, `CometReuseSubquery` creates no Comet nodes, and the two post-columnar rules only act on Comet operators. --- .../latest/understanding-comet-plans.md | 7 ++++++ .../comet/iceberg/IcebergWriteStrategy.scala | 9 +++++++ .../CometIcebergWriteDetectionSuite.scala | 25 +++++++++++++++++++ 3 files changed, 41 insertions(+) diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 21068553ec0..388348376fc 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -241,6 +241,13 @@ 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. +One piece of acceleration is left out in the other direction. Comet's split +Iceberg V2 write (`spark.comet.write.iceberg.splitOperator.enabled`, off by +default) is emitted by a Spark planner strategy rather than by the conversion +rules, so plan-only mode declines it outright to keep the write on Spark. An +Iceberg write Comet would have accelerated is therefore both executed and +reported as Spark. + 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. diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala index c501c900973..7e1dbfd7f37 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala @@ -35,6 +35,15 @@ import org.apache.comet.CometConf case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { override def apply(plan: LogicalPlan): Seq[SparkPlan] = { + // Plan-only mode must leave execution to Spark. Unlike `CometScanRule` and `CometExecRule`, + // this is a planner strategy: it runs during physical planning, before either of them, so + // their short-circuit does not cover it and a write would still be offloaded to Comet. + // Declining here costs the plan-only report the write operators - they are the one piece of + // acceleration it cannot describe - which is the documented trade for not executing them. + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(session.sessionState.conf)) { + return Nil + } + if (!CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.get(session.sessionState.conf)) { return Nil } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index 37ddb1872e8..6fe436adcd7 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -79,6 +79,31 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } } + test("plan-only mode leaves the write with Spark") { + withDetectionCatalog { dir => + createTable(dir, "plan_only", partitionSpec = "") + // `IcebergWriteStrategy` is a planner strategy, so it runs before `CometScanRule` and + // `CometExecRule` and is not covered by their plan-only short-circuit. Without its own + // guard the write is still offloaded, and neither exec it emits is a `CometPlan`, so the + // plan-only suite's "no Comet operators" assertion cannot see it. + // + // `withSQLConf` is declared to return `Unit` on Spark 3.4 and 3.5, so the plan is carried + // out of the block through a var rather than as the block's value. + 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 = "") From 7fa90eed4a61fd5a2fa7fe5de8b0f099e2065762 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 09:51:53 -0600 Subject: [PATCH 11/14] refactor: compose CometScanRule and CometExecRule into one CometRule The two were registered as separate rules, adjacently, in both the columnar and the query-stage-prep paths. Nothing ever ran between them, and neither is useful alone: CometExecRule seeds its native chain only from the nodes CometScanRule produces, so operator conversion over unconverted Spark scans converts nothing. Compose them so the ordering is an invariant of the code rather than of the registration order, and so callers that need the whole conversion have one entry point. Both rules keep their classes, files and tests. --- .../comet/CometSparkSessionExtensions.scala | 24 ++++------ .../org/apache/comet/rules/CometRule.scala | 47 +++++++++++++++++++ 2 files changed, 56 insertions(+), 15 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/rules/CometRule.scala diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 4b37d7b61a8..749a3d6ecb0 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/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.{CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} +import org.apache.comet.rules.{CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} import org.apache.comet.shims.ShimCometSparkSessionExtensions /** @@ -49,7 +49,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * 2. PlanSubqueries -- Spark creates SubqueryExec for scalar subqueries * 3. EnsureRequirements -- Spark inserts shuffles/sorts * 4. ApplyColumnarRulesAndInsertTransitions: - * a. preColumnarTransitions: CometScanRule, CometExecRule + * a. preColumnarTransitions: CometRule (CometScanRule then CometExecRule) * - CometExecRule.convertSubqueryBroadcasts converts SubqueryBroadcastExec to * CometSubqueryBroadcastExec for exchange reuse with Comet broadcasts * b. insertTransitions: ColumnarToRow/RowToColumnar added @@ -62,7 +62,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * {{{ * Initial plan: * PlanAdaptiveSubqueries: creates SubqueryAdaptiveBroadcastExec (SAB) for AQE DPP - * queryStagePreparationRules: CometScanRule, CometExecRule + * queryStagePreparationRules: CometRule (CometScanRule then CometExecRule) * - CometExecRule.convertSubqueryBroadcasts wraps SABs in * CometSubqueryAdaptiveBroadcastExec to prevent Spark's * PlanAdaptiveDynamicPruningFilters from replacing DPP with Literal.TrueLiteral @@ -75,7 +75,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * CometSubqueryBroadcastExec with BroadcastQueryStageExec for broadcast reuse * d. CometReuseSubquery -- deduplicates converted subqueries * 2. postStageCreationRules -> ApplyColumnarRulesAndInsertTransitions: - * a. preColumnarTransitions: CometScanRule, CometExecRule (no-ops, already converted) + * a. preColumnarTransitions: CometRule (no-op, already converted) * b. insertTransitions * c. postColumnarTransitions: RevertNativeForTransitionHeavyStages, * EliminateRedundantTransitions @@ -91,25 +91,19 @@ class CometSparkSessionExtensions with Logging with ShimCometSparkSessionExtensions { override def apply(extensions: SparkSessionExtensions): Unit = { - extensions.injectColumnar { session => CometScanColumnar(session) } - extensions.injectColumnar { session => CometExecColumnar(session) } + extensions.injectColumnar { session => CometColumnar(session) } // Pre-3.5 only: tag AQE DPP regions so the conversion rules below leave them Spark-native. - // Registered before CometScanRule/CometExecRule so tags are in place when conversion runs. + // 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 => CometScanRule(session) } - extensions.injectQueryStagePrepRule { session => CometExecRule(session) } + extensions.injectQueryStagePrepRule { session => CometRule(session) } injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters) injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) } } - case class CometScanColumnar(session: SparkSession) extends ColumnarRule { - override def preColumnarTransitions: Rule[SparkPlan] = CometScanRule(session) - } - - case class CometExecColumnar(session: SparkSession) extends ColumnarRule { - override def preColumnarTransitions: Rule[SparkPlan] = CometExecRule(session) + case class CometColumnar(session: SparkSession) extends ColumnarRule { + override def preColumnarTransitions: Rule[SparkPlan] = CometRule(session) override def postColumnarTransitions: Rule[SparkPlan] = { val rules = diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala new file mode 100644 index 00000000000..ac61a461d3e --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.rules + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.SparkPlan + +/** + * Comet's plan conversion pass: scan conversion followed by operator conversion. + * + * The two were previously registered as separate rules, adjacently, in both the columnar and the + * query-stage-prep paths. Nothing ever ran between them, and neither is useful on its own - + * [[CometExecRule]] seeds its native chain only from the nodes [[CometScanRule]] produces + * (`CometScanExec`, `CometBatchScanExec`, `CometContribScanMarker`), so operator conversion + * against unconverted Spark scans converts nothing. Composing them here makes that ordering an + * invariant of the code rather than of the registration order, and gives callers that need the + * whole conversion - rather than half of it - a single entry point. + * + * The two rules keep their own classes, files and tests; this only fixes how they are sequenced. + * `ruleName` in `spark.comet.explain.transformations` output is still each inner rule's own, + * since this delegates to their `apply`. + */ +case class CometRule(session: SparkSession) extends Rule[SparkPlan] { + + private val scanRule = CometScanRule(session) + private val execRule = CometExecRule(session) + + override def apply(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan)) +} From 8dd66d41adda0afa290e83a271fdf5c1881c10b0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 10:37:02 -0600 Subject: [PATCH 12/14] test: pin the invariant that motivates composing the rules --- .../comet/rules/CometExecRuleSuite.scala | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) 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 9d65dec834d..a746ce7c8e5 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -1086,4 +1086,30 @@ class CometExecRuleSuite extends CometTestBase { } } + test("operator conversion alone converts nothing without scan conversion") { + withTempPath { path => + createTestDataFrame.write.parquet(path.toString) + withTempView("test_data") { + spark.read.parquet(path.toString).createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT id, id * 2 as doubled FROM test_data WHERE id % 2 == 0") + assert(countOperators(sparkPlan, classOf[FileSourceScanExec]) == 1) + + // CometExecRule seeds its native chain only from the nodes CometScanRule produces, so on + // its own it converts nothing: the scan is untouched and every operator above it is + // refused for want of Arrow input. This is why the two are composed into `CometRule` + // rather than registered as independent rules that happen to run in the right order. + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true") { + val execOnly = CometExecRule(spark).apply(stripAQEPlan(sparkPlan)) + assert(countOperators(execOnly, classOf[FileSourceScanExec]) == 1) + assert( + stripAQEPlan(execOnly).collect { case p: CometNativeExec => p }.isEmpty, + s"operator conversion alone should convert nothing, got:\n$execOnly") + } + } + } + } + } From a80ae3fce500cd180dd87eecb3edeff2b1dd7671 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 19:34:08 -0600 Subject: [PATCH 13/14] refactor: simplify plan-only mode and trim its comments Share the post-columnar rule list between CometColumnar and the plan-only preview through CometRule.postColumnarRules, replacing the hand-synced copy. RevertNativeForTransitionHeavyStages takes a wholePlan flag in place of the separate applyToAllStages entry point. Remove the unused imports left in CometExecRule, which failed the scalafix lint, and restore the CometRule docstring and the "scan conversion must run before operator conversion" test from #6082 that the main merge dropped. Condense the comments, the user guide section and the config doc, inline single-use helpers, and factor the plan-only tests onto shared helpers. Check the Iceberg split-write conf before the plan-only conf so the common path does one lookup. --- .../latest/understanding-comet-plans.md | 73 +-- .../scala/org/apache/comet/CometConf.scala | 11 +- .../comet/CometSparkSessionExtensions.scala | 7 +- .../comet/iceberg/IcebergWriteStrategy.scala | 11 +- .../apache/comet/rules/CometExecRule.scala | 4 +- .../org/apache/comet/rules/CometRule.scala | 265 +++------- ...RevertNativeForTransitionHeavyStages.scala | 21 +- .../CometIcebergWriteDetectionSuite.scala | 9 +- .../comet/rules/CometExecRuleSuite.scala | 463 +++++++----------- 9 files changed, 283 insertions(+), 581 deletions(-) diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 388348376fc..6d2b3615103 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -210,57 +210,28 @@ 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. 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. 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, and nor does a plan AQE re-plans wholesale after a -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 -`create_plan` still counts as accelerated. Treat the percentage as an upper -bound. - -One piece of acceleration is left out in the other direction. Comet's split -Iceberg V2 write (`spark.comet.write.iceberg.splitOperator.enabled`, off by -default) is emitted by a Spark planner strategy rather than by the conversion -rules, so plan-only mode declines it outright to keep the write on Spark. An -Iceberg write Comet would have accelerated is therefore both executed and -reported as Spark. - -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. +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 diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 871c140e1ff..d279ce69fde 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -788,12 +788,11 @@ object CometConf extends ShimCometConf { 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.") + .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) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index b750baf7c09..f959d7bf9b5 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/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 /** @@ -108,10 +108,7 @@ class CometSparkSessionExtensions override def preColumnarTransitions: Rule[SparkPlan] = CometRule(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)) + val rules = CometRule.postColumnarRules(session) plan => rules.foldLeft(plan) { case (p, rule) => rule(p) } } } diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala index 7e1dbfd7f37..e8790079797 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala @@ -35,16 +35,11 @@ import org.apache.comet.CometConf case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { override def apply(plan: LogicalPlan): Seq[SparkPlan] = { - // Plan-only mode must leave execution to Spark. Unlike `CometScanRule` and `CometExecRule`, - // this is a planner strategy: it runs during physical planning, before either of them, so - // their short-circuit does not cover it and a write would still be offloaded to Comet. - // Declining here costs the plan-only report the write operators - they are the one piece of - // acceleration it cannot describe - which is the documented trade for not executing them. - if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(session.sessionState.conf)) { + if (!CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.get(session.sessionState.conf)) { return Nil } - - if (!CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.get(session.sessionState.conf)) { + // 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 } 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 afba9186090..54f88916245 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -20,7 +20,6 @@ package org.apache.comet.rules import scala.collection.mutable.ListBuffer -import scala.util.control.NonFatal import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder} @@ -47,9 +46,8 @@ 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, BroadcastExchangeLike, Exchange, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} 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._ diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala index 6dacaa8883c..4a5ac213710 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -35,143 +35,63 @@ 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)) + /** - * 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`. + * `executionId:canonicalPlanHash` keys already reported, LRU-bounded for long-lived drivers. */ - private class BoundedKeySet(limit: Int) { - private val keys: java.util.Map[String, java.lang.Boolean] = + 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() > limit - }) - - /** Adds `key`, returning true if it was not already present. */ - def add(key: String): Boolean = keys.put(key, java.lang.Boolean.TRUE) == null - } - - private val PLAN_ONLY_REPORTED_LIMIT = 1024 - - /** `executionId:planFingerprint` keys that plan-only mode has already reported. */ - private val planOnlyReportedPlans = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT) + eldest: java.util.Map.Entry[String, java.lang.Boolean]): Boolean = size() > 1024 + })) /** - * 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. + * 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` 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. - * - * 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. 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. - */ - private def isReapplication(plan: SparkPlan): Boolean = { - plan.exists(p => - p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] || - 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)) - } - /** * 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. 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: - * - * - 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. + * 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 = { - if (isQueryStage(plan, queryStagePrep, aqeEnabled) || isReapplication(plan)) { + // 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 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. + // Mark even when not reporting, so the final plan built from an empty re-plan is recognized. plan.setTagValue(PLAN_ONLY_REPORTED, ()) - // 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) && + // 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()}")) } } @@ -180,23 +100,20 @@ object CometRule { /** * Comet's plan conversion pass: scan conversion followed by operator conversion. * - * The two were previously registered as separate rules, adjacently, in both the columnar and the - * query-stage-prep paths. Nothing ever ran between them, and neither is useful on its own - - * [[CometExecRule]] seeds its native chain only from the nodes [[CometScanRule]] produces - * (`CometScanExec`, `CometBatchScanExec`, `CometContribScanMarker`), so operator conversion - * against unconverted Spark scans converts nothing. Composing them here makes that ordering an - * invariant of the code rather than of the registration order, and gives plan-only mode a single - * place to stand: one short-circuit, and one call to build the plan it reports on. + * Native scans come only from the nodes [[CometScanRule]] produces (`CometScanExec`, + * `CometBatchScanExec`, `CometContribScanMarker`), so [[CometExecRule]] must run after it. + * Running [[CometExecRule]] alone leaves scans on Spark's readers. Composing the two here makes + * that ordering part of the code instead of the order the rules are registered in, and gives + * callers that need the whole conversion a single entry point. * - * The two rules keep their own classes, files and tests; this only fixes how they are sequenced. - * `ruleName` in `spark.comet.explain.transformations` output is still each inner rule's own, - * since this delegates to their `apply`. + * `spark.comet.explain.transformations` logs each inner rule under its own `ruleName`, since this + * 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 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 - * [[CometRule.shouldReportPlanOnly]]. + * true for the `injectQueryStagePrepRule` instance, which sees the whole initial plan under + * AQE. Only plan-only reporting reads it. */ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) extends Rule[SparkPlan] { @@ -207,36 +124,22 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) override def apply(plan: SparkPlan): SparkPlan = { if (planOnlyApplies(plan)) { reportPlanOnlyCoverage(plan) - return plan + plan + } else { + convert(plan) } - convert(plan) } - /** Scan conversion followed by operator conversion: the plan Comet would execute. */ private def convert(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan)) - /** - * Whether plan-only mode governs this application. - * - * The guards mirror the ones the conversion rules apply to themselves, so that a plan they - * would have left alone anyway is not diverted into a report. In particular plan-only mode is - * scoped to `spark.comet.exec.enabled`: with exec disabled there is no operator conversion to - * describe, and Comet's columnar shuffle should keep being applied as usual. - */ + /** 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() && + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && isCometLoaded(conf) && !ShimCometStreaming.isStreamingPlan(plan) && CometConf.COMET_EXEC_ENABLED.get(conf) - /** - * If `plan` is one plan-only mode has not already described, build the Comet plan we would have - * executed and log it. Called from `apply` in plan-only mode; the built plan is discarded. - * - * Nothing here may fail the query. Plan-only mode exists so that a workload can be assessed - * without taking on risk, and the preview rebuilds and rewrites a plan Spark has already - * prepared, so a plan shape it mishandles has to cost the report rather than the query. - */ + /** Logs the Comet plan for `plan` unless already reported. Never fails the query. */ private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = { try { val executionId = Option( @@ -257,47 +160,26 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) } /** - * The plan Comet would have executed for `plan`. - * - * 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. + * 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 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. + * 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 reverted = RevertNativeForTransitionHeavyStages(session).applyToAllStages(withTransitions) - val preview = EliminateRedundantTransitions(session).apply(reverted) + val preview = CometRule + .postColumnarRules(session, wholePlan = true) + .foldLeft(withTransitions) { case (p, rule) => rule(p) } 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. + * `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 => @@ -306,22 +188,15 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) } 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 => 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. + * `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 => @@ -330,15 +205,9 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) } /** - * `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. + * 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 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 6a2d535b36d..bbdc6db6e57 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/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) @@ -78,19 +82,6 @@ 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/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index 6fe436adcd7..0d7d9935b22 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -82,13 +82,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("plan-only mode leaves the write with Spark") { withDetectionCatalog { dir => createTable(dir, "plan_only", partitionSpec = "") - // `IcebergWriteStrategy` is a planner strategy, so it runs before `CometScanRule` and - // `CometExecRule` and is not covered by their plan-only short-circuit. Without its own - // guard the write is still offloaded, and neither exec it emits is a `CometPlan`, so the - // plan-only suite's "no Comet operators" assertion cannot see it. - // - // `withSQLConf` is declared to return `Unit` on Spark 3.4 and 3.5, so the plan is carried - // out of the block through a var rather than as the block's value. + // 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) { 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 c2c67c58368..73793940a8f 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -1087,79 +1087,42 @@ class CometExecRuleSuite extends CometTestBase { } } + private val PLAN_ONLY_PREFIX = "[Comet plan-only]" + /** - * 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. + * Runs `f` over a Parquet table `tbl`. The source list must be set before the table is created, + * since `withParquetTable` bakes the V1/V2 choice into the temp view. */ - private def runPlanOnlyAndAssertReverted( - sql: String, - useV1: Boolean = true, - aqe: Boolean = true): Unit = { + private def withPlanOnlyTable( + aqe: Boolean = true, + planOnly: Boolean = true, + useV1: Boolean = true)(f: => Unit): 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") { - 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") - } - } + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> planOnly.toString) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl")(f) } } - 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") { - runPlanOnlyAndAssertReverted( - "SELECT _2, count(*) FROM tbl GROUP BY _2", - useV1 = useV1, - aqe = aqe) - } - } - - test("plan-only mode: scalar subquery is also reverted") { - 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") { - 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") - } + /** Plans `sql` in plan-only mode and asserts nothing was offloaded to native. */ + private def runPlanOnlyAndAssertReverted( + sql: String, + useV1: Boolean = true, + aqe: Boolean = true): Unit = { + withPlanOnlyTable(aqe = aqe, useV1 = useV1) { + 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") + val expectedScan = if (useV1) classOf[FileSourceScanExec] else classOf[BatchScanExec] + assert( + executed.exists(p => expectedScan.isInstance(p)), + s"expected ${expectedScan.getSimpleName}, got:\n$executed") } } - private val PLAN_ONLY_PREFIX = "[Comet plan-only]" - /** Runs `f` and returns the `[Comet plan-only]` reports that `CometRule` logged. */ private def capturePlanOnlyReports(f: => Unit): Seq[String] = { val appender = new LogAppender("Comet plan-only reports") @@ -1184,58 +1147,82 @@ class CometExecRuleSuite extends CometTestBase { .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. + /** + * Runs `query` normally, then in plan-only mode, and asserts the one report containing `marker` + * has the same coverage as the executed plan. Returns the executed plan. + */ + private def assertReportMatchesExecuted(query: String, marker: String): SparkPlan = { + val df = sql(query) + df.collect() + val plan = df.queryExecution.executedPlan + val executed = CometCoverageStats.forPlan(plan) + withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val reports = capturePlanOnlyReports(sql(query).collect()) + val matching = reports.filter(_.contains(marker)) + assert( + matching.size == 1, + s"expected one report containing '$marker', got:\n${reports.mkString("\n\n")}") + assert( + coverageOf(matching.head) == + (executed.cometOperators, executed.cometOperators + executed.sparkOperators), + s"report disagrees with the executed plan ($executed):\n${matching.head}") + } + plan + } + + 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") { + runPlanOnlyAndAssertReverted( + "SELECT _2, count(*) FROM tbl GROUP BY _2", + useV1 = useV1, + aqe = aqe) + } + } + + test("plan-only mode: scalar subquery is also reverted") { + 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") { + withPlanOnlyTable(planOnly = false) { + 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") + } + } + + // Subqueries are planned before the outer query, which must still get its own report. 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")) + withPlanOnlyTable(aqe = aqe) { + val reports = capturePlanOnlyReports { + spark.sql("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)").collect() } + // Filter appears only in the outer plan. + assert( + reports.exists(_.contains("Filter")), + s"no report describes the outer query:\n${reports.mkString("\n\n")}") + // One for the subquery, one for the outer query. + assert( + reports.size == 2, + s"expected two reports, 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") { - // SUM and MAX over a non-decimal column are safe to split in either direction across a - // stage boundary, which is what leaves `revertStageIfNeeded` free to revert the stage. - // COUNT is not - it declines Spark-partial to native-final - and an aggregate it declines - // anywhere in the stage suppresses the reversion this test is about. + withSQLConf(CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false") { + // AQE off: Spark applies the post-columnar rules to the whole plan once, as the preview does. + withPlanOnlyTable(aqe = false, planOnly = false) { + // SUM and MAX can be split across a stage boundary in either direction, which leaves the + // stage free to revert. COUNT cannot, and would block the reversion under test. val query = "SELECT _2, sum(_1), max(_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. var unreverted = 0 withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { val df = sql(query) @@ -1247,122 +1234,61 @@ class CometExecRuleSuite extends CometTestBase { 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) - - // Reversion has to actually fire, or the comparison below holds for a preview that - // never ran the post-columnar rules and the test proves nothing. `revertStageIfNeeded` - // declines to revert a stage that produces or consumes an aggregate buffer across a - // stage boundary, so the query above has to keep using aggregates that are safe in - // both directions. + val executed = CometCoverageStats.forPlan(assertReportMatchesExecuted(query, "")) + // Guard against a vacuous test: reversion must actually fire. assert( executed.cometOperators < unreverted, s"stage reversion did not fire, so this test is vacuous: $executed") - - // 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}") - } } } } } - // `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. + // These paths plan without a SQL execution ID and must still report once per query. 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" + withPlanOnlyTable(aqe = aqe) { + 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 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")) - } + 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. + // The outer report must count the subquery's converted operators. AQE off so the preview and + // the executed plan are the same single pass. 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(outer.head) == - (executed.cometOperators, executed.cometOperators + executed.sparkOperators), - s"outer report disagrees with the executed plan ($executed):\n${outer.head}") - } - } + withPlanOnlyTable(aqe = false, planOnly = false) { + val plan = assertReportMatchesExecuted( + "SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)", + marker = "Filter") + assert( + CometCoverageStats.forPlan(plan).cometOperators > 0, + "test query must be partly accelerated for the comparison to mean anything") } } - // 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. + // AQE can collapse a plan whose stage materializes empty into an empty relation that shares + // nothing with the reported plan. 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. + // `toRdd.count` runs without a SQL execution ID, like PySpark's `df.rdd`. (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. + // In the second shape `RemoveRedundantSorts` drops the root sort, so the mark lands on the + // join and the empty re-plan's root inherits the unmarked sort's tags. (shape, query, marker) <- Seq( ( "aggregate", @@ -1379,17 +1305,13 @@ class CometExecRuleSuite extends CometTestBase { 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(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(marker), s"the report does not describe the query:\n${reports.head}") @@ -1397,36 +1319,28 @@ 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")}") + withPlanOnlyTable(aqe = false) { + 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() } + // The `min` subquery, the `max` subquery and the outer query. + assert( + reports.size == 3, + s"expected three reports, got ${reports.size}:\n${reports.mkString("\n\n")}") } } + private val dppQuery = "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" + /** Registers a `fact` table partitioned on the join key plus a small `dim` table. */ private def withDppTables(f: => Unit): Unit = { withTempDir { dir => @@ -1451,70 +1365,37 @@ class CometExecRuleSuite extends CometTestBase { } } - // `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. + // The DPP build must be previewed below its BroadcastExchangeExec so that stage reversion fires + // as in real planning. Reversion is forced on 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. + // AQE off: the preview describes the pre-adaptive plan. 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 + val plan = assertReportMatchesExecuted(dppQuery, marker = "BroadcastHashJoin") // `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}") - } } } } - // 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. + // Under AQE the DPP build is planned mid-execution by `PlanAdaptiveDynamicPruningFilters`, + // after the outer query was reported, and must still get a report. 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()) + val reports = capturePlanOnlyReports(sql(dppQuery).collect()) assert( reports.size == 2, "expected a report for the outer query and one for the DPP build plan, got " + @@ -1526,18 +1407,10 @@ class CometExecRuleSuite extends CometTestBase { } } - // 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. + // Per-stage and per-re-optimization applications must not add reports. 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") { + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withPlanOnlyTable() { 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()) @@ -1548,29 +1421,43 @@ class CometExecRuleSuite extends CometTestBase { } } - test("operator conversion alone converts nothing without scan conversion") { + test("scan conversion must run before operator conversion") { withTempPath { path => createTestDataFrame.write.parquet(path.toString) withTempView("test_data") { spark.read.parquet(path.toString).createOrReplaceTempView("test_data") - val sparkPlan = - createSparkPlan(spark, "SELECT id, id * 2 as doubled FROM test_data WHERE id % 2 == 0") - assert(countOperators(sparkPlan, classOf[FileSourceScanExec]) == 1) + val query = "SELECT id, id * 2 as doubled FROM test_data WHERE id % 2 == 0" + + // One plan per rule application. Fallback reasons are recorded as tags on the Spark + // nodes, and CometNativeScan.isSupported declines a scan already carrying one, so + // reusing the plan the exec rule just refused would hold the second case down. + val forExecRule = stripAQEPlan(createSparkPlan(spark, query)) + val forCometRule = stripAQEPlan(createSparkPlan(spark, query)) + assert(countOperators(forExecRule, classOf[FileSourceScanExec]) == 1) + assert(countOperators(forCometRule, classOf[FileSourceScanExec]) == 1) - // CometExecRule seeds its native chain only from the nodes CometScanRule produces, so on - // its own it converts nothing: the scan is untouched and every operator above it is - // refused for want of Arrow input. This is why the two are composed into `CometRule` - // rather than registered as independent rules that happen to run in the right order. withSQLConf( CometConf.COMET_ENABLED.key -> "true", - CometConf.COMET_EXEC_ENABLED.key -> "true") { - val execOnly = CometExecRule(spark).apply(stripAQEPlan(sparkPlan)) - assert(countOperators(execOnly, classOf[FileSourceScanExec]) == 1) + CometConf.COMET_EXEC_ENABLED.key -> "true", + // Off by default, but pinned here: with it on, CometExecRule bridges the unconverted + // scan with a CometSparkToColumnarExec and converts the operators above it, which is a + // different path from the one under test. + CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "false") { + // CometExecRule builds its native plan up from the nodes CometScanRule produces, so on + // its own it leaves the scan on Spark's reader. + assert( + countOperators( + CometExecRule(spark).apply(forExecRule), + classOf[FileSourceScanExec]) == 1) + // CometRule runs both phases, in that order. This fails if the scan phase is ever + // reordered or dropped. assert( - stripAQEPlan(execOnly).collect { case p: CometNativeExec => p }.isEmpty, - s"operator conversion alone should convert nothing, got:\n$execOnly") + countOperators( + CometRule(spark).apply(forCometRule), + classOf[CometNativeScanExec]) == 1) } } } } + } From 31230882aa51aaf68c7ca7f3b1e54079bf5c05b1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 24 Sep 2026 03:57:14 -0600 Subject: [PATCH 14/14] fix: scope plan-only report dedupe to the query and detect AQE re-plans Render plan-only reports in the verbose format regardless of spark.comet.explain.format, so the coverage summary is never dropped. Tell AQE re-planning apart from initial preparation by whether the rule is running inside AdaptiveSparkPlanExec.reOptimize, instead of guessing from maxRows. An initially empty query is now reported under AQE, and a re-plan that keeps operators above an empty join is no longer reported twice. Deduplicate repeated subqueries within the query being prepared on the current thread rather than within a SQL execution ID, so queries planned without one (toRdd, df.rdd, executedPlan) no longer repeat reports. --- .../apache/comet/ExtendedExplainInfo.scala | 18 +++-- .../org/apache/comet/rules/CometRule.scala | 74 +++++++++++-------- .../comet/rules/CometExecRuleSuite.scala | 63 +++++++++++----- 3 files changed, 101 insertions(+), 54 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index 37535f80500..8d20b33868e 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -42,18 +42,24 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator { def generateExtendedInfo(plan: SparkPlan): String = { CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.get() match { case CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE => - // Generates the extended info in a verbose manner, printing each node along with the - // extended information in a tree display. - val planStats = new CometCoverageStats() - val outString = new StringBuilder() - generateTreeString(getActualPlan(plan), 0, Seq(), 0, outString, planStats) - s"${outString.toString()}\n$planStats" + generateVerboseInfo(plan) case CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_FALLBACK => // Generates the extended info as a list of fallback reasons getFallbackReasons(plan).mkString("\n").trim } } + /** + * The `verbose` format regardless of `spark.comet.explain.format`: each node along with its + * extended information in a tree display, followed by the coverage summary. + */ + def generateVerboseInfo(plan: SparkPlan): String = { + val planStats = new CometCoverageStats() + val outString = new StringBuilder() + generateTreeString(getActualPlan(plan), 0, Seq(), 0, outString, planStats) + s"${outString.toString()}\n$planStats" + } + def getFallbackReasons(plan: SparkPlan): Seq[String] = { fallbackReasons(plan).toSeq.sorted } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala index 4a5ac213710..fa2fb7dc32e 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -19,13 +19,14 @@ package org.apache.comet.rules +import scala.collection.mutable import scala.util.control.NonFatal import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule 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.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, ExecSubqueryExpression, InputAdapter, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, InsertAdaptiveSparkPlan, QueryStageExec} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange} import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery @@ -42,15 +43,13 @@ object CometRule { EliminateRedundantTransitions(session)) /** - * `executionId:canonicalPlanHash` keys already reported, LRU-bounded for long-lived drivers. + * Canonical hashes of the subquery plans reported for the query this thread is preparing. Spark + * prepares a query's subqueries synchronously, before the query itself, so the scope is reset + * when the query's own plan arrives. This works whether or not a SQL execution ID is set. */ - 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 - })) + private val reportedSubqueries = new ThreadLocal[mutable.Set[Int]] { + override def initialValue(): mutable.Set[Int] = mutable.Set.empty + } /** * Marks the root of a reported plan. Catalyst copies tags onto replacement nodes, so later @@ -58,17 +57,30 @@ object CometRule { */ private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported") + /** Where in Spark's planning the rule is running, read off the call stack. */ + private case class PlanningContext(replanning: Boolean, subquery: Boolean) + + private def planningContext(): PlanningContext = { + val frames = Thread.currentThread().getStackTrace + def within(cls: Class[_], method: String): Boolean = + frames.exists(f => f.getMethodName == method && f.getClassName == cls.getName) + PlanningContext( + replanning = within(classOf[AdaptiveSparkPlanExec], "reOptimize"), + // Subqueries are prepared by `PlanSubqueries` and `PlanDynamicPruningFilters` without AQE, + // and by `InsertAdaptiveSparkPlan` and `PlanAdaptiveDynamicPruningFilters` with it. + subquery = within(QueryExecution.getClass, "prepareExecutedPlan") || + within(classOf[InsertAdaptiveSparkPlan], "compileSubquery")) + } + /** * 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. + * is prepared as a plan of its own and gets its own report. */ private[comet] def shouldReportPlanOnly( - executionId: Option[String], plan: SparkPlan, queryStagePrep: Boolean, aqeEnabled: Boolean): Boolean = { @@ -80,19 +92,25 @@ object CometRule { p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] || p.getTagValue(PLAN_ONLY_REPORTED).isDefined) if (isQueryStage || isReapplication) { + // Execution is under way, so any subquery prepared from here on belongs to a new scope. + reportedSubqueries.get().clear() false } else { - // Mark even when not reporting, so the final plan built from an empty re-plan is recognized. + // Mark even when not reporting, so the final plan built from a 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()}")) + val context = planningContext() + if (context.replanning) { + // AQE re-plans a query mid-execution, for example to an empty relation once a stage + // materializes empty, and the result can share no nodes or stages with the reported plan. + false + } else if (context.subquery) { + // A subquery referenced twice is prepared twice, differing only in expression IDs. + // Tags are not part of the canonical form. + reportedSubqueries.get().add(plan.canonicalized.hashCode()) + } else { + reportedSubqueries.get().clear() + true + } } } } @@ -142,16 +160,10 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) /** 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)) { + if (CometRule.shouldReportPlanOnly(plan, queryStagePrep, conf.adaptiveExecutionEnabled)) { val preview = buildPreview(plan, topLevel = true) logWarning( - s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") + s"[Comet plan-only]\n${new ExtendedExplainInfo().generateVerboseInfo(preview)}") } } catch { case NonFatal(e) => 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 73793940a8f..8ad1884acca 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -1288,7 +1288,9 @@ class CometExecRuleSuite extends CometTestBase { "collect" -> (df => df.collect()), "toRdd.count" -> (df => df.queryExecution.toRdd.count())) // In the second shape `RemoveRedundantSorts` drops the root sort, so the mark lands on the - // join and the empty re-plan's root inherits the unmarked sort's tags. + // join and the empty re-plan's root inherits the unmarked sort's tags. In the third the + // global aggregate survives the empty join, so the re-plan is not itself empty. The fourth + // is empty from the start, which must not be mistaken for a re-plan. (shape, query, marker) <- Seq( ( "aggregate", @@ -1300,7 +1302,14 @@ class CometExecRuleSuite extends CometTestBase { |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")) + "SortMergeJoin"), + ( + "global aggregate over a join", + """SELECT count(*) 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""".stripMargin, + "SortMergeJoin"), + ("initially empty", "SELECT id FROM range(0) DISTRIBUTE BY id", "Exchange")) } { test(s"plan-only mode: one report when the plan becomes empty ($shape, AQE=$aqe, $action)") { withSQLConf( @@ -1319,22 +1328,42 @@ class CometExecRuleSuite extends CometTestBase { } } - test("plan-only mode: a subquery referenced twice is reported once") { - withPlanOnlyTable(aqe = false) { - 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() + for { + aqe <- Seq(true, false) + (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: a subquery referenced twice is reported once (AQE=$aqe, $action)") { + withPlanOnlyTable(aqe = aqe) { + val reports = capturePlanOnlyReports { + runIt(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)) + } + // The `min` subquery, the `max` subquery and the outer query. + assert( + reports.size == 3, + s"expected three reports, got ${reports.size}:\n${reports.mkString("\n\n")}") + } + } + } + + test("plan-only mode: the report does not depend on spark.comet.explain.format") { + withSQLConf( + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_FALLBACK) { + withPlanOnlyTable() { + val reports = capturePlanOnlyReports(sql("SELECT _1 + 1 FROM tbl").collect()) + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + val (accelerated, eligible) = coverageOf(reports.head) + assert( + accelerated > 0 && accelerated == eligible, + s"unexpected coverage:\n${reports.head}") } - // The `min` subquery, the `max` subquery and the outer query. - assert( - reports.size == 3, - s"expected three reports, got ${reports.size}:\n${reports.mkString("\n\n")}") } }