diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index b7c88254c9..cbe5d1fd00 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 8210f91b7f..83389fff37 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 diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index a600615b4b..13305aefb9 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 d8fe5b6989..62c0656628 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 c69602fc80..7f6f091537 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 0000000000..b81d096fba --- /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 a524da3af9..cdba25576b 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 1e0cfc79e0..c9107fe350 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 0000000000..cce9d712b6 --- /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")}") + } + } +}