feat: add spark.comet.explain.planOnly.enabled

[andygrove]

Which issue does this PR close?

Closes #5335. Alternate approach to #5345, based on review discussion there.

Heads up: I used an LLM to help draft this. The design is mine, but the code and prose have been shaped with LLM assistance, so review with that in mind.

Rationale for this change

Users evaluating Comet on a workload need a way to estimate how much of it Comet would accelerate without actually changing execution. Turning Comet on and comparing runs carries real risk. This adds a mode that builds the Comet plan Comet would have executed, logs it, and lets Spark run the query unchanged.

What changes are included in this PR?

• New config spark.comet.explain.planOnly.enabled, default off.
• When set, CometScanRule and CometExecRule short-circuit at the top of their apply and return the plan untouched.
• CometExecRule builds a preview by calling both rules' private _apply methods directly (with an explicit forPreview = true parameter that skips the short-circuit on the recursive call), logs the resulting Comet plan at WARN level, then returns the original Spark plan.
• The plan-only branch sits at the top of the exec-enabled block so normalizePlan/RewriteJoin/tagUnsafePartialAggregates are not run and thrown away.
• The WARN is deduped per SQL execution ID (bounded LRU) so under AQE a query gets one report, not one per stage.
• Doc section in understanding-comet-plans.md.

The estimate is Scala-side only. The native plan is never handed to DataFusion, so anything that would have failed inside DataFusion still counts as accelerated. That is called out in the config docstring and the user guide.

How are these changes tested?

New tests in CometExecRuleSuite cover:

• V1 and V2 Parquet scans, each with AQE on and off.
• Scalar subquery, so the recursive preview builds through the subquery plan without leaving any Comet operators behind in the executed plan.
• Same query with the config off, asserting Comet operators do appear (sanity check that we haven't accidentally disabled Comet).

Each test asserts the executed plan has zero CometPlan operators.

[sunchao]

Summary

Reviewed eb514de1ff76785b81a4940d6a2adb3ea149ab07 against a28ac348f5a68500d96fe872b49e453d96e4bab6 across five independent scopes. Reusing the existing conversion rules is a sensible way to keep eligibility checks aligned with normal Comet planning. I found two P2 issues that affect the accuracy and completeness of the report. Both are described inline.

Prior state and problem

The existing explain output describes a plan that has already been converted for Comet execution. Evaluating an unfamiliar workload therefore requires enabling Comet on the execution path. This change adds an opt-in way to inspect potential acceleration while retaining Spark execution.

Design approach

The new configuration defaults to false. The public scan rule leaves the incoming plan alone, while the execution rule builds a disposable preview by invoking scan conversion and execution conversion directly. The explicit forPreview argument prevents that preview from taking the plan-only shortcut again.

Correctness / compatibility analysis

I did not find a verified native-execution leak in the ordinary read path. The main concerns are reporting correctness: nested subquery planning can claim the execution-ID report slot before the outer query, and the preview stops before Scala-side transition insertion and transition-driven reversion.

A local Spark 3.5.2 planning-order probe using the same deduplication logic reproduced the first issue with AQE both off and on. In both cases the scalar aggregate was reported first and the outer filter was suppressed. The second issue is supported by the exact-head rule ordering and the existing RevertNativeForTransitionHeavyStagesSuite regression case. I did not run the full Comet JVM/native suite. git diff --check passed. At the final CI check, 12 checks had succeeded, 5 were running, 1 was queued, 7 were skipped, and none had failed.

Key design decisions

Keeping the feature disabled by default limits ordinary behavior changes. Returning the original plan avoids reconstructing a Spark plan from converted operators. A bounded execution-ID cache also limits retained reporting state, but the report owner must distinguish the root query from recursively prepared subqueries.

Implementation sketch

The implementation adds the configuration and user guide, exposes scan conversion within the rules package, adds the preview/report branch to CometExecRule, and tests V1/V2 scans with AQE on and off plus a scalar subquery and a config-off sanity check. Those tests check plan isolation, but they do not inspect the warning emitted by an actual query action.

Behavioral changes worth calling out

Plan-only mode requires Comet execution to be enabled, emits a driver warning, and deliberately does not call DataFusion's native planner. That documented native-planning limitation is reasonable. It should remain distinct from omitting a known Scala-side reversion or reporting only an inner subquery.

Suggested improvements

Please make the root query own the once-per-execution report, and calculate coverage after the applicable columnar-transition and Comet post-columnar rules. Add action-based log assertions for scalar subqueries under both AQE modes and compare the preview with normal execution when transition reversion is enabled.

[sunchao]

[P2] Keep the report slot for the outer query

Could we avoid consuming the execution-ID entry while Spark is preparing a nested subquery? For a fresh action on SELECT id FROM range(10) WHERE id > (SELECT max(id) FROM range(3)), Spark prepares the scalar subquery before the outer plan under the same execution ID. The subquery records the ID here, so the later outer-query call is suppressed. I reproduced this planning order with the same deduplication logic on Spark 3.5.2 with AQE both disabled and enabled. The only report describes max(id), not the outer scan/filter or the workload being evaluated. Please distinguish root-query reporting from subquery/stage preparation and add an action-based test that captures the warning and checks that the outer plan appears.

[sunchao]

[P2] Include post-columnar reversion in the coverage preview

Could the preview run transition insertion and Comet's post-columnar rules before computing coverage? Normal planning subsequently runs RevertNativeForTransitionHeavyStages and EliminateRedundantTransitions, but this path stops before them. With spark.comet.exec.transitionRevert.enabled=true, spark.comet.exec.transitionRevert.maxTransitions=0, and Comet project execution disabled, the existing regression case shows that the real executed plan has zero CometExec nodes. This preview still counts the operators that the configured Scala-side rule removes, and its transition count is calculated before Spark inserts those transitions. That is separate from the documented uncertainty about DataFusion planning failures. Please compare the report with the real post-columnar plan for this configuration.

[andygrove]

Both fixed in 2f94707. On the first one — I couldn't find a way to identify the root plan at rule-application time. Subqueries are prepared before the outer plan with nothing to distinguish them, and matching the incoming plan against the root QueryExecution's output breaks down for commands and write plans, which would then get no report at all. So instead of one report owned by the root, each independently planned plan gets one: the outer query plus one per separately prepared subquery. Dedupe is now keyed on execution ID and plan, and AQE's per-stage and post-re-optimization applications are skipped as re-plans of something already reported, so the stage spam the execution-ID cache was there to prevent is still gone. Does that seem like a reasonable trade to you, or would you rather see the root identified even if some plan shapes drop out of reporting?

The second one was as you described. The preview now inserts transitions and runs both post-columnar rules. RevertNativeForTransitionHeavyStages needed a new entry point for this: its apply takes the AQE branch when AQE is on, which only judges the topmost stage, whereas the preview holds a plan that hasn't been split into stages, so it needs every shuffle boundary visited. The new test builds your regression configuration, reads the coverage out of the captured report, and compares it against CometCoverageStats for the plan Comet really executes — before the fix the report claimed 4/4 with no transitions against a real 3/4 with one.

What the preview still can't match is AQE re-planning: it describes the pre-adaptive plan and applies the post-columnar rules to it in one pass rather than per stage. That's called out in the user guide alongside the native-planning caveat.

[sunchao]

[P1] Preserve the captured report across Spark 3.x withSQLConf

Could the withSQLConf block be moved inside capturePlanOnlyReports, or could these assertions run inside the configuration block? Spark 3.4 and 3.5 define withSQLConf(...)(f: => Unit): Unit, unlike the generic Spark 4.x helper, so reports is inferred as Unit here. The exact-head Spark 3.4 and Spark 3.5 checks both fail test compilation on the subsequent size, mkString, and head calls. This prevents the supported Spark 3.x test builds from compiling.

[sunchao]

[P2] Deduplicate adaptive reports when the execution ID is absent

Could the no-ID path retain plan-scoped reporting state instead of returning true for every invocation? The public df.rdd.count() path can build and execute AQE stages without installing spark.sql.execution.id. A Spark 3.5.2 probe using this exact decision logic and both rule registrations produced five report decisions for SELECT id % 2 AS k, count(*) AS n FROM range(20) GROUP BY id % 2: initial preparation, the adaptive wrapper, the exchange stage, adaptive re-optimization, and the final stage. A fresh collect() produced one. Because this branch bypasses both the stage check and deduplication, plan-only mode rebuilds previews and emits overlapping coverage summaries for those ordinary RDD-backed workloads, contrary to the documented suppression of stage/re-optimization reports. Please cover df.rdd.count() and planning via executedPlan before an action in the reporting tests.

[sunchao]

[P2] Include converted scalar subqueries in the outer coverage report

Even with AQE disabled, an acceleratable scalar subquery is counted as Spark in the outer report. Its independently converted preview has already been discarded, and these two conversion passes walk ordinary plan children, leaving the original ScalarSubquery.plan in the outer preview. ExtendedExplainInfo then traverses those expression-owned plans and includes their operators in the percentage. For the new scalar-subquery test query, the separate subquery warning can therefore report acceleration that is missing from the outer query's coverage. A constructed-plan probe using the exact-head formatter/serializer and real Comet project nodes reports 1/5 with the untouched subquery versus 2/5 after replacing only its plan. Please carry the converted subquery previews into the outer preview, or exclude separately reported subqueries from that report's counts, and compare its coverage with normal Comet planning.

[coderfender]

Seems like there are CI failures :)

[coderfender]

LGTM pending CI

[andygrove]

All three fixed in 1683a65.

The Spark 3.x compile break was exactly as you described — the assertions now run inside the withSQLConf block rather than binding its result. -Pspark-3.4 and -Pspark-3.5 both test-compile clean here, and CometExecRuleSuite passes on 3.5 as well as on the default 4.1.

On the no-execution-ID path: rather than scoping state to an ID that may not exist, the report is now marked on the plan itself. Catalyst copies tags onto replacement nodes, so the mark survives the rewrites Spark makes between one application of the rule and the next, which is what lets a later application recognise a stage or final-plan pass as a re-plan of something already reported. That turned out to matter beyond the no-ID case: I found a plan with no exchanges gets a final-plan pass under AQE whose hash no longer matches the prep-time plan, because the subquery was substituted in between, so the old dedupe let it through as a second report. I kept the execution-scoped hash alongside the mark for the one case a mark can't cover — a subquery referenced from two places is prepared twice as separate but identical plans. Both df.rdd.count() and executedPlan before an action are covered by tests now, under AQE on and off.

The subquery one was worse than the double-count you found, and getting the outer report to agree with normal planning took two more things. A subquery reaches the outer plan fully prepared, so its aggregates are already inside WholeStageCodegenExec with transitions inserted, and handing that to the conversion pass leaves them unconverted — the preview was reporting a subquery as falling back that Comet actually accelerates. Those artifacts are stripped before previewing. And ReuseExchangeAndSubquery is the last thing QueryExecution.preparations does, after the columnar rules, so the plan the preview sees still has one copy of the subquery per reference; replaying that rule at the top level collapses them the way the executed plan does. With AQE off the outer report is now exactly CometCoverageStats for the plan Comet really executes — 6/7 with 2 transitions on your scalar-subquery query, against 6/12 with 5 transitions before — and there's a test asserting that equality.

One thing I could not make match, and I'd like your read on it. Under AQE the outer report is produced from the prep-rule application, and at that point PlanAdaptiveSubqueries hasn't run, so the subquery is still a logical plan hanging off the expression. Extended explain walks it anyway and counts those logical nodes as un-accelerated Spark operators, so the same query reports 2/7 under AQE against 6/7 off. Reporting operators we never evaluated as fallbacks is the part that bothers me. The alternatives I can see are to exclude un-planned subqueries from the counts, which trades a pessimistic number for an optimistic one, or to report from the post-adaptive application instead of the prep one, which is only reachable for queries AQE never cuts into stages. For now it's documented, with a pointer to the per-subquery reports. Would you rather I pursue one of those, or is documenting it acceptable for a first cut?

The remaining red check is spark-sql-sql_hive-1, which failed downloading apache-maven-3.9.6-bin.zip from Maven Central — unrelated to this PR.

[sunchao]

Re-reviewed 1683a65 with five agents. The 41-test CometExecRuleSuite and 96 normal/preview plan comparisons passed on Spark 4.1.3. Targeted probes against locally compiled JVM and native code reproduced the two reporting issues below; the third comment identifies a test coverage gap.

[sunchao]

[P2] Apply stage reversion to the prepared DPP child

Could the DPP broadcast child be previewed through its own preparation boundary before restoring the exchange wrapper? With AQE disabled, Spark's PlanDynamicPruningFilters prepares the child before constructing BroadcastExchangeExec. This path instead passes that exchange into buildPreview, so applyToAllStages treats it as a boundary and leaves the reconverted child unreverted. I reproduced this on the PR head with a partitioned Parquet fact table joined to a filtered Parquet dimension, spark.comet.exec.transitionRevert.enabled=true, spark.comet.exec.transitionRevert.maxTransitions=0, and spark.comet.exec.project.enabled=false: the outer report says 4/12 accelerated (33%), while CometCoverageStats for the normally executed plan says 2/12 (16%). Both executions returned the same ten correct rows. The separate DPP report correctly shows 0/3, but its child is counted as accelerated again in the outer report. Please preserve the child's normal preparation scope and add a non-AQE DPP coverage comparison.

[sunchao]

[P2] Keep report state when AQE replaces an empty plan

Could report ownership survive replacement of the entire physical tree? With plan-only mode and AQE enabled, spark.sql("SELECT id % 2 AS k, count(*) AS n FROM range(20) WHERE id < 0 GROUP BY id % 2").collect() emits an 83% report for the initial aggregate and then a second 0% report for EmptyRelation. After the empty shuffle materializes, AQE creates a fresh plan containing neither QueryStageExec nor the reporting tag. Its structural hash also differs, so both guards admit another report under the same execution ID. I reproduced the two actual Comet warnings on Spark 4.1.3; AQE disabled emits one. Please retain query/subquery reporting ownership across this rewrite and cover an empty adaptive query.

[sunchao]

[P3] Select V2 before creating the Parquet fixture

Could USE_V1_SOURCE_LIST be configured outside withParquetTable? The fixture reads Parquet and registers tbl before runPlanOnlyAndAssertReverted changes that setting. Changing it afterward does not replace the existing logical relation, so both tests labeled "V2 scan" still exercise FileSourceScanExec. A Spark 4.1.3 probe confirmed this with AQE on and off; creating the fixture after the config change produces BatchScanExec. Actual V2 plan-only execution passed separate checks, but these committed tests do not cover that path.

[andygrove]

All three fixed in c75a87c.

The DPP one was as you described. PlanDynamicPruningFilters prepares the build plan and wraps it in the exchange afterwards, so the preparation boundary — and the stage RevertNativeForTransitionHeavyStages judged — is the exchange's child. The preview now descends through the wrapper and puts it back, and on your configuration the outer report reads 2/12 with 3 transitions, matching CometCoverageStats for the executed plan exactly. There's a test that builds a partitioned fact table joined to a filtered dimension with reversion forced on and asserts that equality.

The empty adaptive plan I reproduced too, and it turned out the hash was not the only guard that missed: the mark could not reach the new tree either, since AQEPropagateEmptyRelation keeps the plan it replaced in EmptyRelationExec's innerChildren rather than its children. Instead of chasing it through there I went with a third guard: once AQE has cut a query stage for an execution, everything that arrives afterwards is AQE re-planning something already reported, so the rest of the execution is suppressed. That is ordering-safe for subqueries because they are all compiled before the first stage is cut. The exception is a subquery AQE only plans once stages are under way — a DPP subquery, since PlanAdaptiveDynamicPruningFilters is a stage optimizer rule — which now gets no report of its own under AQE. That seemed a better trade than an extra bogus 0% report, but say the word if you'd rather keep it. Noted in the user guide, and tested under AQE on and off.

You were right about the V2 tests, and the cause is that withParquetTable resolves the relation through spark.read and registers the result as a temp view, so the source list is baked in at creation. The config now wraps the fixture, and each variant asserts it actually got a FileSourceScanExec or a BatchScanExec so it cannot silently stop covering the path again.

Separately, while building the DPP fixture I hit an unrelated pre-existing bug and filed #5486: with AQE on, DPP in play, and spark.comet.exec.transitionRevert.enabled=true, the query throws SubqueryAdaptiveBroadcastExec does not support the execute() code path. It reproduces on this PR's base commit with plan-only mode never enabled, and only when transitionRevert is on — so it's out of scope here, but it matters for #5207. It's why the new DPP test is non-AQE only, beyond the report being pre-adaptive anyway.

CometExecRuleSuite passes on Spark 3.4 (42, 2 pre-existing version-guard cancellations), 3.5 (43, 1) and 4.1 (44), and RevertNativeForTransitionHeavyStagesSuite, CometScanRuleSuite, CometCoverageStatsSuite and CometDppFallbackRepro3949Suite are green.

[sunchao]

Re-reviewed c75a87c with five agents. The 60 existing tests across five suites and 132 non-AQE coverage comparisons passed on Spark 4.1.3. Targeted probes reproduced the two reporting issues below.

[sunchao]

[P2] Suppress empty AQE reports without an execution ID

The new guard only retains state when an execution ID exists. With AQE and spark.comet.explain.planOnly.enabled=true, an empty grouped query through the JVM bridge used by PySpark's df.rdd still emits the initial 83% report followed by a misleading 0% EmptyRelation report:

SELECT id % 2 AS k, count(*) AS n
FROM range(20)
WHERE id < 0
GROUP BY id % 2

I reproduced this on the current head with Spark 4.1.3 by invoking javaToPython() and counting the resulting RDD; AQE disabled emits one report. The action has no SQL execution ID, and the replacement tree has no reporting tag. Could suppression survive this AQE replacement without requiring an execution ID, with a regression covering this RDD path?

[sunchao]

[P3] Canonicalize nested subqueries before deduplicating reports

Projecting the same nested scalar subquery twice and calling .collect() with AQE disabled produces four reports under one execution ID: the inner aggregate, the outer aggregate twice, and the main query. For a Parquet table tbl with integer columns _1 and _2, the reproducer is:

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

On Spark 4.1.3 at the current head, the two outer-aggregate reports are byte-for-byte identical, and both normal and plan-only physical plans reuse that subquery through ReusedSubqueryExec. The raw plan hash misses this equivalence, causing redundant previews and warnings. Could the key use a canonical fingerprint, with a nested-subquery reuse regression? This is a reporting issue; results and coverage values still match.