diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 32fe297e67d..6d2b3615103 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -208,6 +208,31 @@ 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 plans every query as it normally would, logs the resulting Comet plan and +coverage summary to the driver log with the prefix `[Comet plan-only]`, and then executes the +original plan on Spark. Use it to estimate how much of a workload Comet would accelerate +without changing execution. The report is the same annotated plan that +`spark.comet.explain.format=verbose` produces, and reflects Comet's post-columnar rules, so a +stage Comet would revert to Spark for having too many transitions is reported as reverted. +Requires `spark.comet.exec.enabled=true`. + +Keep the following in mind when reading the reports: + +- Spark plans scalar and dynamic partition pruning subqueries separately from the query that + contains them, so each gets its own report alongside the outer query's. The outer report + also counts its subqueries, so do not add the reports together. +- Only the JVM side of planning runs. Anything that would fail when DataFusion builds the + native plan still counts as accelerated, so treat the percentage as an upper bound. +- Comet's split Iceberg V2 write (`spark.comet.write.iceberg.splitOperator.enabled`) is + declined in plan-only mode, so such writes run on, and are reported as, Spark. +- Under AQE the report describes the plan before any adaptive re-planning, so coverage of the + plan that finally executes can differ. In particular, AQE plans subqueries into the outer + query only after the report is produced, so the outer report counts their operators as + Spark. For subquery-heavy queries, read the per-subquery reports or disable AQE for the + evaluation run. + ## Programmatic Access to Fallback Reasons 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 d371f1ba44c..d279ce69fde 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -785,6 +785,17 @@ 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 logs the plan it would have executed, with a coverage " + + "summary, to the driver log and then lets Spark execute the query unchanged. Native " + + "planning failures are not detected, so the coverage can be optimistic. Requires " + + "`spark.comet.exec.enabled=true`.") + .booleanConf + .createWithDefault(false) + 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 749a3d6ecb0..f959d7bf9b5 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf._ import org.apache.comet.iceberg.IcebergWriteStrategy -import org.apache.comet.rules.{CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} +import org.apache.comet.rules.{CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule} import org.apache.comet.shims.ShimCometSparkSessionExtensions /** @@ -96,7 +96,9 @@ class CometSparkSessionExtensions // Registered before CometRule so tags are in place when conversion runs. // No-op on Spark 3.5+; see CometSpark34AqeDppFallbackRule's class docstring. injectPreSpark35QueryStagePrepRuleShim(extensions, CometSpark34AqeDppFallbackRule) - extensions.injectQueryStagePrepRule { session => CometRule(session) } + extensions.injectQueryStagePrepRule { session => + CometRule(session, queryStagePrep = true) + } injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters) injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) } @@ -106,8 +108,7 @@ class CometSparkSessionExtensions override def preColumnarTransitions: Rule[SparkPlan] = CometRule(session) override def postColumnarTransitions: Rule[SparkPlan] = { - val rules = - Seq(RevertNativeForTransitionHeavyStages(session), EliminateRedundantTransitions(session)) + val rules = CometRule.postColumnarRules(session) plan => rules.foldLeft(plan) { case (p, rule) => rule(p) } } } diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index 37535f80500..8d20b33868e 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -42,18 +42,24 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator { def generateExtendedInfo(plan: SparkPlan): String = { CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.get() match { case CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE => - // Generates the extended info in a verbose manner, printing each node along with the - // extended information in a tree display. - val planStats = new CometCoverageStats() - val outString = new StringBuilder() - generateTreeString(getActualPlan(plan), 0, Seq(), 0, outString, planStats) - s"${outString.toString()}\n$planStats" + generateVerboseInfo(plan) case CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_FALLBACK => // Generates the extended info as a list of fallback reasons getFallbackReasons(plan).mkString("\n").trim } } + /** + * The `verbose` format regardless of `spark.comet.explain.format`: each node along with its + * extended information in a tree display, followed by the coverage summary. + */ + def generateVerboseInfo(plan: SparkPlan): String = { + val planStats = new CometCoverageStats() + val outString = new StringBuilder() + generateTreeString(getActualPlan(plan), 0, Seq(), 0, outString, planStats) + s"${outString.toString()}\n$planStats" + } + def getFallbackReasons(plan: SparkPlan): Seq[String] = { fallbackReasons(plan).toSeq.sorted } diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala index c501c900973..e8790079797 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala @@ -38,6 +38,10 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { if (!CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.get(session.sessionState.conf)) { return Nil } + // Planner strategies run before CometRule, so plan-only mode needs its own guard here. + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(session.sessionState.conf)) { + return Nil + } plan match { case ad: AppendData => diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala index 9b9639ac188..fa2fb7dc32e 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -19,9 +19,101 @@ package org.apache.comet.rules +import scala.collection.mutable +import scala.util.control.NonFatal + import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, ExecSubqueryExpression, InputAdapter, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, InsertAdaptiveSparkPlan, QueryStageExec} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange} +import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery + +import org.apache.comet.{CometConf, ExtendedExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.isCometLoaded +import org.apache.comet.shims.ShimCometStreaming + +object CometRule { + + /** Comet's post-columnar rules, shared by `CometColumnar` and the plan-only preview. */ + def postColumnarRules(session: SparkSession, wholePlan: Boolean = false): Seq[Rule[SparkPlan]] = + Seq( + RevertNativeForTransitionHeavyStages(session, wholePlan), + EliminateRedundantTransitions(session)) + + /** + * Canonical hashes of the subquery plans reported for the query this thread is preparing. Spark + * prepares a query's subqueries synchronously, before the query itself, so the scope is reset + * when the query's own plan arrives. This works whether or not a SQL execution ID is set. + */ + private val reportedSubqueries = new ThreadLocal[mutable.Set[Int]] { + override def initialValue(): mutable.Set[Int] = mutable.Set.empty + } + + /** + * Marks the root of a reported plan. Catalyst copies tags onto replacement nodes, so later + * applications of the rule can recognize the plan however Spark rewrote it in between. + */ + private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.planOnlyReported") + + /** Where in Spark's planning the rule is running, read off the call stack. */ + private case class PlanningContext(replanning: Boolean, subquery: Boolean) + + private def planningContext(): PlanningContext = { + val frames = Thread.currentThread().getStackTrace + def within(cls: Class[_], method: String): Boolean = + frames.exists(f => f.getMethodName == method && f.getClassName == cls.getName) + PlanningContext( + replanning = within(classOf[AdaptiveSparkPlanExec], "reOptimize"), + // Subqueries are prepared by `PlanSubqueries` and `PlanDynamicPruningFilters` without AQE, + // and by `InsertAdaptiveSparkPlan` and `PlanAdaptiveDynamicPruningFilters` with it. + subquery = within(QueryExecution.getClass, "prepareExecutedPlan") || + within(classOf[InsertAdaptiveSparkPlan], "compileSubquery")) + } + + /** + * Whether plan-only mode should report `plan`, marking it reported if so. + * + * Under AQE the rule sees one query several times: the initial plan (prep rule, the one to + * report), the same plan wrapped in `AdaptiveSparkPlanExec`, each query stage, each + * re-optimization and the final plan. Only the first is reported. Each scalar or DPP subquery + * is prepared as a plan of its own and gets its own report. + */ + private[comet] def shouldReportPlanOnly( + plan: SparkPlan, + queryStagePrep: Boolean, + aqeEnabled: Boolean): Boolean = { + // Under AQE the columnar rule sees each new query stage rooted at its Exchange. + val isQueryStage = aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange] + // Already reported: a re-optimized plan holds query stages, and a final plan carries the + // mark. `AdaptiveSparkPlanExec` is a leaf to `exists`, so it is matched directly. + val isReapplication = plan.exists(p => + p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec] || + p.getTagValue(PLAN_ONLY_REPORTED).isDefined) + if (isQueryStage || isReapplication) { + // Execution is under way, so any subquery prepared from here on belongs to a new scope. + reportedSubqueries.get().clear() + false + } else { + // Mark even when not reporting, so the final plan built from a re-plan is recognized. + plan.setTagValue(PLAN_ONLY_REPORTED, ()) + val context = planningContext() + if (context.replanning) { + // AQE re-plans a query mid-execution, for example to an empty relation once a stage + // materializes empty, and the result can share no nodes or stages with the reported plan. + false + } else if (context.subquery) { + // A subquery referenced twice is prepared twice, differing only in expression IDs. + // Tags are not part of the canonical form. + reportedSubqueries.get().add(plan.canonicalized.hashCode()) + } else { + reportedSubqueries.get().clear() + true + } + } + } +} /** * Comet's plan conversion pass: scan conversion followed by operator conversion. @@ -36,11 +128,103 @@ import org.apache.spark.sql.execution.SparkPlan * delegates to their `apply`. Spark's own plan change log sees one rule: query-stage preparation * logs this pass as `org.apache.comet.rules.CometRule`, which is the name * `spark.sql.planChangeLog.rules` has to match. + * + * @param queryStagePrep + * true for the `injectQueryStagePrepRule` instance, which sees the whole initial plan under + * AQE. Only plan-only reporting reads it. */ -case class CometRule(session: SparkSession) extends Rule[SparkPlan] { +case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) + extends Rule[SparkPlan] { private val scanRule = CometScanRule(session) private val execRule = CometExecRule(session) - override def apply(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan)) + override def apply(plan: SparkPlan): SparkPlan = { + if (planOnlyApplies(plan)) { + reportPlanOnlyCoverage(plan) + plan + } else { + convert(plan) + } + } + + private def convert(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan)) + + /** Mirrors the conversion rules' own guards; plan-only is scoped to exec being enabled. */ + private def planOnlyApplies(plan: SparkPlan): Boolean = + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && + isCometLoaded(conf) && + !ShimCometStreaming.isStreamingPlan(plan) && + CometConf.COMET_EXEC_ENABLED.get(conf) + + /** Logs the Comet plan for `plan` unless already reported. Never fails the query. */ + private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = { + try { + if (CometRule.shouldReportPlanOnly(plan, queryStagePrep, conf.adaptiveExecutionEnabled)) { + val preview = buildPreview(plan, topLevel = true) + logWarning( + s"[Comet plan-only]\n${new ExtendedExplainInfo().generateVerboseInfo(preview)}") + } + } catch { + case NonFatal(e) => + logWarning("[Comet plan-only] could not build a coverage report for this query", e) + } + } + + /** + * The plan Comet would execute for `plan`: conversion, columnar transitions, then the + * post-columnar rules, with stage reversion visiting every stage since this is the whole plan. + * + * @param topLevel + * false for subquery plans, which Spark prepares without `ReuseExchangeAndSubquery`. + */ + private def buildPreview(plan: SparkPlan, topLevel: Boolean): SparkPlan = { + val converted = convert(previewSubqueriesOf(plan)) + val withTransitions = + ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted) + val preview = CometRule + .postColumnarRules(session, wholePlan = true) + .foldLeft(withTransitions) { case (p, rule) => rule(p) } + if (topLevel) ReuseExchangeAndSubquery.apply(preview) else preview + } + + /** + * `plan` with each subquery's plan replaced by its preview. Extended explain counts subquery + * operators, and normal planning has already converted them, so leaving them would understate + * coverage. + */ + private def previewSubqueriesOf(plan: SparkPlan): SparkPlan = { + plan.transformAllExpressions { case subquery: ExecSubqueryExpression => + subquery.withNewPlan(previewSubquery(subquery.plan)) + } + } + + private def previewSubquery(subquery: BaseSubqueryExec): BaseSubqueryExec = subquery match { + case reused: ReusedSubqueryExec => reused.copy(child = previewSubquery(reused.child)) + case other => + other.withNewChildren(Seq(previewPreparedPlan(other.child))).asInstanceOf[BaseSubqueryExec] + } + + /** + * `PlanDynamicPruningFilters` prepares a DPP build plan before wrapping it in a + * `BroadcastExchangeExec`, so preview the exchange's child on its own to keep the stage + * boundary that stage reversion saw. + */ + private def previewPreparedPlan(plan: SparkPlan): SparkPlan = plan match { + case exchange: BroadcastExchangeExec => + exchange.withNewChildren(Seq(previewPreparedPlan(exchange.child))) + case other => buildPreview(stripPreparation(other), topLevel = false) + } + + /** + * Removes codegen wrappers and transitions from an already-prepared subquery plan, since the + * conversion rules expect a plan from before those are inserted. [[buildPreview]] re-inserts + * the transitions. + */ + private def stripPreparation(plan: SparkPlan): SparkPlan = plan.transformUp { + case WholeStageCodegenExec(child) => child + case InputAdapter(child) => child + case ColumnarToRowExec(child) => child + case RowToColumnarExec(child) => child + } } diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala index a56e477b035..bbdc6db6e57 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -36,8 +36,12 @@ import org.apache.comet.serde.QueryPlanSerde * Reverts a query stage to Spark row-based execution when it has too many columnar-to-row (C2R) * transitions. Each C2R indicates Comet could not keep execution columnar and had to fall back. * With columnar shuffle enabled, each C2R implies a corresponding R2C round-trip. + * + * @param wholePlan + * visit every stage even under AQE, where Spark normally hands this rule one stage at a time. + * Set by the plan-only preview, which holds the whole plan. */ -case class RevertNativeForTransitionHeavyStages(session: SparkSession) +case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan: Boolean = false) extends Rule[SparkPlan] with Logging { @@ -47,7 +51,7 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession) override def apply(plan: SparkPlan): SparkPlan = { if (!enabled) return plan - if (session.sessionState.conf.adaptiveExecutionEnabled) { + if (session.sessionState.conf.adaptiveExecutionEnabled && !wholePlan) { applyForAQE(plan) } else { applyForNonAQE(plan) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index 37ddb1872e8..0d7d9935b22 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -79,6 +79,26 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } } + test("plan-only mode leaves the write with Spark") { + withDetectionCatalog { dir => + createTable(dir, "plan_only", partitionSpec = "") + // IcebergWriteStrategy runs before CometRule, so it needs its own plan-only guard. + // withSQLConf returns Unit on Spark 3.4/3.5, hence the var. + var plan: SparkPlan = null + withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + plan = captureWritePlan("plan_only", allowWriteFailure = false) { + spark.sql(s"INSERT INTO $catalog.$ns.plan_only VALUES (1, 'us', 1.0)") + } + } + assert( + findWriteExec(plan).isEmpty, + s"plan-only mode must not split the write into Comet's two-operator shape:\n$plan") + assert( + !containsCometWriteExec(plan), + s"plan-only mode must not offload the write to Comet:\n$plan") + } + } + test("SparkWrite reflection helpers all resolve on the current Iceberg runtime") { withDetectionCatalog { dir => createTable(dir, "refl_probe", partitionSpec = "") 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 8bf60977543..8ad1884acca 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -31,11 +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, ShuffleQueryStageExec} 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, ExtendedExplainInfo} +import org.apache.comet.{CometConf, CometCoverageStats, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus, withFallbackReason} import org.apache.comet.serde.{CometAggregateExpressionSerde, Compatible, ExprOuterClass, Unsupported} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} @@ -1086,6 +1087,369 @@ class CometExecRuleSuite extends CometTestBase { } } + private val PLAN_ONLY_PREFIX = "[Comet plan-only]" + + /** + * Runs `f` over a Parquet table `tbl`. The source list must be set before the table is created, + * since `withParquetTable` bakes the V1/V2 choice into the temp view. + */ + private def withPlanOnlyTable( + aqe: Boolean = true, + planOnly: Boolean = true, + useV1: Boolean = true)(f: => Unit): Unit = { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> (if (useV1) "parquet" else ""), + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> planOnly.toString) { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl")(f) + } + } + + /** Plans `sql` in plan-only mode and asserts nothing was offloaded to native. */ + private def runPlanOnlyAndAssertReverted( + sql: String, + useV1: Boolean = true, + aqe: Boolean = true): Unit = { + withPlanOnlyTable(aqe = aqe, useV1 = useV1) { + val executed = stripAQEPlan(spark.sql(sql).queryExecution.executedPlan) + val cometNodes = executed.collect { case p: CometPlan => p } + assert( + cometNodes.isEmpty, + s"plan-only mode must not offload; found Comet operators: $cometNodes") + val expectedScan = if (useV1) classOf[FileSourceScanExec] else classOf[BatchScanExec] + assert( + executed.exists(p => expectedScan.isInstance(p)), + s"expected ${expectedScan.getSimpleName}, got:\n$executed") + } + } + + /** Runs `f` and returns the `[Comet plan-only]` reports that `CometRule` logged. */ + private def capturePlanOnlyReports(f: => Unit): Seq[String] = { + val appender = new LogAppender("Comet plan-only reports") + withLogAppender( + appender, + loggerNames = Seq(classOf[CometRule].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")) + } + + /** + * Runs `query` normally, then in plan-only mode, and asserts the one report containing `marker` + * has the same coverage as the executed plan. Returns the executed plan. + */ + private def assertReportMatchesExecuted(query: String, marker: String): SparkPlan = { + val df = sql(query) + df.collect() + val plan = df.queryExecution.executedPlan + val executed = CometCoverageStats.forPlan(plan) + withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val reports = capturePlanOnlyReports(sql(query).collect()) + val matching = reports.filter(_.contains(marker)) + assert( + matching.size == 1, + s"expected one report containing '$marker', got:\n${reports.mkString("\n\n")}") + assert( + coverageOf(matching.head) == + (executed.cometOperators, executed.cometOperators + executed.sparkOperators), + s"report disagrees with the executed plan ($executed):\n${matching.head}") + } + plan + } + + for { + useV1 <- Seq(true, false) + aqe <- Seq(true, false) + } { + val label = s"${if (useV1) "V1" else "V2"} scan, AQE=$aqe" + test(s"plan-only mode: $label") { + runPlanOnlyAndAssertReverted( + "SELECT _2, count(*) FROM tbl GROUP BY _2", + useV1 = useV1, + aqe = aqe) + } + } + + test("plan-only mode: scalar subquery is also reverted") { + runPlanOnlyAndAssertReverted("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)") + } + + test("plan-only mode: same query with the config off runs on Comet") { + withPlanOnlyTable(planOnly = false) { + val plan = + spark.sql("SELECT _2, count(*) FROM tbl GROUP BY _2").queryExecution.executedPlan + val cometNodes = stripAQEPlan(plan).collect { case p: CometPlan => p } + assert(cometNodes.nonEmpty, "expected Comet operators when plan-only mode is disabled") + } + } + + // Subqueries are planned before the outer query, which must still get its own report. + for (aqe <- Seq(true, false)) { + test(s"plan-only mode: report describes the outer query, not just a subquery (AQE=$aqe)") { + withPlanOnlyTable(aqe = aqe) { + val reports = capturePlanOnlyReports { + spark.sql("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)").collect() + } + // Filter appears only in the outer plan. + assert( + reports.exists(_.contains("Filter")), + s"no report describes the outer query:\n${reports.mkString("\n\n")}") + // One for the subquery, one for the outer query. + assert( + reports.size == 2, + s"expected two reports, got ${reports.size}:\n${reports.mkString("\n\n")}") + } + } + } + + test("plan-only mode: coverage accounts for post-columnar stage reversion") { + withSQLConf(CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false") { + // AQE off: Spark applies the post-columnar rules to the whole plan once, as the preview does. + withPlanOnlyTable(aqe = false, planOnly = false) { + // SUM and MAX can be split across a stage boundary in either direction, which leaves the + // stage free to revert. COUNT cannot, and would block the reversion under test. + val query = "SELECT _2, sum(_1), max(_1) FROM tbl GROUP BY _2" + + var unreverted = 0 + withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val df = sql(query) + df.collect() + unreverted = CometCoverageStats.forPlan(df.queryExecution.executedPlan).cometOperators + assert(unreverted > 0) + } + + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val executed = CometCoverageStats.forPlan(assertReportMatchesExecuted(query, "")) + // Guard against a vacuous test: reversion must actually fire. + assert( + executed.cometOperators < unreverted, + s"stage reversion did not fire, so this test is vacuous: $executed") + } + } + } + } + + // These paths plan without a SQL execution ID and must still report once per query. + for (aqe <- Seq(true, false)) { + test(s"plan-only mode: one report per query without a SQL execution ID (AQE=$aqe)") { + withPlanOnlyTable(aqe = aqe) { + val query = "SELECT _2, count(*) FROM tbl GROUP BY _2" + + val viaRdd = capturePlanOnlyReports(spark.sql(query).rdd.count()) + assert( + viaRdd.size == 1, + s"expected one report for df.rdd.count(), got ${viaRdd.size}:\n" + + viaRdd.mkString("\n\n")) + + val 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")) + } + } + } + + // The outer report must count the subquery's converted operators. AQE off so the preview and + // the executed plan are the same single pass. + test("plan-only mode: outer report coverage matches normal planning for a scalar subquery") { + withPlanOnlyTable(aqe = false, planOnly = false) { + val plan = assertReportMatchesExecuted( + "SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM tbl)", + marker = "Filter") + assert( + CometCoverageStats.forPlan(plan).cometOperators > 0, + "test query must be partly accelerated for the comparison to mean anything") + } + } + + // AQE can collapse a plan whose stage materializes empty into an empty relation that shares + // nothing with the reported plan. + for { + aqe <- Seq(true, false) + // `toRdd.count` runs without a SQL execution ID, like PySpark's `df.rdd`. + (action, runIt) <- Seq[(String, org.apache.spark.sql.DataFrame => Unit)]( + "collect" -> (df => df.collect()), + "toRdd.count" -> (df => df.queryExecution.toRdd.count())) + // In the second shape `RemoveRedundantSorts` drops the root sort, so the mark lands on the + // join and the empty re-plan's root inherits the unmarked sort's tags. In the third the + // global aggregate survives the empty join, so the re-plan is not itself empty. The fourth + // is empty from the start, which must not be mistaken for a re-plan. + (shape, query, marker) <- Seq( + ( + "aggregate", + "SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2", + "HashAggregate"), + ( + "join under a removed root sort", + """SELECT a.id FROM range(0, 20, 1, 2) a + |JOIN range(0, 20, 1, 2) b ON a.id % 7 = b.id % 7 + |WHERE a.id < 0 + |SORT BY a.id % 7""".stripMargin, + "SortMergeJoin"), + ( + "global aggregate over a join", + """SELECT count(*) FROM range(0, 20, 1, 2) a + |JOIN range(0, 20, 1, 2) b ON a.id % 7 = b.id % 7 + |WHERE a.id < 0""".stripMargin, + "SortMergeJoin"), + ("initially empty", "SELECT id FROM range(0) DISTRIBUTE BY id", "Exchange")) + } { + test(s"plan-only mode: one report when the plan becomes empty ($shape, AQE=$aqe, $action)") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + val reports = capturePlanOnlyReports(runIt(spark.sql(query))) + assert( + reports.size == 1, + s"expected one report, got ${reports.size}:\n${reports.mkString("\n\n")}") + assert( + reports.head.contains(marker), + s"the report does not describe the query:\n${reports.head}") + } + } + } + + for { + aqe <- Seq(true, false) + (action, runIt) <- Seq[(String, org.apache.spark.sql.DataFrame => Unit)]( + "collect" -> (df => df.collect()), + "toRdd.count" -> (df => df.queryExecution.toRdd.count())) + } { + test(s"plan-only mode: a subquery referenced twice is reported once (AQE=$aqe, $action)") { + withPlanOnlyTable(aqe = aqe) { + val reports = capturePlanOnlyReports { + runIt(spark.sql("""SELECT _1, + | (SELECT max(_2) FROM tbl + | WHERE _1 > (SELECT min(_2) FROM tbl)) AS a, + | (SELECT max(_2) FROM tbl + | WHERE _1 > (SELECT min(_2) FROM tbl)) AS b + |FROM tbl""".stripMargin)) + } + // The `min` subquery, the `max` subquery and the outer query. + assert( + reports.size == 3, + s"expected three reports, got ${reports.size}:\n${reports.mkString("\n\n")}") + } + } + } + + test("plan-only mode: the report does not depend on spark.comet.explain.format") { + withSQLConf( + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_FALLBACK) { + withPlanOnlyTable() { + val reports = capturePlanOnlyReports(sql("SELECT _1 + 1 FROM tbl").collect()) + assert(reports.size == 1, s"expected one report, got:\n${reports.mkString("\n\n")}") + val (accelerated, eligible) = coverageOf(reports.head) + assert( + accelerated > 0 && accelerated == eligible, + s"unexpected coverage:\n${reports.head}") + } + } + } + + private val dppQuery = "SELECT f.fact_id, f.fact_str, d.dim_str FROM fact f " + + "JOIN dim d ON f.fact_key = d.dim_key WHERE d.dim_id < 10" + + /** Registers a `fact` table partitioned on the join key plus a small `dim` table. */ + private def withDppTables(f: => Unit): Unit = { + withTempDir { dir => + 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) + } + } + + // The DPP build must be previewed below its BroadcastExchangeExec so that stage reversion fires + // as in real planning. Reversion is forced on so a missed reversion changes the number. + test("plan-only mode: outer report coverage matches normal planning for a DPP subquery") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + // AQE off: the preview describes the pre-adaptive plan. + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + 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 plan = assertReportMatchesExecuted(dppQuery, marker = "BroadcastHashJoin") + // `exists` walks children only, and a DPP subquery hangs off the scan's expressions. + assert( + plan.collectWithSubqueries { case p: SubqueryBroadcastExec => p }.nonEmpty, + s"test query must produce a DPP subquery:\n$plan") + } + } + } + + // Under AQE the DPP build is planned mid-execution by `PlanAdaptiveDynamicPruningFilters`, + // after the outer query was reported, and must still get a report. + test("plan-only mode: a DPP subquery planned mid-execution is reported under AQE") { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { + withDppTables { + val reports = capturePlanOnlyReports(sql(dppQuery).collect()) + assert( + reports.size == 2, + "expected a report for the outer query and one for the DPP build plan, got " + + s"${reports.size}:\n${reports.mkString("\n\n")}") + assert( + reports.count(_.contains("BroadcastHashJoin")) == 1, + s"exactly one report should describe the outer query:\n${reports.mkString("\n\n")}") + } + } + } + + // Per-stage and per-re-optimization applications must not add reports. + test("plan-only mode: one report for a multi-stage query under AQE") { + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withPlanOnlyTable() { + val query = "SELECT a._2, count(*) FROM tbl a JOIN tbl b ON a._1 = b._2 " + + "GROUP BY a._2 ORDER BY 1" + val reports = capturePlanOnlyReports(sql(query).collect()) + assert( + reports.size == 1, + s"expected one report, got ${reports.size}:\n${reports.mkString("\n\n")}") + } + } + } + test("scan conversion must run before operator conversion") { withTempPath { path => createTestDataFrame.write.parquet(path.toString)