feat: add spark.comet.explain.planOnly.enabled, reported per query - #5514
feat: add spark.comet.explain.planOnly.enabled, reported per query#5514andygrove wants to merge 2 commits into
Conversation
Report the Comet plan Comet would have executed for a query, without offloading any of it to Comet, so a workload can be assessed without changing how it runs. The conversion rules return the plan untouched while the mode is on, and the report is built afterwards from the plan Spark executed, by a query execution listener. Reporting once the query is over rather than while it is being planned is what keeps this small: Spark applies a planner rule many times for one query - once per query stage and once per adaptive re-optimization under AQE, and separately for every subquery it prepares - and telling those applications apart takes a state machine. A listener fires once per action, holding the finished plan, so there is nothing to tell apart. Building the report from the plan that ran means undoing the part of Spark's preparation that follows the conversion rules: the adaptive wrappers, the codegen wrappers, the columnar transitions, and the exchange reuse. It 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. Tests assert both properties that matter: the executed plan holds no Comet operator, and the report's coverage equals CometCoverageStats for the plan Comet really executes - for an aggregate, a shuffled join, a scalar subquery, a DPP subquery and a reverted stage, with AQE on and off.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
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.
| */ | ||
| def register(session: SparkSession): Unit = { | ||
| if (registeredSessions.add(session)) { | ||
| session.listenerManager.register(new CometPlanOnlyListener) |
There was a problem hiding this comment.
[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.
| val conf = session.sessionState.conf | ||
| if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && isCometLoaded(conf) && |
There was a problem hiding this comment.
[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.
| // conversion would never reach the subtree while the coverage count - which unwraps the | ||
| // wrapper - still counts every operator in it as Spark. Undo the reuse and let both copies | ||
| // convert, which is what the counts of a real Comet run reflect. | ||
| case reused: ReusedExchangeExec => stripPreparation(reused.child) |
There was a problem hiding this comment.
[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.
2 participants
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 aqueryStagePrepflag 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
QueryExecutionListenerinstead. 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?
spark.comet.explain.planOnly.enabled, default off.CometScanRuleandCometExecRulereturn the plan untouched at the top ofapplywhile the mode is on, so Spark plans and executes the query exactly as it would with Comet off.CometExecRulealso registers the session's listener, so a session carries one only if the mode is used, and the config stays togglable mid-session.CometPlanOnlybuilds the report fromqe.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.RevertNativeForTransitionHeavyStagesgainsapplyToAllStages, because the preview holds a whole plan where AQE would have handed that rule one stage at a time.CREATE VIEW,SHOW TABLES) are not reported; a session runs enough of them that one 0% report each would bury the rest.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:CometPlanoperator and that the intended scan path was actually exercised.CometCoverageStatsfor 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:
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,RevertNativeForTransitionHeavyStagesSuiteandCometCoverageStatsSuiteare green (64 tests) on the default profile,test-compileis clean onspark-3.4,spark-3.5,spark-4.0andspark-4.2, and scalafix passes with the semantic rules.