diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 79fea7fb881..4966d38c8ce 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -418,6 +418,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 5473d3988a8..51c608a2561 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -191,6 +191,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 32fe297e67d..0c40a4313af 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -199,6 +199,59 @@ 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 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 +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 c015d2d4e66..e18fd590a32 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -741,6 +741,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 db5373e1093..083366b7e71 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -672,6 +672,16 @@ 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) + // 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) if (showTransformations && !newPlan.fastEquals(plan)) { logInfo(s""" @@ -682,7 +692,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..637028d46e5 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala @@ -0,0 +1,413 @@ +/* + * 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.catalyst.expressions.{Attribute, ExprId} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +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} +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 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 { + + 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") + } + } + + /** + * 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. + * + * `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, + 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 + * runs on the listener bus with no access to the planning thread's configuration. + */ + 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 = { + plan.setTagValue(PLAN_ONLY_SETTINGS, snapshot(session.sessionState.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(snapshot(qe.sparkSession.sessionState.conf)) + } + + /** + * 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 settings = settingsFor(qe) + if (settings.enabled && settings.cometLoaded && settings.execEnabled && + !isMetadataOnly(qe.executedPlan)) { + 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) => + 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. + // + // 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) + case RowToColumnarExec(child) => stripPreparation(child) + 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 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 + val source = stripped.output + if (target.length != source.length || target.zip(source).forall { case (t, s) => + t.exprId == s.exprId + }) { + return stripped + } + 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) + } + } + } + + /** + * `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 b07052a1877..be582961db8 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 9a941c28ffc..6664f75d100 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -54,6 +54,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..96206971158 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/rules/CometPlanOnlySuite.scala @@ -0,0 +1,524 @@ +/* + * 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, 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 +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`. + * + * 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")}") + } + } + } + + // 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 { + 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")}") + } + } + } + } + + /** + * 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 { + 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 anything done here happens strictly + // before `CometPlanOnly` gets to look at the settings. + whileParked(session) + gate.countDown() + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } + appender.loggingEvents + .map(_.getMessage.getFormattedMessage) + .filter(_.startsWith(PLAN_ONLY_PREFIX)) + .toSeq + } 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. + // + // 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") + } + } + } + } + + // 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")}") + } + } +}