feat: add spark.comet.explain.planOnly.enabled, reported per query

[andygrove]

Which issue does this PR close?

Closes #5335. Alternate approach to #5394 (which was itself an alternate to #5345).

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.

#5394 does this by reporting from inside CometExecRule. That works, but Spark applies a planner rule many times for one query — once per query stage and once per adaptive re-optimization under AQE, plus once for every subquery it prepares separately — so most of that PR is machinery deciding which application owns the report: a bounded LRU of execution-id/plan-hash keys, a tag copied along Catalyst rewrites, a query-stage check, an empty-re-plan check, and a queryStagePrep flag threaded in from the extension. Five review findings on that PR are variations of "the wrong application claimed the report", and each needed another guard.

This PR reports from a QueryExecutionListener instead. One callback per action, holding the finished plan, so there is nothing to tell apart and none of that machinery exists.

What changes are included in this PR?

• New config spark.comet.explain.planOnly.enabled, default off.
• CometScanRule and CometExecRule return the plan untouched at the top of apply while the mode is on, so Spark plans and executes the query exactly as it would with Comet off. CometExecRule also registers the session's listener, so a session carries one only if the mode is used, and the config stays togglable mid-session.
• CometPlanOnly builds the report from qe.executedPlan: it undoes the preparation that follows the conversion rules (adaptive wrappers, query stages, AQEShuffleReadExec, codegen wrappers, columnar transitions, exchange reuse), previews the plans behind subquery expressions, converts, then replays transition insertion and Comet's post-columnar rules — so a stage Comet would have handed back to Spark is reported as handed back.
• RevertNativeForTransitionHeavyStages gains applyToAllStages, because the preview holds a whole plan where AQE would have handed that rule one stage at a time.
• Metadata-only statements (CREATE VIEW, SHOW TABLES) are not reported; a session runs enough of them that one 0% report each would bury the rest.
• Nothing in the reporting path can fail the query: it runs off the query thread and swallows non-fatal failures with a warning.

Measured in non-comment code lines the two approaches are the same size (123 vs 122 in main sources); the state machine's cost was mostly the prose needed to explain it. What differs is where the remaining complexity sits. Here it is one function that normalizes a plan, and when it gets a shape wrong the symptom is a coverage number that disagrees with the real plan — which a test catches mechanically. In #5394 a wrong guard shows up as a missing or duplicated report, which nothing catches until someone reads the log.

How are these changes tested?

New CometPlanOnlySuite, 20 tests. Two properties:

• The query runs on Spark. V1 and V2 scans, AQE on and off, asserting the executed plan holds no CometPlan operator and that the intended scan path was actually exercised.
• The report describes what Comet would really have executed. The same query is run with Comet enabled and the report's coverage compared against CometCoverageStats for the plan that ran, for an aggregate, a shuffled join, a scalar subquery, a DPP subquery and a stage forced through transition reversion. AQE on and off, except DPP (non-AQE only, because AQE + DPP + spark.comet.exec.transitionRevert.enabled fails with "SubqueryAdaptiveBroadcastExec does not support the execute() code path" #5486 still breaks that combination).

Plus: one report per action for a multi-stage query with a subquery, two actions reported twice, an RDD action reported once, an adaptive query that collapses to an empty relation reported once, metadata-only statements not reported, and the config off leaving Comet in charge.

Two behaviours found while writing those tests are documented rather than fixed:

• Under AQE the report describes the plan AQE settled on, which is the plan Comet would have been asked to run. If a stage materialized empty and AQE replaced the query with an empty relation, that is what gets reported.
• Under AQE the transition count can be one lower than a real Comet run's, because Spark inserts transitions one stage at a time while the report is a single pass over the flattened plan. The operator counts behind the percentage are unaffected, and the tests compare transitions with AQE off.

The estimate remains Scala-side only — the plan is never handed to DataFusion, so a DataFusion planning failure still counts as accelerated. That is called out in the config docstring and the user guide.

CometPlanOnlySuite, CometExecRuleSuite, CometScanRuleSuite, RevertNativeForTransitionHeavyStagesSuite and CometCoverageStatsSuite are green (64 tests) on the default profile, test-compile is clean on spark-3.4, spark-3.5, spark-4.0 and spark-4.2, and scalafix passes with the semantic rules.

[sunchao]

Reviewed the full change at fa77513. Three reporting issues remain below. Validation included the existing merge CI logs and focused Spark 4.0.4 component checks with unchanged reporter code; the local checks did not run Comet conversion or JNI.

[sunchao]

[P2] Cover RDD actions outside the named SQL callback path

Registering only this listener leaves spark.sql(...).rdd.count() without a report on Spark 3.4/3.5: the new RDD test captures zero reports and is the sole failure in both the 3.4 job and 3.5 job. Those RDD actions do not emit the named SQL completion event this listener requires. The 4.x single-count pass is not a per-action control either: in the Spark 4.0.4 component check, obtaining df.rdd emits one report, while two subsequent count() actions add none. Please cover the RDD execution lifecycle before promising one report per action.

[andygrove]

Fixed in 8ccff51, by correcting the claim rather than by covering RDD actions — and I reproduced your CI result locally first, on both profiles:

=== 4.1 ===  - an RDD action is reported once ... succeeded 1
=== 3.5 ===  - an RDD action is reported once *** FAILED ***
             ArrayBuffer() had size 0 instead of expected size 1

So the split is exactly as you described. Reporting is driven by QueryExecutionListener, which Spark fires from the Dataset action path, so it is per SQL execution, not per action. On 4.0+ obtaining df.rdd runs a query of its own and is reported at that point; on 3.4/3.5 it is not reported at all; and on both, the RDD actions that follow add nothing because no new SQL execution starts.

I did not add a second reporting path. Doing it properly means keying on job starts and then reconciling against this listener so an ordinary query is not reported twice, which is a good deal more machinery than this mode justifies — so the honest move was to stop promising something it does not do. The user guide now says "once per SQL execution" with a paragraph on what that means for RDD work, and the class doc says the same and records why the RDD path is not covered.

The test is now RDD actions are outside the reported path. It takes three actions on the same RDD and asserts the version-appropriate outcome (one report on 4.0+, none on 3.x), which pins the split rather than leaving it to be rediscovered, and asserts the part that is version-independent: the RDD's own actions contribute nothing. CometPlanOnlySuite is 22/22 on both the default 4.1 profile and -Pspark-3.5.

[sunchao]

[P2] Retain the action's settings for delayed reporting

A caller can enable plan-only, run collect(), and restore the setting after the action returns while this asynchronous callback is still queued. This reads the session's later flag and silently drops the report for a query that ran in plan-only mode. Other Comet settings changed in that gap also affect the preview. A gated Spark 4.0.4 listener control produces one report when the flag stays enabled and zero when it is restored before callback delivery. The new tests drain the bus inside withSQLConf, so they hide this ordering. Keep the action's eligibility and planning settings with its report instead of re-reading mutable session state.

[andygrove]

Fixed in 8ccff51. The settings now travel with the query instead of being read back off the session.

CometExecRule tags the plan it declines to convert with a PlanOnlySettings snapshot — plan-only enabled, Comet loaded, exec enabled — taken synchronously while the query is being planned, so it is inside whatever withSQLConf the caller wrapped the action in. report reads the tag off the executed plan, searching the adaptive wrappers' own plans too, since under AQE the tagged plan hangs off AdaptiveSparkPlanExec rather than appearing among its children. It falls back to the session only when no tag is found, which covers a plan that reached the listener without passing through the rule.

The test was the hard part, and you had already identified why. My first attempt did what the existing tests do — leave the withSQLConf block, then drain — and it passed against the unfixed code, because the bus normally drains during collect() and the callback sees the flag still on. A test that passes either way is worse than none, so I did not keep it.

The version that works forces the ordering. It runs on spark.newSession(), which gives a deterministic registration order: CometPlanOnly registers its listener lazily on first use, so on a session that has never run a plan-only query a gating QueryExecutionListener registered first also runs first, and can park the bus on a latch. The setting is restored while the callback sits in that latch, then the latch opens.

Against the previous code it fails as you predicted:

List() had size 0 instead of expected size 1
expected the query planned under plan-only mode to still be reported

[sunchao]

[P2] Preserve reused exchange output IDs when expanding it

In a Parquet self-join with exchange reuse enabled, the wrapper can expose fresh IDs such as k#16 while its shared child produces k#3. Returning that child leaves the parent sort referring to the discarded ID; Comet's attribute binder then declines the sort and its consuming join stays on Spark, so the preview understates coverage for otherwise supported work. The actual Spark plan binds before normalization; invoking this unchanged method loses the sort/join bindings with AQE both off and on, while disabling reuse preserves them. Retain or remap the wrapper's output IDs while exposing its subtree for conversion.

[andygrove]

Fixed in 8ccff51. stripPreparation still undoes the reuse — that part was deliberate and the comment explains why — but it now carries the wrapper's output IDs across:

case reused: ReusedExchangeExec =>
  restoreReusedOutput(reused, stripPreparation(reused.child))

restoreReusedOutput is a no-op when the IDs already agree, which is the common case; when they do not, it re-aliases positionally — the same correspondence ReusedExchangeExec itself relies on — using a ProjectExec of Aliases carrying the wrapper's exprIds. That is one extra operator in the report per re-aliased reuse, which I judged the better trade: a projection of aliases is something Comet converts, so it costs a point of denominator, whereas losing the sort and the join above it was costing the preview real coverage.

The regression test is a reused exchange keeps its output IDs so consumers still convert: a self-join over two identical grouped subqueries with EXCHANGE_REUSE_ENABLED=true and broadcast disabled, asserting the join converts in the preview and that no Sort is left on Spark. Against the previous code it fails, so it reproduces the shape you described rather than just documenting it.

I used AQE off for the test since that is the simpler of the two configurations you mentioned; the fix is in stripPreparation, which both paths go through.

[sunchao]

Correctness

Rechecked dd54fbd2 against the assigned base 7f1e0018 and prior reviewed head fa775130. The added and removed lines in all nine authored files are unchanged after accounting for their respective merge-bases. CometPlanOnly and CometPlanOnlySuite are also byte-for-byte unchanged. The three existing P2 findings remain applicable. I am keeping the feedback in those threads.

[P2] Report RDD action completion

The RDD-action finding is reproduced by current CI. The Spark 3.4 execution job passed 775 tests and failed only the RDD reporting test. The Spark 3.5 execution job passed 788 and failed the same test. Both captured zero reports at suite line 208. The inspected Spark 4.0 execution job passed all 835 executed tests, including that single-count case. Maintained Spark 3.5 source still constructs Dataset.rdd without a named SQL action. Maintained 4.0 wraps the lazy RDD accessor, not each later RDD action, so that passing control does not establish per-action reporting. Executor-side SQLExecutionRDD configuration propagation does not supply a driver listener callback.

[P2] Preserve the action's settings

The delayed-settings finding still reads the live session configuration in the asynchronous callback. No query-local eligibility or configuration snapshot was added.

[P2] Preserve reused-exchange output IDs

The reused-exchange finding still discards the wrapper's output IDs when expanding its child. In maintained Spark 3.5 and 4.0, those IDs are explicitly preserved by the wrapper, and attribute binding still resolves by exprId. The current Comet binder continues to fall back when the reference cannot be resolved. The merge therefore does not address either issue.

The three inspected jobs executed merge 2f5e16e4, whose parents are exactly the assigned base and head and whose whole tree equals the head. The current snapshot contains 62 successful, 9 skipped and 2 failed checks. This is fresh CI evidence, separate from the historical component probes, which I did not rerun. No local JVM/native test or benchmark was run. Maintained Spark 3.4/4.1 source remains unavailable, so the 3.4 CI failure does not fill that source gap. There is no update-specific expression change involving types, nulls, overflow or ANSI behavior. Ordinary execution and AQE/subquery report handling retain the same source-level limitations as before.

Performance

The update introduces no new reporting work or native hot-path change. The preview still runs conversion and rendering on the driver listener queue. Its coverage percentage remains a Scala-side estimate, not measured acceleration or proof that DataFusion would accept the plan. No new performance claim or changed implementation warrants a separate microbenchmark in this follow-up. No speedup is inferred from the passing tests.

Design

The current-base integration retains the early return in both conversion rules and the preview's explicit conversion, transition insertion and stage-reversion sequence. It does not change the listener's lifetime or event boundary. The existing fixes still need to align reporting with actual action completion, preserve the action's settings and retain reused-exchange attribute identity. Resolving those points remains necessary before relying on the per-action coverage report.

Abstraction & complexity

No new plan-only abstraction or state was added since the previous review. The existing normalization and subquery helpers remain unchanged, and applyToAllStages still delegates to the existing non-AQE stage traversal. I found no additional P1/P2 finding in this update that warrants another inline thread.

[sunchao]

Correctness

Rechecked 8ccff51abc7afdd2d4d17aac3677c96e3ddeb3d2 against 7f1e00189b1ed86f1cb5acd872d97fce694482b1 and the four-file increment from dd54fbd2. The base is unchanged. The update improves all three prior findings, but two existing P2 findings remain partially addressed. There are no new inline findings.

The RDD reporting concern is addressed by narrowing the documented contract to SQL executions. That matches maintained Spark 3.5/4.0: the 3.5 RDD accessor does not emit the named listener event, while 4.0 wraps the lazy accessor rather than each later RDD action. The revised test checks three actions on the same RDD and the version-dependent report count. This resolves the scope mismatch without adding another event path.

[P2] Use the captured settings throughout conversion

The delayed-settings finding is only partly fixed. The new tag keeps the three eligibility flags, and the gated test now verifies that restoring the plan-only flag does not drop the report. However, report passes the original session to previewOf without restoring a query-local configuration. Both conversion rules call isCometLoaded(conf) again. If spark.comet.enabled changes from true to false while the callback is held, the saved flags allow reporting but both rules return the Spark plan, producing an unaccelerated preview for the earlier enabled configuration. Operator flags and transition-reversion settings can likewise change its coverage.

Maintained Spark's listener calls the callback without a saved SQLConf. Also, SparkPlan.conf reads its captured session's current configuration, and isOperatorEnabled explicitly uses op.conf. Saving a few flags, or only changing the listener's active session, therefore does not preserve every preview read. Could the snapshot cover the conversion and rendering configuration, with both rule-level reads and plan-bound reads using it? Extend the gated test to restore a conversion setting and compare the reported coverage with the unchanged-settings control, rather than asserting only that one log line exists. This remaining case is verified from the current caller chain, not a new local runtime experiment.

[P2] Keep exchange identity restoration independent of projection support

The reused-exchange finding is fixed for the default projection-enabled case. Positional aliases preserve the wrapper's IDs, which is consistent with maintained Spark's reuse mapping. The same normalization reaches reused query stages under AQE.

The new ProjectExec introduces a remaining configuration-dependent failure: with spark.comet.exec.project.enabled=false, CometProjectExec.enabledConfig rejects this synthetic projection. tryConvertToComet then supplies no native child to the parent sort, whose serde requires one, and the consuming join also falls back. The intermediate projection cannot use the leaf-only Spark-to-Arrow rescue. A real conversion of those exchange branches has no such identity-restoration projection. Thus the preview can still lose a supported sort/join solely because it added a node for bookkeeping. Could identity be restored without requiring the user's projection setting to be enabled? A comparison against the real plan with projection disabled would cover this remaining case. The current reuse test leaves projection enabled and only checks AQE off.

Validation and limits

The five Linux execution jobs each pass all 22 CometPlanOnlySuite tests, including the three updated controls. These jobs executed merge 8a831daa, with parents 424c31aa and this head. That newer base differs from the assigned base, and its whole tree differs from head in 121 files. The reporter, suite, scan rule, reversion rule and guide match exactly. The relevant conversion/configuration regions were also compared. This supports those tests without claiming whole-tree head validation or independently reconciled native artifact provenance.

At 2026-09-09T19:02:23Z, checks show 54 successful, nine skipped, five failed and one queued. The syntactic lint job and four semantic lint jobs all request removing the unused s prefix at CometPlanOnlySuite.scala:272. CI is not green. No local product test or benchmark was run. Canonical semantics were checked on maintained Spark 3.5/4.0. Maintained 3.4/4.1 source gaps remain, and CI execution on other versions does not replace that source review.

Performance

The update adds a small settings tag and a plan search on the listener thread. Re-aliasing adds a preview-only projection when IDs differ. Neither executes in the query's native data path. The existing conversion and rendering work still runs on the listener queue. No measured throughput improvement is claimed, and the passing correctness tests are not a benchmark.

Design

The narrower RDD contract and deterministic listener gate are useful changes. A query-owned snapshot fits the asynchronous reporting boundary, but it must govern the whole preview. Identity restoration also needs to preserve conversion eligibility under the user's operator settings. Those are refinements to the existing two findings, with no need for another reporting path.

Abstraction & complexity

The settings record and focused output-remapping helper are understandable. The alias projection is not transparent to the planner: it adds an eligible operator to the count and participates in the projection support gate. Keep identity bookkeeping from altering the work being measured. No separate abstraction or complexity P1/P2 is proposed.

[andygrove]

Both remaining P2s are addressed in ac8e5a1, and the lint failure at CometPlanOnlySuite.scala:272 is gone (an s prefix on a string with no substitution).

Use the captured settings throughout conversion

You are right that the three-flag snapshot only covered the gate, and I went looking for a way to make it govern the whole preview before concluding it cannot be done. SparkPlan.conf is:

override def conf: SQLConf = {
  if (session != null) { session.sessionState.conf } else { super.conf }
}

with @transient final val session = SparkSession.getActiveSession.orNull captured at construction. So op.conf is the live conf of whichever session each node was built under. SQLConf.withExistingConf cannot redirect it, because it never consults the thread-local. A cloned session with the snapshot as its conf cannot either, because the nodes already exist and withNewChildren returns this when the children are unchanged, so most of the tree is never reconstructed. And CometExecRule reads op.conf in six places, including the per-operator gate at handler.enabledConfig.forall(_.get(op.conf)), the strict-fallback check and the shuffle checks. Injecting the snapshot would have covered isCometLoaded and the two rules and missed all of those, which is the partial fix you were objecting to, only harder to see.

So the snapshot is now used as a check rather than an override. It records every SQL conf that was set, not just Comet's, and at report time anything that differs from the live session means the preview would describe a configuration the query never ran under. The report is skipped and the diagnostic names the keys that moved. Better nothing than coverage numbers for a configuration that never executed.

The gate is deliberately excluded from that comparison, so the case the snapshot was added for still works: restoring spark.comet.explain.planOnly.enabled before the callback runs does not drop the report. That flag decides only whether to report, and the snapshot remains the authority for it.

The gated test is now two tests sharing one helper, so the comparison you asked for is against a control rather than against nothing:

- a report survives the setting being restored before the callback runs
- a report is skipped when a conversion setting changed before the callback runs

The first still asserts one report and now also asserts it carries Comet accelerated numbers. The second flips spark.comet.exec.sort.enabled instead of the gate, and asserts one line that says settings changed since it was planned, that it names the changed key, and that it carries no coverage numbers.

Keep exchange identity restoration independent of projection support

Fixed by not adding a node at all. The subtree's own attribute IDs are rewritten to the wrapper's, everywhere they appear beneath it, instead of being re-labelled by a ProjectExec of aliases. That reads no user setting and adds no eligible operator, so it also answers your abstraction point about identity bookkeeping moving the number the report exists to state.

The reuse test is now parameterized on spark.comet.exec.project.enabled. Your prediction was exact; with the projection approach still in place and only that flag flipped:

- a reused exchange keeps its output IDs so consumers still convert (project.enabled=true)
- a reused exchange keeps its output IDs so consumers still convert (project.enabled=false) *** FAILED ***
  Comet accelerated 10 out of 14 eligible operators (71%) ... did not contain "CometHashJoin"

Both pass with the rewrite. All 24 CometPlanOnlySuite tests pass on the default profile, spotless and scalastyle are clean.

Two things about the rewrite worth stating, since it is more invasive-looking than an added node. It walks the whole subtree rather than just the root because almost every operator derives output from its children, so the producing leaf has to be rewritten for the root to present the new IDs; that is safe because attribute IDs are globally unique, so an ID in the root's output means the same attribute wherever it appears below, and intermediate attributes not in the root's output are untouched. And the direction is load-bearing: rewriting the consumers to reference the shared subtree's IDs instead would merge the two copies of a self-join onto one set of attributes, which is precisely what ReusedExchangeExec's re-aliasing exists to prevent.

[unikdahal]

Small wording nit: should this say "once per SQL execution" instead of "once per action"?

The RDD case documented in CometPlanOnly shows that the two aren't always the same.

[unikdahal]

I think this fallback can be risky.

If a normal query runs with plan-only off, it won't have the tag. If plan-only gets enabled before the async listener callback runs, this can pick up the new session config and report that older query as plan-only.

Would it be safer to just skip reporting when the tag is missing?

[andygrove]

superseded by #5394