You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/25 00:19:39 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-5394-4453c57249aa4bf9c47fd4ff88bd13a09cd2a9f8 created (now a53912f572)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-5394-4453c57249aa4bf9c47fd4ff88bd13a09cd2a9f8
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at a53912f572 feat: add spark.comet.explain.planOnly.enabled (#5394)
This branch includes the following new commits:
new a53912f572 feat: add spark.comet.explain.planOnly.enabled (#5394)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: feat: add spark.comet.explain.planOnly.enabled (#5394)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5394-4453c57249aa4bf9c47fd4ff88bd13a09cd2a9f8
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit a53912f572569be5a8212258c53f43a354c8f57a
Author: Andy Grove <ag...@apache.org>
AuthorDate: Fri Sep 25 00:19:16 2026 +0000
feat: add spark.comet.explain.planOnly.enabled (#5394)
* feat: add spark.comet.explain.planOnly.enabled
Builds the Comet plan Comet would have executed and logs it to the driver
log, then discards it and lets Spark run the query unchanged. Both
CometScanRule and CometExecRule short-circuit when the config is set;
CometExecRule uses thread-local bypass flags on both rules to force normal
behavior while it constructs the preview for the report. The WARN is
deduped per SQL execution ID so AQE emits one report per query.
Closes #5335.
* refactor: simplify plan-only wiring
- Collapse two ThreadLocals into one shared `planOnlyPreviewInProgress`
(dropped `CometScanRule.forceApply`/`withForceApply`).
- Move the plan-only branch to the top of the exec-enabled block so
`normalizePlan`/`RewriteJoin`/`tagUnsafePartialAggregates` are not run
and thrown away.
- Replace hand-rolled `LinkedHashSet` LRU with a synchronized
`LinkedHashMap` + `removeEldestEntry`, matching
`IcebergPlanDataInjector.commonCache`. Drops the `clearPlanOnlyReported`
test helper.
- Collapse V1/V2 × AQE test variants into one loop; fold `assertNoComet`
into `runPlanOnlyAndAssertReverted` and drop redundant `collect`s.
* refactor: replace plan-only ThreadLocal with a forPreview parameter
`_apply` on both rules now takes/uses an explicit `forPreview` flag rather
than reading a thread-local. `reportPlanOnlyCoverage` calls the private
`_apply` methods directly, skipping the `Rule.apply` short-circuits, so
the recursive preview reads the parameter instead of ambient state.
Widens `CometScanRule._apply` to `private[rules]` so `CometExecRule` can
invoke it. Drops `planOnlyPreviewInProgress` and `withPreview`.
* Address review: report the outer query and include post-columnar rules
Plan-only reporting keyed its once-per-query slot on the SQL execution ID
alone. Spark prepares scalar and DPP subqueries as their own top-level plans
before the outer plan reaches the conversion rules, so a nested subquery took
the slot and the outer query - the plan being evaluated - was never reported.
Reports are now keyed on the execution ID and the plan, so the outer query
always gets one, while AQE's per-stage and per-re-optimization applications of
the rule are recognised as re-plans and stay quiet.
The preview also stopped at operator conversion, before Spark inserts columnar
transitions and Comet runs its post-columnar rules. It therefore counted
operators that RevertNativeForTransitionHeavyStages removes and counted
transitions before they existed. The preview now inserts transitions and
applies both post-columnar rules, using a new applyToAllStages entry point so
that every shuffle boundary is visited rather than only the topmost stage.
Tests capture the logged report and assert that the outer query appears under
both AQE modes, and that the reported coverage matches the coverage of the
plan Comet really executes when stage reversion fires.
* Address review: fix Spark 3.x test compile, dedupe without an execution ID, preview subqueries
Move the coverage assertions inside their withSQLConf block. Spark 3.4 and 3.5
declare withSQLConf as returning Unit, so binding its result inferred Unit and
broke test compilation on every supported Spark 3.x build.
Replace the execution-ID keyed report slots with a mark on the plan itself.
df.rdd.count() and reading executedPlan without an action can plan, and in the
first case execute AQE stages, with no execution ID installed, and the old no-ID
path reported unconditionally. Catalyst copies tags onto replacement nodes, so
the mark survives the rewrites between one rule application and the next and
identifies an AQE stage or final-plan pass as a re-plan of something already
reported. The execution-scoped hash is kept for the one case a mark cannot
cover: a subquery referenced twice is prepared twice, as separate but identical
plans.
Preview the plan behind each subquery expression before computing coverage.
Extended explain counts the plans owned by a node's expressions, and normal
planning has already converted those by the time the outer plan arrives, so the
outer report counted accelerated subquery operators as Spark. Two things were
needed to make the numbers line up: a prepared subquery arrives with codegen
wrappers and transitions already inserted, which blocks conversion, so those are
stripped first; and ReuseExchangeAndSubquery is replayed at the top level, since
it is the last preparation step and otherwise each copy of a subquery is counted
separately. With AQE off the outer report now equals CometCoverageStats for the
plan Comet really executes.
Under AQE the outer report still describes the pre-adaptive plan, in which
subqueries have not been planned yet and so count as Spark. Called out in the
user guide.
* Address review: DPP preparation scope, empty adaptive plans, V2 scan coverage
Preview a DPP subquery inside its BroadcastExchangeExec rather than around it.
PlanDynamicPruningFilters prepares the build plan and wraps it in the exchange
afterwards, so the plan that went through the post-columnar rules is the
exchange's child. Handing the exchange to the preview left that child a stage
whose top boundary was the exchange, so RevertNativeForTransitionHeavyStages
never fired and the outer report claimed 4/12 accelerated where the executed
plan has 2/12.
Suppress reports for the rest of an execution once AQE has cut a query stage.
When a shuffle materializes empty, AQE re-plans from the logical plan and can
collapse the tree to an empty relation; that plan shares no nodes with the one
already reported, so the mark does not reach it, and it holds no query stages
either, so the structural checks do not either. Its hash differs too, so both
guards admitted a second 0% report under the same execution ID. Subqueries are
all compiled before the first stage is cut, so this costs no report except for a
subquery AQE only plans once stages are under way, which the user guide now
notes.
Set USE_V1_SOURCE_LIST before creating the Parquet fixture in the plan-only scan
tests. withParquetTable resolves the relation through spark.read and registers
the result as a temp view, so changing the source list afterwards left both the
V1 and the V2 variant planning a FileSourceScanExec. The scan node is now
asserted so this cannot regress unnoticed.
New tests cover an empty adaptive query under AQE on and off, and compare the
outer report for a DPP subquery against CometCoverageStats for the plan Comet
really executes with stage reversion forced on.
* fix: suppress AQE re-plan reports without an execution ID and dedupe reused subqueries
Mark the logical plan a reported physical plan was built from, alongside the
existing mark on the physical tree. When a stage materializes empty AQE
re-plans from the logical plan and can collapse the whole tree to an empty
relation, leaving a physical tree that shares no node with the one reported and
holds no query stage, so only the logical mark reaches it. The mark needs no
SQL execution ID, which the previous stage-tracking guard did, so the duplicate
0% report is now suppressed on the RDD path PySpark's `df.rdd` takes as well.
That also removes the reason the stage-tracking guard existed, and with it the
suppression of every report that followed the first stage cut, so a subquery
AQE only plans once stages are under way - a DPP subquery - gets its report
back.
Key the per-execution dedupe on the canonicalized plan. A subquery referenced
twice is prepared once per reference as plans differing only in expression IDs,
which Spark then collapses to one `ReusedSubqueryExec`; the raw structural hash
saw two plans and reported the same work twice.
* fix: recognize an AQE re-plan that collapsed to an empty relation
A mark on the plan already reported cannot reach the plan AQE builds when a
stage materializes empty, and marking the logical plan only reached it while
the physical root's logical link happened to be the logical root. When a
preparation rule drops the root operator first - `RemoveRedundantSorts`
removing a sort above a sort merge join - the mark lands on the join below it,
empty propagation replaces the join and then the sort, and the replacement root
inherits the unmarked sort's tags.
Recognize that plan by what it is instead. A re-optimized plan holds a
`QueryStageExec` for every stage that materialized; the only way to eliminate
all of them is empty propagation, which leaves a plan Catalyst records as
producing zero rows, so the logical link answers the question directly. The
check is confined to the query-stage-prep rule, so a genuinely empty query is
still reported.
This subsumes the logical mark, which is dropped.
* fix: drop a redundant interpolator and guard the plan-only report
The interpolator broke both scalafix lint jobs. The guard keeps a preview
failure from failing the query, which is the one thing plan-only mode
promises not to do.
* fix: keep plan-only mode from offloading Iceberg V2 writes
`IcebergWriteStrategy` is registered with `injectPlannerStrategy`, not as a
`Rule[SparkPlan]`, so it runs during physical planning ahead of both
conversion rules and their plan-only short-circuit never reaches it. An
Iceberg V2 write was still emitted as Comet's two-operator shape and
executed natively, which is exactly what plan-only mode promises not to do.
Neither `IcebergCommitExec` nor `IcebergWriteExec` extends `CometPlan`, so
the plan-only suite's "no Comet operators" assertion could not see it.
Decline the strategy in plan-only mode. The cost is that the report no
longer describes the write as accelerated; that is documented alongside the
existing native-planner caveat.
Everything else registered in `CometSparkSessionExtensions.apply` is already
inert once no Comet nodes exist: `CometPlanAdaptiveDynamicPruningFilters`
only matches Comet scans, `CometReuseSubquery` creates no Comet nodes, and
the two post-columnar rules only act on Comet operators.
* refactor: compose CometScanRule and CometExecRule into one CometRule
The two were registered as separate rules, adjacently, in both the columnar
and the query-stage-prep paths. Nothing ever ran between them, and neither is
useful alone: CometExecRule seeds its native chain only from the nodes
CometScanRule produces, so operator conversion over unconverted Spark scans
converts nothing.
Compose them so the ordering is an invariant of the code rather than of the
registration order, and so callers that need the whole conversion have one
entry point. Both rules keep their classes, files and tests.
* test: pin the invariant that motivates composing the rules
* refactor: simplify plan-only mode and trim its comments
Share the post-columnar rule list between CometColumnar and the plan-only
preview through CometRule.postColumnarRules, replacing the hand-synced copy.
RevertNativeForTransitionHeavyStages takes a wholePlan flag in place of the
separate applyToAllStages entry point.
Remove the unused imports left in CometExecRule, which failed the scalafix
lint, and restore the CometRule docstring and the "scan conversion must run
before operator conversion" test from #6082 that the main merge dropped.
Condense the comments, the user guide section and the config doc, inline
single-use helpers, and factor the plan-only tests onto shared helpers.
Check the Iceberg split-write conf before the plan-only conf so the common
path does one lookup.
* fix: scope plan-only report dedupe to the query and detect AQE re-plans
Render plan-only reports in the verbose format regardless of
spark.comet.explain.format, so the coverage summary is never dropped.
Tell AQE re-planning apart from initial preparation by whether the rule is
running inside AdaptiveSparkPlanExec.reOptimize, instead of guessing from
maxRows. An initially empty query is now reported under AQE, and a re-plan
that keeps operators above an empty join is no longer reported twice.
Deduplicate repeated subqueries within the query being prepared on the
current thread rather than within a SQL execution ID, so queries planned
without one (toRdd, df.rdd, executedPlan) no longer repeat reports.
---
.../user-guide/latest/understanding-comet-plans.md | 25 ++
.../main/scala/org/apache/comet/CometConf.scala | 11 +
.../apache/comet/CometSparkSessionExtensions.scala | 9 +-
.../org/apache/comet/ExtendedExplainInfo.scala | 18 +-
.../comet/iceberg/IcebergWriteStrategy.scala | 4 +
.../scala/org/apache/comet/rules/CometRule.scala | 190 ++++++++++-
.../RevertNativeForTransitionHeavyStages.scala | 8 +-
.../comet/CometIcebergWriteDetectionSuite.scala | 20 ++
.../apache/comet/rules/CometExecRuleSuite.scala | 366 ++++++++++++++++++++-
9 files changed, 635 insertions(+), 16 deletions(-)
diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md
index 32fe297e67..6d2b361510 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 a83bef6b45..f0187c8cf0 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -884,6 +884,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 f274248a2b..7da6ed5813 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 37535f8050..8d20b33868 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 8ceb4e6e48..46c3dbca2e 100644
--- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala
+++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala
@@ -42,6 +42,10 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy {
if (!isCometLoaded(conf) || !CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.get(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 9b9639ac18..fa2fb7dc32 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 a56e477b03..bbdc6db6e5 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 37ddb1872e..0d7d9935b2 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 d59a9d3207..1769663dbc 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -32,11 +32,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, DoubleType, FloatType, 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, QueryPlanSerde, Unsupported}
import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator}
@@ -1272,6 +1273,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)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org