From ab5c3915752d83c01154ab99c8bfb99e37f56a61 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 27 Aug 2026 16:15:11 -0600 Subject: [PATCH 1/5] feat: add spark.comet.explain.planOnly.enabled Report the Comet plan Comet would have executed for a query, without offloading any of it to Comet, so a workload can be assessed without changing how it runs. The conversion rules return the plan untouched while the mode is on, and the report is built afterwards from the plan Spark executed, by a query execution listener. Reporting once the query is over rather than while it is being planned is what keeps this small: Spark applies a planner rule many times for one query - once per query stage and once per adaptive re-optimization under AQE, and separately for every subquery it prepares - and telling those applications apart takes a state machine. A listener fires once per action, holding the finished plan, so there is nothing to tell apart. Building the report from the plan that ran means undoing the part of Spark's preparation that follows the conversion rules: the adaptive wrappers, the codegen wrappers, the columnar transitions, and the exchange reuse. It then replays transition insertion and Comet's post-columnar rules, so a stage Comet would have handed back to Spark is reported as handed back. Tests assert both properties that matter: the executed plan holds no Comet operator, and the report's coverage equals CometCoverageStats for the plan Comet really executes - for an aggregate, a shuffled join, a scalar subquery, a DPP subquery and a reverted stage, with AQE on and off. --- .../latest/understanding-comet-plans.md | 43 ++ .../scala/org/apache/comet/CometConf.scala | 13 + .../apache/comet/rules/CometExecRule.scala | 10 +- .../apache/comet/rules/CometPlanOnly.scala | 243 ++++++++++++ .../apache/comet/rules/CometScanRule.scala | 6 +- ...RevertNativeForTransitionHeavyStages.scala | 12 + .../comet/rules/CometPlanOnlySuite.scala | 375 ++++++++++++++++++ 7 files changed, 700 insertions(+), 2 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala create mode 100644 spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index a600615b4b3..13305aefb97 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -199,6 +199,49 @@ from both operator counts: Counting reused exchanges once is tracked as item 3 of [#5203](https://github.com/apache/datafusion-comet/issues/5203). +### `spark.comet.explain.planOnly.enabled` + +When enabled, Comet leaves every query for Spark to execute and, once the query +finishes, logs the Comet plan it would have executed along with a coverage +summary. Use this to evaluate how much of a workload Comet would accelerate +without changing how that workload runs. + +Comet's conversion rules return the plan untouched while this is on, so the plan +Spark executes is the plan it would have built with Comet switched off. The +report is built afterwards from the plan that ran, and then discarded. + +The log line is prefixed with `[Comet plan-only]` and carries the same annotated +plan and summary that `spark.comet.explain.format=verbose` produces for a real +Comet plan. It is written once per action, so a plan that is built but never +executed — `df.explain()`, or reading `queryExecution.executedPlan` — is not +reported, and a DataFrame collected twice is reported twice. + +The preview goes through the whole Comet planning sequence rather than operator +conversion alone: Spark's columnar transitions are inserted and Comet's +post-columnar rules (`RevertNativeForTransitionHeavyStages`, +`EliminateRedundantTransitions`) are applied, so a stage Comet would have handed +back to Spark for having too many transitions is reported as handed back. + +Two things to keep in mind when reading the percentage: + +- The estimate reflects Scala-side conversion only. The 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 the report describes the plan AQE settled on, not the query as + written, because that is the plan Comet would have been asked to run. If a + stage materialized empty and AQE replaced the query with an empty relation, + that is what gets reported. The transition count can also be slightly lower + than a real Comet run's, because Spark inserts transitions one query stage at + a time whereas the report is produced from the whole plan in one pass. The + operator counts behind the percentage are not affected. + +A reused exchange is expanded in the report, so its subtree appears once per +reference. That matches how coverage is counted for a real Comet plan; see the +note under `spark.comet.explain.format` above. + +The config requires `spark.comet.exec.enabled=true`. With Comet exec disabled +the rule that arranges the report does not run. + ### `spark.comet.explain.native.enabled` When enabled, each executor task logs the DataFusion plan it executes, diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d8fe5b69890..62c0656628e 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -661,6 +661,19 @@ 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 leaves the query for Spark to execute and afterwards logs the " + + "Comet plan it would have executed, with a coverage summary. Use this to evaluate how " + + "much of a workload Comet would accelerate without changing execution. The estimate is " + + "Scala-side only; the plan is never handed to DataFusion, so native planning failures " + + "are not surfaced and the acceleration percentage can be optimistic. Reported once per " + + "action, so a plan built but never executed is not reported. Requires " + + "`spark.comet.exec.enabled=true`. Disabled by default.") + .booleanConf + .createWithDefault(false) + val COMET_EXPLAIN_FALLBACK_LOG_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.explain.fallback.log.enabled") .withAlternative("spark.comet.logFallbackReasons.enabled") 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 c69602fc80b..7f6f0915376 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -574,6 +574,14 @@ case class CometExecRule(session: SparkSession) } override def apply(plan: SparkPlan): SparkPlan = { + // Plan-only mode: leave the plan alone, so Spark executes exactly what it would with Comet + // off, and arrange for the Comet plan to be reported once the query is over. See + // `CometPlanOnly` for why the report is not built here. + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) { + CometPlanOnly.register(session) + return plan + } + val newPlan = _apply(plan) if (showTransformations && !newPlan.fastEquals(plan)) { logInfo(s""" @@ -584,7 +592,7 @@ case class CometExecRule(session: SparkSession) newPlan } - private def _apply(plan: SparkPlan): SparkPlan = { + private[rules] def _apply(plan: SparkPlan): SparkPlan = { // We shouldn't transform Spark query plan if Comet is not loaded. if (!isCometLoaded(conf)) return plan diff --git a/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala b/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala new file mode 100644 index 00000000000..b81d096fba8 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala @@ -0,0 +1,243 @@ +/* + * 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 scala.util.control.NonFatal + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, CommandResultExec, ExecSubqueryExpression, InputAdapter, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} +import org.apache.spark.sql.execution.command.ExecutedCommandExec +import org.apache.spark.sql.execution.datasources.v2.V2CommandExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec} +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.{CometConf, ExtendedExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.isCometLoaded + +/** + * Plan-only mode: report the Comet plan Comet would have executed for a query, without offloading + * any of it to Comet. See `spark.comet.explain.planOnly.enabled`. + * + * The conversion rules leave the plan alone while the mode is on, so Spark plans and executes the + * query exactly as it would with Comet switched off. The report is built afterwards, from the + * plan Spark actually executed, and thrown away. Comet code therefore cannot reach the query: not + * by planning it, and not by failing while describing it. + * + * Reporting once the query is over, rather than while it is being planned, is what keeps this + * simple. Spark applies a planner rule many times for one query - once per query stage and once + * per adaptive re-optimization under AQE, and separately for every subquery it prepares - and + * telling those applications apart takes a good deal of bookkeeping. A query execution listener + * fires once per action, holding the finished plan, so there is nothing to tell apart. + */ +object CometPlanOnly extends Logging { + + private val REPORT_PREFIX = "[Comet plan-only]" + + /** + * Sessions that already have a listener registered. Weakly held so a session that goes away is + * not kept alive by this, and so a long-lived driver retains no more state than the sessions it + * is running. + */ + private val registeredSessions: java.util.Set[SparkSession] = + java.util.Collections.synchronizedSet( + java.util.Collections.newSetFromMap( + new java.util.WeakHashMap[SparkSession, java.lang.Boolean]())) + + /** + * Registers this session's plan-only listener, if it does not have one yet. + * + * Called from `CometExecRule` rather than at session creation so that a session never carries a + * listener unless plan-only mode is actually used, and so the config can be turned on part way + * through a session. + */ + def register(session: SparkSession): Unit = { + if (registeredSessions.add(session)) { + session.listenerManager.register(new CometPlanOnlyListener) + logInfo(s"$REPORT_PREFIX registered a plan-only reporter for this session") + } + } + + /** + * Logs the Comet plan Comet would have executed for `qe`. + * + * Nothing here may fail the query, which has finished by this point but whose action would + * still see an exception thrown from a listener. Plan-only mode exists to let a workload be + * assessed without taking on risk, so a plan shape the preview mishandles has to cost the + * report rather than the query. + */ + private def report(qe: QueryExecution): Unit = { + val session = qe.sparkSession + // The listener bus thread has no active session, and the conversion rules read their configs + // from the active one. Without this the preview would be built from default config values. + val previous = SparkSession.getActiveSession + SparkSession.setActiveSession(session) + try { + val conf = session.sessionState.conf + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && isCometLoaded(conf) && + CometConf.COMET_EXEC_ENABLED.get(conf) && !isMetadataOnly(qe.executedPlan)) { + val preview = previewOf(session, qe.executedPlan) + logWarning(s"$REPORT_PREFIX\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") + } + } catch { + case NonFatal(e) => + logWarning(s"$REPORT_PREFIX could not build a coverage report for this query", e) + } finally { + previous match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + } + } + + /** + * Whether `plan` only touches metadata - `CREATE VIEW`, `SHOW TABLES`, `SET`. + * + * There is nothing to accelerate in one, and a session runs enough of them that reporting each + * as 0% would bury the reports worth reading. A command that carries a query below it - `INSERT + * ... SELECT`, `CREATE TABLE AS SELECT`, a V2 append - has that query as a child and is + * reported. + */ + private def isMetadataOnly(plan: SparkPlan): Boolean = plan match { + case _: ExecutedCommandExec | _: CommandResultExec => true + case command: V2CommandExec => command.children.isEmpty + case _ => false + } + + /** + * The plan Comet would have executed for `plan`, which Spark has finished preparing and + * running. + * + * Conversion is only the first half of Comet planning. 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. + * + * `RevertNativeForTransitionHeavyStages` is applied with `applyToAllStages` because this holds + * a whole plan, whereas under AQE Spark hands that rule one stage at a time. + */ + private def previewOf(session: SparkSession, plan: SparkPlan): SparkPlan = { + val prepared = previewSubqueriesOf(session, stripPreparation(plan)) + val converted = CometExecRule(session)._apply(CometScanRule(session)._apply(prepared)) + val withTransitions = + ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted) + val reverted = RevertNativeForTransitionHeavyStages(session).applyToAllStages(withTransitions) + EliminateRedundantTransitions(session).apply(reverted) + } + + /** + * `plan` as the conversion rules would have seen it, with everything Spark added after them + * removed: the adaptive wrappers, the whole-stage codegen wrappers, and the columnar + * transitions. + * + * Taking the plan Spark executed and undoing this much of its preparation is what buys the + * accuracy this mode needs. The alternative, describing the plan as it stood before + * preparation, describes a plan AQE may have replanned beyond recognition: stages coalesced, + * joins switched from sort merge to broadcast, an empty side pruned away. + */ + private def stripPreparation(plan: SparkPlan): SparkPlan = plan match { + // Under AQE the executed plan is a wrapper holding the plan AQE settled on. Its query stages + // hold their own plans off to one side, out of `children`, so an ordinary transform would not + // reach into them. + case adaptive: AdaptiveSparkPlanExec => stripPreparation(adaptive.executedPlan) + case stage: QueryStageExec => stripPreparation(stage.plan) + // A runtime partition-coalescing wrapper over a shuffle stage. It has no counterpart in a plan + // that has not been through AQE, and the conversion rules judge a shuffle by the exchange, so + // it goes with the stage it wraps. + case read: AQEShuffleReadExec => stripPreparation(read.child) + // `ReuseExchangeAndSubquery` is the last thing Spark's preparation does, after the columnar + // rules, so in a real Comet run the exchange behind a `ReusedExchangeExec` has already been + // converted. Here it has not, and the wrapper is a leaf as far as a transform is concerned, so + // conversion would never reach the subtree while the coverage count - which unwraps the + // wrapper - still counts every operator in it as Spark. Undo the reuse and let both copies + // convert, which is what the counts of a real Comet run reflect. + case reused: ReusedExchangeExec => stripPreparation(reused.child) + case WholeStageCodegenExec(child) => stripPreparation(child) + case InputAdapter(child) => stripPreparation(child) + case ColumnarToRowExec(child) => stripPreparation(child) + case RowToColumnarExec(child) => stripPreparation(child) + case other => other.withNewChildren(other.children.map(stripPreparation)) + } + + /** + * `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. Spark prepares a subquery + * as a plan in its own right and substitutes it into the outer plan, so leaving those plans + * untouched here would report every subquery operator as un-accelerated Spark and understate + * coverage against what Comet really executes. + */ + private def previewSubqueriesOf(session: SparkSession, plan: SparkPlan): SparkPlan = { + plan.transformAllExpressions { case subquery: ExecSubqueryExpression => + subquery.withNewPlan(previewSubquery(session, subquery.plan)) + } + } + + private def previewSubquery( + session: SparkSession, + subquery: BaseSubqueryExec): BaseSubqueryExec = + subquery match { + // Reuse bookkeeping: the plan to preview is one level further down. + case reused: ReusedSubqueryExec => + reused.copy(child = previewSubquery(session, reused.child)) + case other => + other + .withNewChildren(Seq(previewSubqueryPlan(session, other.child))) + .asInstanceOf[BaseSubqueryExec] + } + + /** + * Preview the plan behind a subquery, which Spark prepared as a plan in its own right. + * + * A dynamic partition pruning 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 would leave its child a stage bounded at the top by + * the exchange, which stops the reversion firing, and the report would then count 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 previewSubqueryPlan(session: SparkSession, plan: SparkPlan): SparkPlan = + plan match { + case stage: QueryStageExec => previewSubqueryPlan(session, stage.plan) + case exchange: BroadcastExchangeExec => + exchange.withNewChildren(Seq(previewSubqueryPlan(session, exchange.child))) + case other => previewOf(session, other) + } + + /** The listener that reports one plan per query. */ + private class CometPlanOnlyListener extends QueryExecutionListener { + + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + report(qe) + } + + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = { + // Report anyway: the query was planned, which is all this mode describes. + report(qe) + } + } +} 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 a524da3af92..cdba25576b2 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,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 alone, so Spark scans exactly as it would with Comet off. + // `CometPlanOnly` calls `_apply` on a copy of the plan Spark executed, once the query is over. + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) return plan + val newPlan = _apply(plan) if (showTransformations && !newPlan.fastEquals(plan)) { logInfo(s""" @@ -74,7 +78,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 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..c9107fe3509 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -52,6 +52,18 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession) } } + /** + * 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. A caller holding a whole plan + * that has not been split into stages - the plan-only preview in `CometPlanOnly` - needs every + * shuffle boundary visited to see the reversions that the real per-stage applications made. + */ + private[rules] def applyToAllStages(plan: SparkPlan): SparkPlan = { + if (!enabled) plan else applyForNonAQE(plan) + } + private def applyForAQE(plan: SparkPlan): SparkPlan = { plan match { case _: BroadcastExchangeLike => plan diff --git a/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala new file mode 100644 index 00000000000..cce9d712b62 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala @@ -0,0 +1,375 @@ +/* + * 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.logging.log4j.Level +import org.apache.spark.CometListenerBusUtils +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.comet.CometPlan +import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan, SubqueryBroadcastExec} +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometCoverageStats} + +/** + * Tests for plan-only mode: `spark.comet.explain.planOnly.enabled`. + * + * Two properties matter. The query must run exactly as it would with Comet off, which is checked + * by asserting the executed plan holds no Comet operator. And the report must describe the plan + * Comet would really have executed, which is checked by running the same query with Comet enabled + * and comparing the report's coverage against `CometCoverageStats` for the plan that ran. + */ +class CometPlanOnlySuite extends CometTestBase { + + private val PLAN_ONLY_PREFIX = "[Comet plan-only]" + + private val reporterLogger = CometPlanOnly.getClass.getName.stripSuffix("$") + + /** + * Runs `f` and returns the plan-only reports logged for it. + * + * The report is written from the listener bus, so the bus has to be drained before the appender + * is installed - or a report for an action that ran earlier, the fixture's view creation say, + * lands in the window - and again before the log is read, or the reports for `f`'s own actions + * may not have been written yet. + */ + private def capturePlanOnlyReports(f: => Unit): Seq[String] = { + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + val appender = new LogAppender("Comet plan-only reports") + withLogAppender(appender, loggerNames = Seq(reporterLogger), level = Some(Level.WARN)) { + f + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } + 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 transition count in a plan-only report. */ + private def transitionsOf(report: String): Int = { + """contains (\d+) transitions""".r + .findFirstMatchIn(report) + .map(_.group(1).toInt) + .getOrElse(fail(s"report has no transition count:\n$report")) + } + + private def planOnlyConf(aqe: Boolean, useV1: Boolean): Seq[(String, String)] = Seq( + 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") + + // `collect` here is `AdaptiveSparkPlanHelper.collect`, which descends into query stages; a plain + // `SparkPlan.collect` stops at them, because a stage holds its plan outside `children`. + private def cometOperatorsOf(plan: SparkPlan): Seq[SparkPlan] = + collect(plan) { case p: CometPlan => p } + + for { + useV1 <- Seq(true, false) + aqe <- Seq(true, false) + } { + test(s"the query runs on Spark (${if (useV1) "V1" else "V2"} scan, AQE=$aqe)") { + // The source list has to be set before the fixture reads the table: `withParquetTable` + // resolves the relation through `spark.read` and registers the result as a temp view, so + // changing it afterwards leaves a V1 relation in place and the V2 case would not be covered. + withSQLConf(planOnlyConf(aqe, useV1): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val query = "SELECT _2, count(*) FROM tbl GROUP BY _2" + + // Sanity check on the fixture: with the config off Comet does accelerate this query, so + // the assertion below is about plan-only mode and not about an unrelated fallback. Only + // for V1: Comet declines a plain V2 Parquet scan ("Unsupported scan: ParquetScan"), so + // the V2 case is here to cover `CometScanRule`'s V2 branch, not to show acceleration. + if (useV1) { + val normal = sql(query) + normal.collect() + assert(cometOperatorsOf(normal.queryExecution.executedPlan).nonEmpty) + } + + withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val df = sql(query) + val reports = capturePlanOnlyReports(df.collect()) + val executed = df.queryExecution.executedPlan + assert( + cometOperatorsOf(executed).isEmpty, + s"plan-only mode left Comet operators in the executed plan:\n$executed") + // The fixture must exercise the scan path the test name claims. + val scans = collect(executed) { + case p: FileSourceScanExec => p + case p: BatchScanExec => p + } + if (useV1) { + assert( + scans.exists(_.isInstanceOf[FileSourceScanExec]), + s"expected a V1 scan:\n$executed") + } else { + assert( + scans.exists(_.isInstanceOf[BatchScanExec]), + s"expected a V2 scan:\n$executed") + } + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + } + } + } + } + } + + test("the config off leaves Comet running the query") { + withSQLConf( + planOnlyConf(aqe = true, useV1 = true) :+ + (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "false"): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val df = sql("SELECT _2, count(*) FROM tbl GROUP BY _2") + val reports = capturePlanOnlyReports(df.collect()) + assert(cometOperatorsOf(df.queryExecution.executedPlan).nonEmpty) + assert(reports.isEmpty, s"expected no report, got:\n${reports.mkString("\n\n")}") + } + } + } + + // One report per action, whatever Spark does to the plan in between. Under AQE one query reaches + // the conversion rules once per query stage and once per adaptive re-optimization on top of the + // initial planning, and a query with subqueries reaches them once per subquery as well; none of + // that is visible from a query execution listener. + for (aqe <- Seq(true, false)) { + test(s"one report per action for a multi-stage query with a subquery (AQE=$aqe)") { + withSQLConf( + planOnlyConf(aqe, useV1 = true) ++ Seq( + // Force a shuffled join so the plan has more than one shuffle boundary. + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true"): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + // The filter has to leave rows behind: an empty stage lets AQE replace the whole plan + // with an empty relation, and the report would then describe that instead of the query. + val query = "SELECT a._2, count(*) FROM tbl a JOIN tbl b ON a._1 = b._2 " + + "WHERE a._1 >= (SELECT min(_2) FROM tbl) GROUP BY a._2 ORDER BY 1" + val reports = capturePlanOnlyReports(sql(query).collect()) + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + // The report describes the whole query, subquery included. + assert(reports.head.contains("HashAggregate"), s"report:\n${reports.head}") + } + } + } + } + + test("two actions on the same query are reported twice") { + withSQLConf( + planOnlyConf(aqe = true, useV1 = true) :+ + (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true"): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val df = sql("SELECT _2, count(*) FROM tbl GROUP BY _2") + val reports = capturePlanOnlyReports { + df.collect() + df.collect() + } + assert(reports.size == 2, s"expected two reports, got:\n${reports.mkString("\n\n")}") + } + } + } + + // `df.rdd` - the path PySpark's `df.rdd` takes through `Dataset.javaToPython` - plans a second + // query of its own and runs it under its own execution id, so it is reported too, once. + test("an RDD action is reported once") { + withSQLConf( + planOnlyConf(aqe = true, useV1 = 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 _2, count(*) FROM tbl GROUP BY _2").rdd.count() + } + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + assert(reports.head.contains("HashAggregate"), s"report:\n${reports.head}") + } + } + } + + // A session runs plenty of metadata-only statements, and one 0% report each would bury the + // reports worth reading. + test("a metadata-only command is not reported") { + withSQLConf( + planOnlyConf(aqe = true, useV1 = true) :+ + (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true"): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + withTempView("v") { + val reports = capturePlanOnlyReports { + sql("CREATE OR REPLACE TEMP VIEW v AS SELECT _1 FROM tbl") + sql("SHOW TABLES").collect() + } + assert(reports.isEmpty, s"expected no report, got:\n${reports.mkString("\n\n")}") + } + } + } + } + + /** + * Asserts that the report for `query` in plan-only mode agrees with the coverage of the plan + * Comet really executes for it. + */ + private def assertReportMatchesRealPlan( + query: String, + compareTransitions: Boolean = true): Unit = { + val df = sql(query) + df.collect() + val executed = CometCoverageStats.forPlan(df.queryExecution.executedPlan) + assert( + executed.cometOperators > 0, + "the query must be partly accelerated for the comparison to mean anything:\n" + + df.queryExecution.executedPlan) + + 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}") + if (compareTransitions) { + assert( + transitionsOf(reports.head) == executed.transitions, + s"report disagrees with the executed plan on transitions ($executed):\n${reports.head}") + } + } + } + + for { + aqe <- Seq(true, false) + (shape, query) <- Seq( + "aggregate" -> "SELECT _2, count(*), sum(_1) FROM tbl GROUP BY _2", + "shuffled join" -> "SELECT a._2, count(*) FROM tbl a JOIN tbl b ON a._1 = b._2 GROUP BY a._2", + "scalar subquery" -> "SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)") + } { + test(s"report coverage matches the plan Comet executes ($shape, AQE=$aqe)") { + withSQLConf( + planOnlyConf(aqe, useV1 = true) :+ + (SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1"): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + assertReportMatchesRealPlan(query) + } + } + } + } + + // A stage Comet would have handed back to Spark for having too many transitions must be reported + // as handed back. Reversion is forced on here, and Comet project execution off so that the + // aggregate leaves transitions behind for the rule to count. + for (aqe <- Seq(true, false)) { + test(s"report accounts for post-columnar stage reversion (AQE=$aqe)") { + withSQLConf( + planOnlyConf(aqe, useV1 = true) ++ Seq( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0", + 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" + + // With reversion off Comet accelerates part of this plan, so a report that skipped the + // post-columnar rules would not match the executed plan below. + withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val df = sql(query) + df.collect() + assert(CometCoverageStats.forPlan(df.queryExecution.executedPlan).cometOperators > 0) + } + + // Transitions are not compared under AQE: Spark inserts them one query stage at a + // time, so a stage that reversion handed back to Spark keeps a transition above the + // stage boundary below it that a single pass over the flattened plan does not produce. + // The operator counts, which are what the coverage percentage is built from, do match. + assertReportMatchesRealPlan(query, compareTransitions = !aqe) + } + } + } + } + + /** 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) + } + } + + // A dynamic partition pruning subquery is prepared by `PlanDynamicPruningFilters`, which prepares + // the build plan and only then wraps it in the broadcast exchange, so the stage the post-columnar + // rules judged is the exchange's child. Reversion is forced on so that getting that boundary + // wrong changes the number. + test("report coverage matches the plan Comet executes (DPP subquery)") { + withSQLConf( + planOnlyConf(aqe = false, useV1 = true) ++ Seq( + SQLConf.DYNAMIC_PARTITION_PRUNING_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() + // `exists` walks children only, and a DPP subquery hangs off the scan's expressions. + assert( + df.queryExecution.executedPlan.collectWithSubqueries { case p: SubqueryBroadcastExec => + p + }.nonEmpty, + s"the query must produce a DPP subquery:\n${df.queryExecution.executedPlan}") + + assertReportMatchesRealPlan(query) + } + } + } + + // A query AQE re-plans wholesale once a stage materializes empty. The report describes the plan + // AQE settled on, and there is still exactly one of them. + test("an adaptive query that collapses to nothing is reported once") { + withSQLConf( + planOnlyConf(aqe = true, useV1 = true) ++ Seq( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true"): _*) { + val query = "SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2" + val reports = capturePlanOnlyReports(spark.sql(query).collect()) + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + } + } +} From fa77513059655f70ded212d0036214e6a0a36537 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 27 Aug 2026 16:43:50 -0600 Subject: [PATCH 2/5] ci: register CometPlanOnlySuite in the PR build workflows --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index b7c88254c93..cbe5d1fd00d 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -358,6 +358,7 @@ jobs: org.apache.comet.rules.CometScanContribSuite org.apache.comet.rules.CometScanSchemeFallbackSuite org.apache.comet.rules.CometExecRuleSuite + org.apache.comet.rules.CometPlanOnlySuite org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite org.apache.spark.sql.CometTPCDSQuerySuite org.apache.spark.sql.CometTPCDSQueryTestSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 8210f91b7f9..83389fff373 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -174,6 +174,7 @@ jobs: org.apache.comet.rules.CometScanContribSuite org.apache.comet.rules.CometScanSchemeFallbackSuite org.apache.comet.rules.CometExecRuleSuite + org.apache.comet.rules.CometPlanOnlySuite org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite org.apache.spark.sql.CometTPCDSQuerySuite org.apache.spark.sql.CometTPCDSQueryTestSuite From 8ccff51abc7afdd2d4d17aac3677c96e3ddeb3d2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 9 Sep 2026 10:08:07 -0600 Subject: [PATCH 3/5] fix: keep plan-only reports tied to their query's own settings Snapshots the plan-time settings on the plan so an asynchronous report is not decided by session state the caller has since restored, preserves a reused exchange's output IDs when undoing the reuse for the preview, and corrects the per-action claim: reporting is per SQL execution, so RDD actions are outside it. --- .../latest/understanding-comet-plans.md | 16 ++- .../apache/comet/rules/CometExecRule.scala | 4 +- .../apache/comet/rules/CometPlanOnly.scala | 121 +++++++++++++++++- .../comet/rules/CometPlanOnlySuite.scala | 110 +++++++++++++++- 4 files changed, 235 insertions(+), 16 deletions(-) diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index dd98fb2ff2e..0c40a4313af 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -212,9 +212,19 @@ report is built afterwards from the plan that ran, and then discarded. The log line is prefixed with `[Comet plan-only]` and carries the same annotated plan and summary that `spark.comet.explain.format=verbose` produces for a real -Comet plan. It is written once per action, so a plan that is built but never -executed — `df.explain()`, or reading `queryExecution.executedPlan` — is not -reported, and a DataFrame collected twice is reported twice. +Comet plan. It is written once per SQL execution, so a plan that is built but +never executed — `df.explain()`, or reading `queryExecution.executedPlan` — is +not reported, and a DataFrame collected twice is reported twice. + +"Per SQL execution" rather than "per action" is a real distinction for RDD work. +The report comes from a `QueryExecutionListener`, which Spark drives from the +Dataset action path, so actions taken on `df.rdd` are outside it. On Spark 4.0 +and later, obtaining `df.rdd` runs a query of its own and is reported once at +that point; later actions on the resulting RDD add nothing, because no new SQL +execution starts. On Spark 3.4 and 3.5 obtaining `df.rdd` is not reported at +all. Either way the RDD's own actions are never counted individually, so treat +an RDD-heavy workload's report as covering the plans it built, not the jobs it +ran. The preview goes through the whole Comet planning sequence rather than operator conversion alone: Spark's columnar transitions are inserted and Comet'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 2fc359d6d0f..083366b7e71 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -677,7 +677,9 @@ case class CometExecRule(session: SparkSession) // `CometPlanOnly` for why the report is not built here. if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) { CometPlanOnly.register(session) - return plan + // Snapshot the settings with the plan: the report is delivered asynchronously, long after + // a `withSQLConf` block around the action may have restored them. + return CometPlanOnly.tagSettings(session, plan) } val newPlan = _apply(plan) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala b/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala index b81d096fba8..93c225c9e36 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala @@ -23,7 +23,9 @@ import scala.util.control.NonFatal import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, CommandResultExec, ExecSubqueryExpression, InputAdapter, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.catalyst.expressions.Alias +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, CommandResultExec, ExecSubqueryExpression, InputAdapter, ProjectExec, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.command.ExecutedCommandExec import org.apache.spark.sql.execution.datasources.v2.V2CommandExec @@ -46,7 +48,14 @@ import org.apache.comet.CometSparkSessionExtensions.isCometLoaded * simple. Spark applies a planner rule many times for one query - once per query stage and once * per adaptive re-optimization under AQE, and separately for every subquery it prepares - and * telling those applications apart takes a good deal of bookkeeping. A query execution listener - * fires once per action, holding the finished plan, so there is nothing to tell apart. + * fires once per SQL execution, holding the finished plan, so there is nothing to tell apart. + * + * Per SQL execution, not per action: Spark drives `QueryExecutionListener` from the Dataset + * action path, so work taken through `df.rdd` is outside it. On Spark 4.0+ obtaining `df.rdd` + * runs a query of its own and is reported at that point, while the RDD actions that follow are + * not; on 3.4/3.5 obtaining it is not reported either. Covering those would mean a second + * reporting path keyed on job starts, and reconciling it against this one so an ordinary query is + * not reported twice. */ object CometPlanOnly extends Logging { @@ -76,6 +85,72 @@ object CometPlanOnly extends Logging { } } + /** + * The settings that decide whether a query gets a report and how its preview is built, as they + * stood while the query was being planned. + * + * Reporting is asynchronous, so by the time the listener runs the caller may have restored or + * changed any of these - a `withSQLConf` block that runs `collect()` and exits before the + * callback is delivered is enough. Reading them back off the session then decides one query's + * report using another query's settings, which usually means dropping it. Snapshotting at plan + * time keeps the decision with the query it belongs to. + */ + private case class PlanOnlySettings( + enabled: Boolean, + cometLoaded: Boolean, + execEnabled: Boolean) + + /** + * Set on the plan `CometExecRule` saw while plan-only mode was on. Read back in `report`, which + * runs on the listener bus with no access to the planning thread's configuration. + */ + private val PLAN_ONLY_SETTINGS = new TreeNodeTag[PlanOnlySettings]("CometPlanOnlySettings") + + /** + * Record the plan-time settings on `plan` and return it unchanged. + * + * Called by `CometExecRule` on the plan it is declining to convert. + */ + def tagSettings(session: SparkSession, plan: SparkPlan): SparkPlan = { + val conf = session.sessionState.conf + plan.setTagValue( + PLAN_ONLY_SETTINGS, + PlanOnlySettings( + enabled = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf), + cometLoaded = isCometLoaded(conf), + execEnabled = CometConf.COMET_EXEC_ENABLED.get(conf))) + plan + } + + /** + * The settings recorded for this query, searching the adaptive wrappers' own plans as well: + * under AQE the tagged plan is the one AQE was handed, which hangs off `AdaptiveSparkPlanExec` + * rather than appearing among its children. + * + * Falls back to reading the session when no tag is found, which covers a plan that reached the + * listener without passing through `CometExecRule`. + */ + private def settingsFor(qe: QueryExecution): PlanOnlySettings = { + def search(plan: SparkPlan): Option[PlanOnlySettings] = + plan + .getTagValue(PLAN_ONLY_SETTINGS) + .orElse((plan match { + case adaptive: AdaptiveSparkPlanExec => + Seq(adaptive.inputPlan, adaptive.initialPlan, adaptive.executedPlan) + case stage: QueryStageExec => Seq(stage.plan) + case _ => Seq.empty + }).flatMap(search).headOption) + .orElse(plan.children.flatMap(search).headOption) + + search(qe.executedPlan).getOrElse { + val conf = qe.sparkSession.sessionState.conf + PlanOnlySettings( + enabled = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf), + cometLoaded = isCometLoaded(conf), + execEnabled = CometConf.COMET_EXEC_ENABLED.get(conf)) + } + } + /** * Logs the Comet plan Comet would have executed for `qe`. * @@ -91,9 +166,9 @@ object CometPlanOnly extends Logging { val previous = SparkSession.getActiveSession SparkSession.setActiveSession(session) try { - val conf = session.sessionState.conf - if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && isCometLoaded(conf) && - CometConf.COMET_EXEC_ENABLED.get(conf) && !isMetadataOnly(qe.executedPlan)) { + val settings = settingsFor(qe) + if (settings.enabled && settings.cometLoaded && settings.execEnabled && + !isMetadataOnly(qe.executedPlan)) { val preview = previewOf(session, qe.executedPlan) logWarning(s"$REPORT_PREFIX\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") } @@ -171,7 +246,15 @@ object CometPlanOnly extends Logging { // conversion would never reach the subtree while the coverage count - which unwraps the // wrapper - still counts every operator in it as Spark. Undo the reuse and let both copies // convert, which is what the counts of a real Comet run reflect. - case reused: ReusedExchangeExec => stripPreparation(reused.child) + // + // The wrapper's own output IDs have to survive that, though. `ReusedExchangeExec` re-aliases + // the shared child's output, so a self-join can expose `k#16` here while the child produces + // `k#3`. Returning the child bare leaves the parent sort referencing an attribute nothing in + // its subtree produces; Comet's attribute binder then declines the sort, the join above it + // stays on Spark, and the preview understates coverage for work Comet would have accelerated. + // Re-alias positionally, the same correspondence `ReusedExchangeExec` itself relies on. + case reused: ReusedExchangeExec => + restoreReusedOutput(reused, stripPreparation(reused.child)) case WholeStageCodegenExec(child) => stripPreparation(child) case InputAdapter(child) => stripPreparation(child) case ColumnarToRowExec(child) => stripPreparation(child) @@ -179,6 +262,32 @@ object CometPlanOnly extends Logging { case other => other.withNewChildren(other.children.map(stripPreparation)) } + /** + * `stripped`, the subtree that was behind a `ReusedExchangeExec`, presenting the wrapper's + * output attributes instead of its own. + * + * A no-op when the IDs already agree, which is the common case: reuse only re-aliases when the + * two copies were planned with different attribute IDs. Otherwise a `ProjectExec` of aliases + * carries the wrapper's `exprId`s, which is the cheapest node that can re-label an output + * without disturbing the subtree the preview is trying to measure. It does add one operator to + * the report, but a projection of aliases is one Comet converts, so the coverage percentage is + * not skewed the way losing the parent sort and its join was. + */ + private def restoreReusedOutput(reused: ReusedExchangeExec, stripped: SparkPlan): SparkPlan = { + val target = reused.output + val source = stripped.output + if (target.length != source.length || target.zip(source).forall { case (t, s) => + t.exprId == s.exprId + }) { + return stripped + } + val aliases = target.zip(source).map { case (t, s) => + if (t.exprId == s.exprId) s + else Alias(s, t.name)(exprId = t.exprId, qualifier = t.qualifier) + } + ProjectExec(aliases, stripped) + } + /** * `plan` with the plan behind each of its subquery expressions replaced by that plan's own * preview. diff --git a/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala index cce9d712b62..cff7b270f17 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala @@ -23,11 +23,13 @@ import org.apache.logging.log4j.Level import org.apache.spark.CometListenerBusUtils import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.CometPlan -import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan, SubqueryBroadcastExec} +import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution, SparkPlan, SubqueryBroadcastExec} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.util.QueryExecutionListener import org.apache.comet.{CometConf, CometCoverageStats} +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus /** * Tests for plan-only mode: `spark.comet.explain.planOnly.enabled`. @@ -195,18 +197,114 @@ class CometPlanOnlySuite extends CometTestBase { } } - // `df.rdd` - the path PySpark's `df.rdd` takes through `Dataset.javaToPython` - plans a second - // query of its own and runs it under its own execution id, so it is reported too, once. - test("an RDD action is reported once") { + // Reporting is driven by `QueryExecutionListener`, which Spark fires from the Dataset action + // path, so RDD actions are outside it. What `df.rdd` itself does differs by version: on 4.0+ it + // runs a query of its own under its own execution id and is reported once, on 3.4/3.5 it is not + // reported at all. Either way the RDD's own actions add nothing, which is the part that makes + // "one report per action" the wrong description of this mode. Asserted rather than skipped so + // the version split is pinned instead of rediscovered. + test("RDD actions are outside the reported path") { withSQLConf( planOnlyConf(aqe = true, useV1 = 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 _2, count(*) FROM tbl GROUP BY _2").rdd.count() + val rdd = spark.sql("SELECT _2, count(*) FROM tbl GROUP BY _2").rdd + // Three actions on the same RDD. None of them starts a SQL execution, so between them + // they contribute nothing beyond whatever obtaining `rdd` already did. + rdd.count() + rdd.count() + rdd.collect() + } + if (isSpark40Plus) { + assert( + reports.size == 1, + s"expected obtaining `.rdd` to be reported once, got:\n${reports.mkString("\n\n")}") + assert(reports.head.contains("HashAggregate"), s"report:\n${reports.head}") + } else { + assert( + reports.isEmpty, + s"expected no report on Spark 3.x, got:\n${reports.mkString("\n\n")}") + } + } + } + } + + // The report is delivered asynchronously, so a caller can enable plan-only, run the action, and + // restore the setting while the callback is still queued. Reading the flag back off the session + // at that point drops the report for a query that really did run in plan-only mode. + // + // The ordering has to be forced or the test proves nothing: normally the bus drains during + // `collect()` and the callback sees the flag still on. A fresh session gives a deterministic + // registration order - `CometPlanOnly` registers its listener lazily on first use, so on a + // session that has never run a plan-only query the gating listener below is registered first + // and therefore runs first, holding the bus until the setting has been restored. + test("a report survives the setting being restored before the callback runs") { + val session = spark.newSession() + val gate = new java.util.concurrent.CountDownLatch(1) + val gating = new QueryExecutionListener { + override def onSuccess(name: String, qe: QueryExecution, durationNs: Long): Unit = { + gate.await(60, java.util.concurrent.TimeUnit.SECONDS) + } + override def onFailure(name: String, qe: QueryExecution, e: Exception): Unit = () + } + session.listenerManager.register(gating) + try { + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + val appender = new LogAppender("Comet plan-only reports") + withLogAppender(appender, loggerNames = Seq(reporterLogger), level = Some(Level.WARN)) { + val conf = planOnlyConf(aqe = true, useV1 = true) :+ + (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") + conf.foreach { case (k, v) => session.conf.set(k, v) } + session.sql("SELECT id, count(*) FROM range(100) GROUP BY id").collect() + // The callback is parked in `gating.onSuccess`, so restoring the setting here happens + // strictly before `CometPlanOnly` gets to look at it. + session.conf.unset(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key) + gate.countDown() + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } + val reports = appender.loggingEvents + .map(_.getMessage.getFormattedMessage) + .filter(_.startsWith(PLAN_ONLY_PREFIX)) + .toSeq + assert( + reports.size == 1, + s"expected the query planned under plan-only mode to still be reported, got:\n" + + reports.mkString("\n\n")) + } finally { + gate.countDown() + session.listenerManager.unregister(gating) + } + } + + // `ReusedExchangeExec` re-aliases the shared child's output, so a self-join can expose `k#16` + // while its child produces `k#3`. Undoing the reuse without carrying those IDs across leaves + // the parent sort referencing an attribute nothing below it produces, and Comet's binder then + // declines the sort and the join above it - understating coverage for work Comet would have run. + test("a reused exchange keeps its output IDs so consumers still convert") { + withSQLConf( + planOnlyConf(aqe = false, useV1 = true) ++ Seq( + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1"): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val reports = capturePlanOnlyReports { + spark + .sql("""SELECT a._2, b._2 FROM + | (SELECT _2, count(*) c FROM tbl GROUP BY _2) a + | JOIN + | (SELECT _2, count(*) c FROM tbl GROUP BY _2) b + | ON a._2 = b._2""".stripMargin) + .collect() } assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") - assert(reports.head.contains("HashAggregate"), s"report:\n${reports.head}") + val report = reports.head + assert( + report.contains("CometSortMergeJoin") || report.contains("CometHashJoin"), + s"the join over a reused exchange should convert in the preview, got:\n$report") + assert( + !report.contains("\nSort ") && !report.contains("+- Sort "), + s"no Sort should be left on Spark by a lost attribute binding, got:\n$report") } } } From ac8e5a1c86329d520e658fca20c2ce8ab765d8bc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 9 Sep 2026 14:19:59 -0600 Subject: [PATCH 4/5] fix: keep a plan-only preview from mixing two configurations Two review follow-ups on the asynchronous reporting boundary. The settings snapshot covered three flags, but the preview reruns the conversion rules, which read per-operator gates, the strict-fallback and shuffle settings off op.conf, and Spark settings such as ANSI mode. It cannot be injected: SparkPlan.conf is the captured session's live conf, so neither a thread-local SQLConf nor a cloned session redirects op.conf. The snapshot now records every set SQL conf and is used to check that the configuration the preview will read is still the query's own; when it is not, the report is skipped and names the settings that moved. The reporting gate itself stays excluded, so restoring plan-only mode before the callback still does not drop the report. Reused-exchange identity is now restored by rewriting the subtree's attribute IDs rather than adding a ProjectExec of aliases. The synthetic projection failed CometProjectExec.enabledConfig with spark.comet.exec.project.enabled=false, so the parent sort lost its native child and the consuming join fell back, dropping coverage from 14/14 to 10/14 purely because of a bookkeeping node. The rewrite adds no operator, so it also stops the identity bookkeeping moving the coverage percentage. Also drops an interpolator prefix on a string with no substitution, which the lint jobs flagged. --- .../apache/comet/rules/CometPlanOnly.scala | 117 +++++++++++---- .../comet/rules/CometPlanOnlySuite.scala | 135 ++++++++++++------ 2 files changed, 182 insertions(+), 70 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala b/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala index 93c225c9e36..637028d46e5 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala @@ -23,13 +23,14 @@ import scala.util.control.NonFatal import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.Alias +import org.apache.spark.sql.catalyst.expressions.{Attribute, ExprId} import org.apache.spark.sql.catalyst.trees.TreeNodeTag -import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, CommandResultExec, ExecSubqueryExpression, InputAdapter, ProjectExec, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, CommandResultExec, ExecSubqueryExpression, InputAdapter, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.command.ExecutedCommandExec import org.apache.spark.sql.execution.datasources.v2.V2CommandExec import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.util.QueryExecutionListener import org.apache.comet.{CometConf, ExtendedExplainInfo} @@ -94,11 +95,41 @@ object CometPlanOnly extends Logging { * callback is delivered is enough. Reading them back off the session then decides one query's * report using another query's settings, which usually means dropping it. Snapshotting at plan * time keeps the decision with the query it belongs to. + * + * `sqlConfs` is every SQL conf that was explicitly set, not just Comet's. The preview reruns + * the conversion rules, and those read far more than three flags: `isCometLoaded` and the two + * rules take the session conf, `CometExecRule` reads per-operator gates and the strict-fallback + * and shuffle settings off `op.conf`, and the serde reads Spark settings such as ANSI mode and + * the session time zone. `op.conf` is `session.sessionState.conf` of the session each node + * captured when it was constructed (`SparkPlan.conf`), so no thread-local override and no + * cloned session can redirect it. The snapshot therefore cannot be *injected* into the preview; + * it is used to check that the configuration the preview will read is still the one the query + * was planned under, and the report is skipped when it is not. See [[changedSince]]. */ private case class PlanOnlySettings( enabled: Boolean, cometLoaded: Boolean, - execEnabled: Boolean) + execEnabled: Boolean, + sqlConfs: Map[String, String]) { + + /** + * The settings that have changed since this snapshot was taken, other than the reporting gate + * itself. + * + * `spark.comet.explain.planOnly.enabled` is deliberately excluded: turning plan-only mode + * off, or leaving the `withSQLConf` block that turned it on, must not drop the report for a + * query that was planned while it was on. That flag only gates reporting, and the snapshot is + * the authority for it. Everything else shapes the preview, so a difference there means the + * preview would describe a configuration the query never ran under. + */ + def changedSince(conf: SQLConf): Seq[String] = { + val now = conf.getAllConfs + val gate = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key + (sqlConfs.keySet ++ now.keySet).toSeq.sorted + .filterNot(_ == gate) + .filter(key => sqlConfs.get(key) != now.get(key)) + } + } /** * Set on the plan `CometExecRule` saw while plan-only mode was on. Read back in `report`, which @@ -106,19 +137,20 @@ object CometPlanOnly extends Logging { */ private val PLAN_ONLY_SETTINGS = new TreeNodeTag[PlanOnlySettings]("CometPlanOnlySettings") + private def snapshot(conf: SQLConf): PlanOnlySettings = + PlanOnlySettings( + enabled = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf), + cometLoaded = isCometLoaded(conf), + execEnabled = CometConf.COMET_EXEC_ENABLED.get(conf), + sqlConfs = conf.getAllConfs) + /** * Record the plan-time settings on `plan` and return it unchanged. * * Called by `CometExecRule` on the plan it is declining to convert. */ def tagSettings(session: SparkSession, plan: SparkPlan): SparkPlan = { - val conf = session.sessionState.conf - plan.setTagValue( - PLAN_ONLY_SETTINGS, - PlanOnlySettings( - enabled = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf), - cometLoaded = isCometLoaded(conf), - execEnabled = CometConf.COMET_EXEC_ENABLED.get(conf))) + plan.setTagValue(PLAN_ONLY_SETTINGS, snapshot(session.sessionState.conf)) plan } @@ -142,13 +174,7 @@ object CometPlanOnly extends Logging { }).flatMap(search).headOption) .orElse(plan.children.flatMap(search).headOption) - search(qe.executedPlan).getOrElse { - val conf = qe.sparkSession.sessionState.conf - PlanOnlySettings( - enabled = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf), - cometLoaded = isCometLoaded(conf), - execEnabled = CometConf.COMET_EXEC_ENABLED.get(conf)) - } + search(qe.executedPlan).getOrElse(snapshot(qe.sparkSession.sessionState.conf)) } /** @@ -169,8 +195,19 @@ object CometPlanOnly extends Logging { val settings = settingsFor(qe) if (settings.enabled && settings.cometLoaded && settings.execEnabled && !isMetadataOnly(qe.executedPlan)) { - val preview = previewOf(session, qe.executedPlan) - logWarning(s"$REPORT_PREFIX\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") + val changed = settings.changedSince(session.sessionState.conf) + if (changed.nonEmpty) { + // Better no report than one describing a configuration the query never ran under. The + // preview cannot be built under the snapshot, so the only honest options are to report + // what the query was planned with or to say why nothing is being reported. + logWarning( + s"$REPORT_PREFIX not reporting this query: settings changed since it was planned, " + + s"so a preview would describe a different configuration (${changed.mkString(", ")})") + } else { + val preview = previewOf(session, qe.executedPlan) + logWarning( + s"$REPORT_PREFIX\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") + } } } catch { case NonFatal(e) => @@ -267,11 +304,30 @@ object CometPlanOnly extends Logging { * output attributes instead of its own. * * A no-op when the IDs already agree, which is the common case: reuse only re-aliases when the - * two copies were planned with different attribute IDs. Otherwise a `ProjectExec` of aliases - * carries the wrapper's `exprId`s, which is the cheapest node that can re-label an output - * without disturbing the subtree the preview is trying to measure. It does add one operator to - * the report, but a projection of aliases is one Comet converts, so the coverage percentage is - * not skewed the way losing the parent sort and its join was. + * two copies were planned with different attribute IDs. + * + * Otherwise the subtree's own attributes are rewritten to the wrapper's, everywhere they appear + * beneath it, rather than re-labelled by an added `ProjectExec`. Two reasons, and both of them + * are about the added node changing the thing being measured: + * + * - a projection participates in `spark.comet.exec.project.enabled`. With that off, the + * synthetic node fails `CometProjectExec.enabledConfig`, the parent sort then has no native + * child, and the consuming join falls back too - so the preview would lose a supported + * sort/join purely because of a bookkeeping node, which is the failure this method exists + * to avoid. A real conversion of those branches has no such node in it. + * - it counts as an eligible operator, so it moves the coverage percentage the report exists + * to state. + * + * The rewrite runs on the whole subtree, not just its root, because almost every operator + * derives `output` from its children: the producing leaf has to be rewritten for the root to + * present the new IDs. That is safe because attribute IDs are globally unique, so an ID in the + * root's output means the same attribute wherever it appears below, and intermediate attributes + * that are not in the root's output are left alone. It is also safe to do at all only because + * this plan is a throwaway preview that is never executed. + * + * The direction matters. Rewriting the consumers instead, to reference the shared subtree's + * IDs, would merge the two copies of a self-join onto one set of attributes, which is exactly + * what `ReusedExchangeExec`'s re-aliasing exists to prevent. */ private def restoreReusedOutput(reused: ReusedExchangeExec, stripped: SparkPlan): SparkPlan = { val target = reused.output @@ -281,11 +337,16 @@ object CometPlanOnly extends Logging { }) { return stripped } - val aliases = target.zip(source).map { case (t, s) => - if (t.exprId == s.exprId) s - else Alias(s, t.name)(exprId = t.exprId, qualifier = t.qualifier) + val rewrites: Map[ExprId, Attribute] = source + .zip(target) + .collect { case (s, t) if s.exprId != t.exprId => s.exprId -> t } + .toMap + stripped.transformUp { case node => + node.transformExpressions { + case a: Attribute if rewrites.contains(a.exprId) => + rewrites(a.exprId) + } } - ProjectExec(aliases, stripped) } /** diff --git a/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala index cff7b270f17..96206971158 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala @@ -21,7 +21,7 @@ package org.apache.comet.rules import org.apache.logging.log4j.Level import org.apache.spark.CometListenerBusUtils -import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.{CometTestBase, SparkSession} import org.apache.spark.sql.comet.CometPlan import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution, SparkPlan, SubqueryBroadcastExec} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec @@ -230,16 +230,18 @@ class CometPlanOnlySuite extends CometTestBase { } } - // The report is delivered asynchronously, so a caller can enable plan-only, run the action, and - // restore the setting while the callback is still queued. Reading the flag back off the session - // at that point drops the report for a query that really did run in plan-only mode. - // - // The ordering has to be forced or the test proves nothing: normally the bus drains during - // `collect()` and the callback sees the flag still on. A fresh session gives a deterministic - // registration order - `CometPlanOnly` registers its listener lazily on first use, so on a - // session that has never run a plan-only query the gating listener below is registered first - // and therefore runs first, holding the bus until the setting has been restored. - test("a report survives the setting being restored before the callback runs") { + /** + * Runs `SELECT id, count(*) FROM range(100) GROUP BY id` on a fresh session with plan-only mode + * on, holding the listener bus until `whileParked` has run, and returns the reports. + * + * The ordering has to be forced or a test using this proves nothing: normally the bus drains + * during `collect()` and the callback sees the settings still in place. A fresh session gives a + * deterministic registration order - `CometPlanOnly` registers its listener lazily on first + * use, so on a session that has never run a plan-only query the gating listener below is + * registered first and therefore runs first, holding the bus until `whileParked` has returned. + */ + private def reportsWithSettingsChangedBeforeCallback( + whileParked: SparkSession => Unit): Seq[String] = { val session = spark.newSession() val gate = new java.util.concurrent.CountDownLatch(1) val gating = new QueryExecutionListener { @@ -257,54 +259,103 @@ class CometPlanOnlySuite extends CometTestBase { (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") conf.foreach { case (k, v) => session.conf.set(k, v) } session.sql("SELECT id, count(*) FROM range(100) GROUP BY id").collect() - // The callback is parked in `gating.onSuccess`, so restoring the setting here happens - // strictly before `CometPlanOnly` gets to look at it. - session.conf.unset(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key) + // The callback is parked in `gating.onSuccess`, so anything done here happens strictly + // before `CometPlanOnly` gets to look at the settings. + whileParked(session) gate.countDown() CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) } - val reports = appender.loggingEvents + appender.loggingEvents .map(_.getMessage.getFormattedMessage) .filter(_.startsWith(PLAN_ONLY_PREFIX)) .toSeq - assert( - reports.size == 1, - s"expected the query planned under plan-only mode to still be reported, got:\n" + - reports.mkString("\n\n")) } finally { gate.countDown() session.listenerManager.unregister(gating) } } + // The report is delivered asynchronously, so a caller can enable plan-only, run the action, and + // restore the setting while the callback is still queued. Reading the flag back off the session + // at that point drops the report for a query that really did run in plan-only mode. The gate + // itself is the one setting the snapshot overrides, so restoring it must not cost the report. + test("a report survives the setting being restored before the callback runs") { + val reports = reportsWithSettingsChangedBeforeCallback { session => + session.conf.unset(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key) + } + assert( + reports.size == 1, + "expected the query planned under plan-only mode to still be reported, got:\n" + + reports.mkString("\n\n")) + assert( + reports.head.contains("Comet accelerated"), + s"expected a coverage report, got:\n${reports.head}") + } + + // Every other setting shapes the preview rather than gating it, and the snapshot cannot be + // injected into the conversion: `SparkPlan.conf` is the captured session's live conf, so + // `op.conf` reads it whatever the listener does with active sessions or thread-locals. Reporting + // anyway would describe a configuration the query never ran under, so the report is skipped and + // says which settings moved. Compared here against the unchanged-settings control above, rather + // than only asserting that some line was logged. + test("a report is skipped when a conversion setting changed before the callback runs") { + val changed = reportsWithSettingsChangedBeforeCallback { session => + // Not the gate: a setting the preview's conversion actually reads. + session.conf.set(CometConf.COMET_EXEC_SORT_ENABLED.key, "false") + } + assert(changed.size == 1, s"expected one line, got:\n${changed.mkString("\n\n")}") + assert( + changed.head.contains("settings changed since it was planned"), + s"expected the skip diagnostic, got:\n${changed.head}") + assert( + changed.head.contains(CometConf.COMET_EXEC_SORT_ENABLED.key), + s"expected the diagnostic to name the changed setting, got:\n${changed.head}") + assert( + !changed.head.contains("Comet accelerated"), + s"expected no coverage numbers from a configuration the query never ran under, got:\n" + + changed.head) + } + // `ReusedExchangeExec` re-aliases the shared child's output, so a self-join can expose `k#16` // while its child produces `k#3`. Undoing the reuse without carrying those IDs across leaves // the parent sort referencing an attribute nothing below it produces, and Comet's binder then // declines the sort and the join above it - understating coverage for work Comet would have run. - test("a reused exchange keeps its output IDs so consumers still convert") { - withSQLConf( - planOnlyConf(aqe = false, useV1 = true) ++ Seq( - CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true", - SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1"): _*) { - withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { - val reports = capturePlanOnlyReports { - spark - .sql("""SELECT a._2, b._2 FROM - | (SELECT _2, count(*) c FROM tbl GROUP BY _2) a - | JOIN - | (SELECT _2, count(*) c FROM tbl GROUP BY _2) b - | ON a._2 = b._2""".stripMargin) - .collect() + // + // Run with projection both enabled and disabled. Restoring identity with an added `ProjectExec` + // passed the first and failed the second: the synthetic node fails + // `CometProjectExec.enabledConfig`, so the parent sort loses its native child and the join falls + // back, losing a supported sort/join purely because of a bookkeeping node that a real conversion + // of those branches would not contain. Rewriting the subtree's attribute IDs instead adds no node + // and reads no user setting. + Seq(true, false).foreach { projectEnabled => + test( + "a reused exchange keeps its output IDs so consumers still convert " + + s"(project.enabled=$projectEnabled)") { + withSQLConf( + planOnlyConf(aqe = false, useV1 = true) ++ Seq( + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> projectEnabled.toString, + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1"): _*) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val reports = capturePlanOnlyReports { + spark + .sql("""SELECT a._2, b._2 FROM + | (SELECT _2, count(*) c FROM tbl GROUP BY _2) a + | JOIN + | (SELECT _2, count(*) c FROM tbl GROUP BY _2) b + | ON a._2 = b._2""".stripMargin) + .collect() + } + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + val report = reports.head + assert( + report.contains("CometSortMergeJoin") || report.contains("CometHashJoin"), + s"the join over a reused exchange should convert in the preview, got:\n$report") + assert( + !report.contains("\nSort ") && !report.contains("+- Sort "), + s"no Sort should be left on Spark by a lost attribute binding, got:\n$report") } - assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") - val report = reports.head - assert( - report.contains("CometSortMergeJoin") || report.contains("CometHashJoin"), - s"the join over a reused exchange should convert in the preview, got:\n$report") - assert( - !report.contains("\nSort ") && !report.contains("+- Sort "), - s"no Sort should be left on Spark by a lost attribute binding, got:\n$report") } } } From 43669546f6b6e41ad075ae5cec18cb27fc9a61e5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 17 Sep 2026 18:37:37 -0600 Subject: [PATCH 5/5] fix: drop a redundant string interpolator that scalafix rejects RedundantSyntax flagged an `s"..."` with no interpolation in CometPlanOnlySuite, failing both the syntactic scalafix job and all four lint-java jobs. --- .../test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala index 96206971158..5d0e26f1c57 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala @@ -312,7 +312,7 @@ class CometPlanOnlySuite extends CometTestBase { s"expected the diagnostic to name the changed setting, got:\n${changed.head}") assert( !changed.head.contains("Comet accelerated"), - s"expected no coverage numbers from a configuration the query never ran under, got:\n" + + "expected no coverage numbers from a configuration the query never ran under, got:\n" + changed.head) }