diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index a600615b4b..336fc5f6d4 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -208,6 +208,55 @@ operators were arranged after Comet's serialization). See the [Metrics Guide](metrics.md) for details on the DataFusion metrics that appear in this output. +### `spark.comet.explain.planOnly.enabled` + +When enabled, Comet runs its full conversion pass on every query and logs the +resulting Comet plan and coverage summary to the driver log, then reverts to +executing the plan on Spark instead of offloading anything to native. Use this +to evaluate how much of a workload Comet would accelerate without changing the +execution. + +The log line is prefixed with `[Comet plan-only]` and includes the same +annotated plan and summary as `spark.comet.explain.format=verbose` produces +against a normal Comet plan. The preview goes through the whole Comet planning +sequence, not just operator conversion: Spark's columnar transitions are +inserted and Comet's post-columnar rules +(`RevertNativeForTransitionHeavyStages`, `EliminateRedundantTransitions`) are +applied, so a stage that Comet would have reverted to Spark for having too many +transitions is reported as reverted. + +Spark prepares some plans on their own, ahead of the query that contains them — +scalar subqueries and dynamic partition pruning subqueries, for instance — so a +query gets one report per independently planned plan: one for the outer query, +plus one per such subquery. The outer report counts its subqueries too, the same +way normal Comet planning does, so the reports for one query describe +overlapping sets of operators and their counts should not be added up. Repeat +applications of the same plan are not reported again: under AQE, neither the +per-stage applications nor the applications that follow each adaptive +re-optimization add reports, and nor does a plan AQE re-plans wholesale after a +stage materializes empty. One consequence under AQE is that a subquery Spark +only plans once query stages are under way — a DPP subquery, for instance — has +no report of its own. + +The estimate reflects Scala-side conversion only. The native plan is never +handed to DataFusion, so anything that would have failed in DataFusion's +`create_plan` still counts as accelerated. Treat the percentage as an upper +bound. + +Under AQE the report is an estimate for a second reason: it describes the plan +as it stands before any adaptive re-planning, and the post-columnar rules are +applied to that whole plan at once rather than to each stage as it is created. +Coverage of the plan AQE finally executes can differ. One case is worth calling +out, because it moves the number the other way: AQE does not plan a subquery +into the outer plan until after the report has been produced, so the outer +report counts a subquery's operators as un-accelerated Spark even where Comet +would accelerate them. For a subquery-heavy query under AQE, read the +per-subquery reports rather than the outer percentage, or turn AQE off for the +evaluation run. + +The config requires `spark.comet.exec.enabled=true`. With Comet exec disabled +the rule that emits the report does not run. + ## Programmatic Access to Fallback Reasons The configs above route fallback reasons to logs or the SQL UI. If you want diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index f40d0cbdd9..94f8fafd86 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -688,6 +688,18 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_EXPLAIN_PLAN_ONLY_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.explain.planOnly.enabled") + .category(CATEGORY_EXEC_EXPLAIN) + .doc("When enabled, Comet builds the Comet plan it would have executed and logs it to " + + "the driver log, then discards it and lets Spark execute the query. Use this to " + + "evaluate how much of a workload Comet would accelerate without changing execution. " + + "The estimate is Scala-side only; native planning failures are not surfaced, so the " + + "acceleration percentage can be optimistic. Requires `spark.comet.exec.enabled=true`. " + + "Disabled by default.") + .booleanConf + .createWithDefault(false) + val COMET_STRICT_FALLBACK_REASONS: ConfigEntry[Boolean] = conf("spark.comet.explain.fallback.strict.enabled") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 4dc3619433..22a450d0a0 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -97,7 +97,9 @@ class CometSparkSessionExtensions // No-op on Spark 3.5+; see CometSpark34AqeDppFallbackRule's class docstring. injectPreSpark35QueryStagePrepRuleShim(extensions, CometSpark34AqeDppFallbackRule) extensions.injectQueryStagePrepRule { session => CometScanRule(session) } - extensions.injectQueryStagePrepRule { session => CometExecRule(session) } + extensions.injectQueryStagePrepRule { session => + CometExecRule(session, queryStagePrep = true) + } injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters) injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) } @@ -111,6 +113,8 @@ class CometSparkSessionExtensions override def preColumnarTransitions: Rule[SparkPlan] = CometExecRule(session) override def postColumnarTransitions: Rule[SparkPlan] = { + // Keep in sync with `CometExecRule.reportPlanOnlyCoverage`, which replays these rules over + // the plan it previews so that plan-only reports describe the plan that would have run. val rules = Seq(RevertNativeForTransitionHeavyStages(session), EliminateRedundantTransitions(session)) plan => rules.foldLeft(plan) { case (p, rule) => rule(p) } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index cbf6c6a1e8..f9625ab36f 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} import org.apache.spark.sql.execution.datasources.WriteFilesExec @@ -43,8 +43,9 @@ import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, V2CommandEx import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan import org.apache.spark.sql.execution.datasources.v2.json.JsonScan import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange, ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -115,12 +116,143 @@ object CometExecRule { */ val SKIP_COMET_BROADCAST_TAG: org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit] = org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometBroadcast") + + /** + * A bounded set of keys, used for plan-only reporting state. Evicts in LRU order once `limit` + * keys are held, so a long-lived driver retains a fixed amount of reporting state. Same + * synchronized-`LinkedHashMap` pattern used by `IcebergPlanDataInjector.commonCache`. + */ + private class BoundedKeySet(limit: Int) { + private val keys: java.util.Map[String, java.lang.Boolean] = + java.util.Collections.synchronizedMap( + new java.util.LinkedHashMap[String, java.lang.Boolean](16, 0.75f, true) { + override def removeEldestEntry( + eldest: java.util.Map.Entry[String, java.lang.Boolean]): Boolean = size() > limit + }) + + /** Adds `key`, returning true if it was not already present. */ + def add(key: String): Boolean = keys.put(key, java.lang.Boolean.TRUE) == null + + def contains(key: String): Boolean = keys.containsKey(key) + } + + private val PLAN_ONLY_REPORTED_LIMIT = 1024 + + /** `executionId:planFingerprint` keys that plan-only mode has already reported. */ + private val planOnlyReportedPlans = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT) + + /** Execution IDs for which AQE has begun cutting the plan into query stages. */ + private val planOnlyStagedExecutions = new BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT) + + /** + * Set on the root of every plan that plan-only mode has reported. Catalyst copies a node's tags + * onto the node that replaces it (`TreeNode.copyTagsFrom`), so the mark survives the rewrites + * Spark applies between one application of this rule and the next, which is what lets a later + * application recognize a plan it has already described. + */ + private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported") + + /** + * Whether `plan` is the plan of a query stage AQE has just cut, which reaches the columnar rule + * rooted at the `Exchange` the cut was made at. + * + * With AQE off a plan rooted at an `Exchange` is an ordinary plan - `df.repartition(n)`, say - + * and must still be reported, hence the `aqeEnabled` guard. + */ + private def isQueryStage( + plan: SparkPlan, + queryStagePrep: Boolean, + aqeEnabled: Boolean): Boolean = { + aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange] + } + + /** + * Whether `plan` is an application of this rule to a plan already reported, rather than to a + * plan the user is asking about. The state is on the plan itself, so this holds whether or not + * the application carries a SQL execution ID - `df.rdd.count()` and reading `executedPlan` + * without an action both plan, and in the first case execute AQE stages, with no execution ID + * set. + * + * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the wrapper is matched + * by name rather than through its contents. A re-optimized plan holds `QueryStageExec` nodes + * for the stages already materialized. A final plan for a query AQE never had to cut into + * stages holds neither, and is recognized by the mark left when it was first reported. + * + * @param queryStagePrep + * whether the calling rule instance is registered as a query-stage-prep rule. + */ + private def isReapplication(plan: SparkPlan): Boolean = { + plan.exists(p => + p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] || + p.getTagValue(PLAN_ONLY_REPORTED).isDefined) + } + + /** + * Whether plan-only mode should report `plan`, marking it reported if so. + * + * Spark applies this rule many times while executing one query, and only some of those + * applications correspond to a plan the user is asking about. Under AQE one query reaches the + * rules at least five times: + * + * - the initial plan, through the query-stage-prep rule - the one to report; + * - the same plan again as a columnar rule, now wrapped in `AdaptiveSparkPlanExec`; + * - each query stage as it is created, again as a columnar rule; + * - the re-optimized plan after each stage materializes, through the prep rule; + * - the final plan once every stage has materialized, as a columnar rule. + * + * Three mechanisms sort those out, because no one of them covers every shape: + * + * - Each scalar subquery and DPP subquery is prepared as its own top-level plan, and that + * happens *before* the outer plan reaches the conversion rules. A single report slot per + * SQL execution therefore let a nested subquery consume the slot and suppressed the outer + * plan, which is the plan being evaluated. Marking plans individually gives the outer plan + * its own report and each separately prepared subquery theirs; see [[isReapplication]]. + * - A mark cannot survive AQE replacing the physical tree wholesale, which is what happens + * when a stage materializes empty and the plan collapses to an empty relation: the new plan + * shares no nodes with the one reported and holds no query stages either. Once AQE has cut + * a stage for an execution, though, everything that arrives afterwards is AQE re-planning + * something already reported, and subqueries are all compiled before the first stage is + * cut, so suppressing the rest of the execution costs no report. + * - A subquery referenced from more than one place in the outer plan is prepared once per + * reference, as a separate but identical plan each time. Neither of the above catches + * those, so within one SQL execution the plan's structural hash dedupes them. + * + * @param queryStagePrep + * whether the calling rule instance is registered as a query-stage-prep rule. + */ + private[comet] def shouldReportPlanOnly( + executionId: Option[String], + plan: SparkPlan, + queryStagePrep: Boolean, + aqeEnabled: Boolean): Boolean = { + if (isQueryStage(plan, queryStagePrep, aqeEnabled)) { + // Record that AQE has started executing this query before dropping the stage itself. + executionId.foreach(planOnlyStagedExecutions.add) + false + } else if (isReapplication(plan)) { + false + } else if (executionId.exists(planOnlyStagedExecutions.contains)) { + false + } else { + plan.setTagValue(PLAN_ONLY_REPORTED, ()) + // Node tags are not part of a plan's structural hash, so the mark set above does not + // perturb the key. Without an execution ID there is nothing to scope the state to, and the + // checks above have already ruled out the repeat applications AQE makes, so report. + executionId.forall(id => planOnlyReportedPlans.add(s"$id:${plan.hashCode()}")) + } + } } /** * Spark physical optimizer rule for replacing Spark operators with Comet operators. + * + * @param queryStagePrep + * true for the instance registered with `injectQueryStagePrepRule`, which under AQE sees the + * whole initial plan, and false for the one registered as a columnar rule, which under AQE sees + * one query stage at a time. Only plan-only reporting reads this; see + * [[CometExecRule.shouldReportPlanOnly]]. */ -case class CometExecRule(session: SparkSession) +case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) extends Rule[SparkPlan] with ShimSubqueryBroadcast { @@ -563,7 +695,7 @@ case class CometExecRule(session: SparkSession) } override def apply(plan: SparkPlan): SparkPlan = { - val newPlan = _apply(plan) + val newPlan = _apply(plan, forPreview = false) if (showTransformations && !newPlan.fastEquals(plan)) { logInfo(s""" |=== Applying Rule $ruleName === @@ -573,7 +705,109 @@ case class CometExecRule(session: SparkSession) newPlan } - private def _apply(plan: SparkPlan): SparkPlan = { + /** + * Build the Comet plan we would have executed and log it. Called from `_apply` in plan-only + * mode; the built plan is discarded. + */ + private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = { + val preview = buildPreview(plan, topLevel = true) + logWarning(s"[Comet plan-only]\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") + } + + /** + * The plan Comet would have executed for `plan`. Passes `forPreview = true` through the nested + * calls so both conversion rules run their normal transforms instead of short-circuiting. + * + * Conversion is only the first half of Comet planning. Normally Spark then inserts the columnar + * transitions and runs Comet's post-columnar rules (see + * `CometSparkSessionExtensions.CometExecColumnar.postColumnarTransitions`), which can revert + * whole stages back to Spark and drop redundant transitions. Those steps run here too, so the + * report describes the plan that would really have executed and counts the transitions that + * would really have been there, rather than the pre-transition conversion result. + * + * `RevertNativeForTransitionHeavyStages` is applied with `applyToAllStages` because the preview + * holds the whole plan at once, whereas under AQE Spark hands that rule one stage at a time. + * + * @param topLevel + * false when previewing the plan behind a subquery expression. `ReuseExchangeAndSubquery` is + * the last step of Spark's preparation and `QueryExecution.preparations` omits it for a + * subquery, so the preview follows suit. + */ + private def buildPreview(plan: SparkPlan, topLevel: Boolean): SparkPlan = { + val converted = + _apply(CometScanRule(session)._apply(previewSubqueriesOf(plan)), forPreview = true) + val withTransitions = + ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted) + val reverted = RevertNativeForTransitionHeavyStages(session).applyToAllStages(withTransitions) + val preview = EliminateRedundantTransitions(session).apply(reverted) + if (topLevel) ReuseExchangeAndSubquery.apply(preview) else preview + } + + /** + * `plan` with the plan behind each of its subquery expressions replaced by that plan's own + * preview. + * + * Extended explain walks a node's `innerChildren`, which for a `SparkPlan` are the plans owned + * by its expressions, and counts their operators towards the report. Normal planning has + * already converted those plans by the time the outer plan reaches this rule - Spark prepares a + * scalar subquery through the full preparation sequence, columnar rules included, before + * substituting it into the outer plan - so leaving them untouched here would report every + * subquery operator as un-accelerated Spark and understate coverage relative to what Comet + * really executes. + * + * Each subquery is also reported in its own right, because Spark prepares it as a top-level + * plan of its own; those reports and the counts here therefore describe overlapping sets of + * operators. + */ + private def previewSubqueriesOf(plan: SparkPlan): SparkPlan = { + plan.transformAllExpressions { case subquery: ExecSubqueryExpression => + subquery.withNewPlan(previewSubquery(subquery.plan)) + } + } + + private def previewSubquery(subquery: BaseSubqueryExec): BaseSubqueryExec = subquery match { + // Reuse bookkeeping: the plan to preview is one level further down. + case reused: ReusedSubqueryExec => reused.copy(child = previewSubquery(reused.child)) + case other => + other.withNewChildren(Seq(previewPreparedPlan(other.child))).asInstanceOf[BaseSubqueryExec] + } + + /** + * Preview `plan`, which Spark prepared as a plan in its own right. + * + * A DPP subquery is the exception to that framing: `PlanDynamicPruningFilters` prepares the + * build plan and only then wraps it in a `BroadcastExchangeExec`, so the plan that went through + * the post-columnar rules - and the stage `RevertNativeForTransitionHeavyStages` judged - is + * the exchange's child, not the exchange. Previewing the exchange instead leaves its child a + * stage bounded at the top by the exchange, which stops the reversion firing, and the report + * then counts operators as accelerated that the executed plan runs on Spark. Descend through + * the wrapper and put it back, so the preview keeps the boundary Spark's preparation used. + */ + private def previewPreparedPlan(plan: SparkPlan): SparkPlan = plan match { + case exchange: BroadcastExchangeExec => + exchange.withNewChildren(Seq(previewPreparedPlan(exchange.child))) + case other => buildPreview(stripPreparation(other), topLevel = false) + } + + /** + * `plan` with the artifacts of a finished plan preparation removed: whole-stage codegen + * wrappers and the columnar transitions Spark inserted. + * + * A subquery arrives inside the outer plan fully prepared - + * `ApplyColumnarRulesAndInsertTransitions` and `CollapseCodegenStages` have both run over it - + * whereas the conversion rules only ever see a plan midway through preparation. A + * `HashAggregateExec` still wrapped in `WholeStageCodegenExec` is left unconverted, so + * previewing the prepared form would report a subquery as falling back that Comet in fact + * accelerates. [[buildPreview]] re-inserts the transitions once conversion is done. + */ + private def stripPreparation(plan: SparkPlan): SparkPlan = plan.transformUp { + case WholeStageCodegenExec(child) => child + case InputAdapter(child) => child + case ColumnarToRowExec(child) => child + case RowToColumnarExec(child) => child + } + + private def _apply(plan: SparkPlan, forPreview: Boolean): SparkPlan = { // We shouldn't transform Spark query plan if Comet is not loaded. if (!isCometLoaded(conf)) return plan @@ -589,6 +823,24 @@ case class CometExecRule(session: SparkSession) plan } } else { + // Plan-only mode: build the Comet plan Comet would have executed, log it, and return + // the original plan unchanged. `CometScanRule` also short-circuits in this mode, so + // `plan` is still pure Spark; `reportPlanOnlyCoverage` rebuilds a scan-wrapped copy for + // the preview. Placed before `normalizePlan`/`RewriteJoin`/`tagUnsafePartialAggregates` + // so their work is not wasted on the discarded outer pass. + if (!forPreview && CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) { + val executionId = Option( + session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)) + if (CometExecRule.shouldReportPlanOnly( + executionId, + plan, + queryStagePrep, + conf.adaptiveExecutionEnabled)) { + reportPlanOnlyCoverage(plan) + } + return plan + } + val normalizedPlan = normalizePlan(plan) val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { 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 2c15bb5e4e..e852ddc1e4 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -64,6 +64,12 @@ case class CometScanRule(session: SparkSession) private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get() override def apply(plan: SparkPlan): SparkPlan = { + // Plan-only mode: leave the plan untouched. `CometExecRule.reportPlanOnlyCoverage` calls + // `_apply` directly to bypass this short-circuit when it needs the wrapping for the + // preview plan. + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) { + return plan + } val newPlan = _apply(plan) if (showTransformations && !newPlan.fastEquals(plan)) { logInfo(s""" @@ -74,7 +80,7 @@ case class CometScanRule(session: SparkSession) newPlan } - private def _apply(plan: SparkPlan): SparkPlan = { + private[rules] def _apply(plan: SparkPlan): SparkPlan = { if (!isCometLoaded(conf)) return plan // Comet does not support structured streaming. The parallel guard in 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..88783c6572 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -76,6 +76,19 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession) .getOrElse(withRevertedStages) } + /** + * Applies the revert decision to every stage of `plan`, regardless of whether AQE is enabled. + * + * `apply` picks the AQE branch when AQE is on because Spark hands it a single query stage at a + * time there, so only the topmost stage of `plan` is considered. Callers holding a whole plan + * that has not been split into stages - the plan-only preview in + * `CometExecRule.reportPlanOnlyCoverage` - need every shuffle boundary visited to see the + * reversions that the real per-stage applications would make. + */ + private[rules] def applyToAllStages(plan: SparkPlan): SparkPlan = { + if (!enabled) plan else applyForNonAQE(plan) + } + /** * Reverts the stage if C2R count exceeds threshold. Wraps in R2C if exchange needs columnar. */ diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 5444a89fa3..611c7c0e37 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -21,6 +21,7 @@ package org.apache.comet.rules import scala.util.Random +import org.apache.logging.log4j.Level import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo} @@ -30,10 +31,12 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} -import org.apache.comet.{CometConf, CometExplainInfo} +import org.apache.comet.{CometConf, CometCoverageStats, CometExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} @@ -797,4 +800,348 @@ class CometExecRuleSuite extends CometTestBase { } } + /** + * With plan-only mode enabled, plan `sql` over a fresh Parquet table and assert nothing was + * offloaded to native. `useV1` toggles between `USE_V1_SOURCE_LIST=parquet` (the V1 + * `FileSourceScanExec` path) and `USE_V1_SOURCE_LIST=""` (the V2 `BatchScanExec` path). + * + * The source list has to be set before the table is created. `withParquetTable` resolves the + * relation through `spark.read` and registers the result as a temp view, so the choice of V1 or + * V2 is baked in at that point; changing the config afterwards leaves both variants planning a + * `FileSourceScanExec`. The scan node is asserted below so that this cannot regress unnoticed. + */ + private def runPlanOnlyAndAssertReverted( + sql: String, + useV1: Boolean = true, + aqe: Boolean = true): Unit = { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> (if (useV1) "parquet" else ""), + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val executed = stripAQEPlan(spark.sql(sql).queryExecution.executedPlan) + val cometNodes = executed.collect { case p: CometPlan => p } + assert( + cometNodes.isEmpty, + s"plan-only mode must not offload; found Comet operators: $cometNodes") + if (useV1) { + assert( + executed.exists(_.isInstanceOf[FileSourceScanExec]), + s"expected the V1 scan path, got:\n$executed") + } else { + assert( + executed.exists(_.isInstanceOf[BatchScanExec]), + s"expected the V2 scan path, got:\n$executed") + } + } + } + } + + for { + useV1 <- Seq(true, false) + aqe <- Seq(true, false) + } { + val label = s"${if (useV1) "V1" else "V2"} scan, AQE=$aqe" + test(s"plan-only mode: $label") { + runPlanOnlyAndAssertReverted( + "SELECT _2, count(*) FROM tbl GROUP BY _2", + useV1 = useV1, + aqe = aqe) + } + } + + test("plan-only mode: scalar subquery is also reverted") { + runPlanOnlyAndAssertReverted("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)") + } + + test("plan-only mode: same query with the config off runs on Comet") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "false") { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val plan = + spark.sql("SELECT _2, count(*) FROM tbl GROUP BY _2").queryExecution.executedPlan + val cometNodes = stripAQEPlan(plan).collect { case p: CometPlan => p } + assert(cometNodes.nonEmpty, "expected Comet operators when plan-only mode is disabled") + } + } + } + + private val PLAN_ONLY_PREFIX = "[Comet plan-only]" + + /** Runs `f` and returns the `[Comet plan-only]` reports that `CometExecRule` logged. */ + private def capturePlanOnlyReports(f: => Unit): Seq[String] = { + val appender = new LogAppender("Comet plan-only reports") + withLogAppender( + appender, + loggerNames = Seq(classOf[CometExecRule].getName), + level = Some(Level.WARN)) { + f + } + appender.loggingEvents + .map(_.getMessage.getFormattedMessage) + .filter(_.startsWith(PLAN_ONLY_PREFIX)) + .toSeq + } + + /** The `Comet accelerated N out of M eligible operators` counts in a plan-only report. */ + private def coverageOf(report: String): (Int, Int) = { + val pattern = """Comet accelerated (\d+) out of (\d+) eligible operators""".r + pattern + .findFirstMatchIn(report) + .map(m => (m.group(1).toInt, m.group(2).toInt)) + .getOrElse(fail(s"report has no coverage summary:\n$report")) + } + + // The outer query is planned after any subquery it contains, so a report slot owned by the + // first plan Spark prepares would describe the subquery and never the query being evaluated. + for (aqe <- Seq(true, false)) { + test(s"plan-only mode: report describes the outer query, not just a subquery (AQE=$aqe)") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val reports = capturePlanOnlyReports { + spark.sql("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)").collect() + } + assert(reports.nonEmpty, "expected a plan-only report") + // The outer plan's Filter appears in no subquery plan, so its presence proves the + // outer query was reported and not suppressed by the subquery's earlier planning. + assert( + reports.exists(_.contains("Filter")), + s"no report describes the outer query:\n${reports.mkString("\n\n")}") + assert( + reports.distinct.size == reports.size, + s"the same plan was reported more than once:\n${reports.mkString("\n\n")}") + // Expected: one report for the subquery plan, one for the outer plan. AQE applies the + // rule again per stage and per re-optimization; those must not add reports. + assert( + reports.size <= 4, + s"expected a report per planned plan, got ${reports.size}:\n" + + reports.mkString("\n\n")) + } + } + } + } + + test("plan-only mode: coverage accounts for post-columnar stage reversion") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + // AQE off so that Spark applies the post-columnar rules to the whole plan exactly once, + // which is what the preview does, making the two directly comparable. + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false") { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val query = "SELECT _2, count(*), sum(_1) FROM tbl GROUP BY _2" + + // Comet accelerates part of this plan when the stage is left alone, so a preview that + // stopped before the post-columnar rules would report a non-zero count below. + withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val df = sql(query) + df.collect() + assert(CometCoverageStats.forPlan(df.queryExecution.executedPlan).cometOperators > 0) + } + + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + // What Comet really executes with reversion enabled. + val df = sql(query) + df.collect() + val executed = CometCoverageStats.forPlan(df.queryExecution.executedPlan) + + // The assertions stay inside the config block: on Spark 3.4 and 3.5 `withSQLConf` is + // declared to return `Unit`, so a value cannot be carried out of one. + withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val reports = capturePlanOnlyReports(sql(query).collect()) + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + assert( + coverageOf(reports.head) == + (executed.cometOperators, executed.cometOperators + executed.sparkOperators), + s"report disagrees with the executed plan ($executed):\n${reports.head}") + } + } + } + } + } + + // `df.rdd.count()` and reading `executedPlan` without running an action can plan - and in the + // first case execute AQE stages - without a SQL execution ID installed, so the reporting state + // cannot be scoped to one. Each of these must still produce exactly one report per query, not + // one per stage and re-optimization. + for (aqe <- Seq(true, false)) { + test(s"plan-only mode: one report per query without a SQL execution ID (AQE=$aqe)") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val query = "SELECT _2, count(*) FROM tbl GROUP BY _2" + + val viaRdd = capturePlanOnlyReports(spark.sql(query).rdd.count()) + assert( + viaRdd.size == 1, + s"expected one report for df.rdd.count(), got ${viaRdd.size}:\n" + + viaRdd.mkString("\n\n")) + + val viaExecutedPlan = + capturePlanOnlyReports(spark.sql(query).queryExecution.executedPlan) + assert( + viaExecutedPlan.size == 1, + s"expected one report for executedPlan, got ${viaExecutedPlan.size}:\n" + + viaExecutedPlan.mkString("\n\n")) + } + } + } + } + + // A scalar subquery is planned as a top-level plan of its own and substituted into the outer + // plan, and extended explain counts the plans owned by a node's expressions. The outer preview + // therefore has to preview its subqueries too, or it reports operators Comet does accelerate as + // Spark. AQE is off here so the preview and the executed plan are the same single pass; see the + // user guide for why the two can differ under AQE. + test("plan-only mode: outer report coverage matches normal planning for a scalar subquery") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val query = "SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)" + + // What Comet really executes, subquery operators included. + val df = sql(query) + df.collect() + val executed = CometCoverageStats.forPlan(df.queryExecution.executedPlan) + assert( + executed.cometOperators > 0, + "test query must be partly accelerated for the comparison to mean anything") + + withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val reports = capturePlanOnlyReports(sql(query).collect()) + // One report for the separately planned subquery, one for the outer query. Only the + // outer one describes the Filter. + val outer = reports.filter(_.contains("Filter")) + assert( + outer.size == 1, + s"expected exactly one report for the outer query, got ${outer.size}:\n" + + reports.mkString("\n\n")) + assert( + coverageOf(outer.head) == + (executed.cometOperators, executed.cometOperators + executed.sparkOperators), + s"outer report disagrees with the executed plan ($executed):\n${outer.head}") + } + } + } + } + + // When a shuffle materializes empty, AQE re-plans from the logical plan and can collapse the + // whole tree to an empty relation. That plan shares no nodes with the one already reported and + // holds no query stages either, so neither the mark nor the stage check recognizes it. + for (aqe <- Seq(true, false)) { + test(s"plan-only mode: one report when the plan becomes empty (AQE=$aqe)") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val reports = capturePlanOnlyReports { + spark + .sql("SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2") + .collect() + } + assert( + reports.size == 1, + s"expected one report, got ${reports.size}:\n${reports.mkString("\n\n")}") + // The one report must describe the query, not the empty relation AQE replaced it with. + assert( + reports.head.contains("HashAggregate"), + s"the report does not describe the query:\n${reports.head}") + } + } + } + + /** Registers a `fact` table partitioned on the join key plus a small `dim` table. */ + private def withDppTables(f: => Unit): Unit = { + withTempDir { dir => + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + val sess = spark + import sess.implicits._ + (0 until 400) + .map(i => (i, i % 10, s"f$i")) + .toDF("fact_id", "fact_key", "fact_str") + .write + .partitionBy("fact_key") + .parquet(s"${dir.getAbsolutePath}/fact") + (0 until 10) + .map(i => (i, i, s"d$i")) + .toDF("dim_id", "dim_key", "dim_str") + .write + .parquet(s"${dir.getAbsolutePath}/dim") + } + spark.read.parquet(s"${dir.getAbsolutePath}/fact").createOrReplaceTempView("fact") + spark.read.parquet(s"${dir.getAbsolutePath}/dim").createOrReplaceTempView("dim") + withTempView("fact", "dim")(f) + } + } + + // `PlanDynamicPruningFilters` prepares the DPP build plan and only wraps it in a + // `BroadcastExchangeExec` afterwards, so the stage `RevertNativeForTransitionHeavyStages` judged + // is the exchange's child. A preview that hands it the exchange instead leaves that child's top + // boundary open, the reversion never fires, and the report claims acceleration the executed plan + // does not have. Reversion is forced on here so a missed reversion changes the number. + test("plan-only mode: outer report coverage matches normal planning for a DPP subquery") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + // AQE off: the preview describes the pre-adaptive plan, so only the non-adaptive plan is + // comparable. Comet also cannot currently run this query with AQE on and reversion enabled. + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false") { + withDppTables { + val query = "SELECT f.fact_id, f.fact_str, d.dim_str FROM fact f " + + "JOIN dim d ON f.fact_key = d.dim_key WHERE d.dim_id < 10" + + val df = sql(query) + df.collect() + val plan = df.queryExecution.executedPlan + // `exists` walks children only, and a DPP subquery hangs off the scan's expressions. + assert( + plan.collectWithSubqueries { case p: SubqueryBroadcastExec => p }.nonEmpty, + s"test query must produce a DPP subquery:\n$plan") + val executed = CometCoverageStats.forPlan(plan) + + withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val reports = capturePlanOnlyReports(sql(query).collect()) + // One report for the separately prepared DPP build plan, one for the outer query. Only + // the outer one describes the join. + val outer = reports.filter(_.contains("BroadcastHashJoin")) + assert( + outer.size == 1, + s"expected exactly one report for the outer query, got ${outer.size}:\n" + + reports.mkString("\n\n")) + assert( + coverageOf(outer.head) == + (executed.cometOperators, executed.cometOperators + executed.sparkOperators), + s"outer report disagrees with the executed plan ($executed):\n${outer.head}") + } + } + } + } + }