feat: add spark.comet.explain.planOnly.enabled - #5394
Conversation
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 apache#5335.
- 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.
`_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`.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
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.
| if (!forPreview && CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get()) { | ||
| val executionId = Option( | ||
| session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)) | ||
| if (CometExecRule.markPlanOnlyReported(executionId)) { |
There was a problem hiding this comment.
[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.
| * both rules run their normal transforms instead of short-circuiting. | ||
| */ | ||
| private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = { | ||
| val preview = _apply(CometScanRule(session)._apply(plan), forPreview = true) |
There was a problem hiding this comment.
[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.
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.
|
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 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. |
| df.collect() | ||
| val executed = CometCoverageStats.forPlan(df.queryExecution.executedPlan) | ||
|
|
||
| val reports = withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") { |
There was a problem hiding this comment.
[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.
| plan: SparkPlan, | ||
| queryStagePrep: Boolean): Boolean = { | ||
| executionId match { | ||
| case None => |
There was a problem hiding this comment.
[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.
| * holds the whole plan at once, whereas under AQE Spark hands that rule one stage at a time. | ||
| */ | ||
| private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = { | ||
| val converted = _apply(CometScanRule(session)._apply(plan), forPreview = true) |
There was a problem hiding this comment.
[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.
2 participants
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?
spark.comet.explain.planOnly.enabled, default off.CometScanRuleandCometExecRuleshort-circuit at the top of theirapplyand return the plan untouched.CometExecRulebuilds a preview by calling both rules' private_applymethods directly (with an explicitforPreview = trueparameter that skips the short-circuit on the recursive call), logs the resulting Comet plan at WARN level, then returns the original Spark plan.normalizePlan/RewriteJoin/tagUnsafePartialAggregatesare not run and thrown away.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
CometExecRuleSuitecover:Each test asserts the executed plan has zero
CometPlanoperators.