refactor: compose CometScanRule and CometExecRule into a single CometRule

[andygrove]

Which issue does this PR close?

Closes #6081.

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

CometScanRule and CometExecRule are registered as two independent rules, adjacently, in both the columnar path and the query-stage-prep path. They are really one pass split in two, held together by registration order rather than by anything in the code.

Nothing ever runs between them: both registrations happen back to back inside one apply, so no other extension can interleave. And neither does the whole job on its own. CometExecRule.transform builds its native plan up from the nodes CometScanRule produces, so without scan conversion the scan stays on Spark's reader — everything above it is either refused for want of Arrow input or, with spark.comet.convert.parquet.enabled=true, converted over a CometSparkToColumnarExec bridge instead of a native scan. Scan conversion is not independently useful either: the V1 path is gated on spark.comet.exec.enabled, and CometScanExec is a planning intermediate whose doExecute throws, so CometExecRule has to replace it before execution either way.

The split leaks into anything that needs the whole conversion rather than half of it. It is what forces #5394 to widen both rules' _apply to package visibility and thread a forPreview flag through CometExecRule, so one rule can reach into the other and rebuild the conversion for a report. I plan to rebase #5394 onto this and drop that machinery.

What changes are included in this PR?

• New CometRule, which applies CometScanRule then CometExecRule.
• CometSparkSessionExtensions registers one ColumnarRule and one query-stage-prep rule instead of two of each. CometScanColumnar and CometExecColumnar collapse into CometColumnar; nothing referenced either by name.
• RevertNativeForTransitionHeavyStagesSuite.applyFullColumnarPipeline chained the two rules by hand; it now calls CometRule.
• Rule-order docstring and plugin_overview.md updated.

Both rules keep their own classes, files and tests. This only changes how they are sequenced.

One user-visible change: the plan change log

spark.sql.planChangeLog.rules filters on Rule.ruleName, and AdaptiveSparkPlanExec.applyPhysicalRules logs one entry per registered query-stage-prep rule. So anyone filtering that conf on org.apache.comet.rules.CometScanRule or org.apache.comet.rules.CometExecRule will stop seeing output and needs org.apache.comet.rules.CometRule instead.

spark.comet.explain.transformations is unaffected: CometRule delegates to each inner rule's apply, so each still logs under its own ruleName.

How are these changes tested?

Existing coverage, plus one new test in CometExecRuleSuite pinning the ordering invariant. With spark.comet.convert.parquet.enabled pinned off, CometExecRule alone leaves the FileSourceScanExec in place, while CometRule on the same query produces a CometNativeScanExec. I checked it is not vacuous by mutating CometRule: it fails both when the scan phase is dropped and when the two phases are swapped. RevertNativeForTransitionHeavyStagesSuite now exercises CometRule too.

The load-bearing evidence is plan stability. CometTPCDSV1_4_PlanStabilitySuite (97) and CometTPCDSV2_7_PlanStabilitySuite (32) both pass, so every golden plan is byte-identical to before.

Also green locally: CometScanRuleSuite, CometExecRuleSuite, RevertNativeForTransitionHeavyStagesSuite, CometScanSchemeFallbackSuite (66), CometExecSuite (146), and test-compile on the default profile and -Pspark-3.4.

This touches rule registration, so it wants the broader suites before merge — I have applied run-all-spark-profiles and run-spark-4.1-tests.

[rich7420]

LGTM

[mbutrovich]

Thanks @andygrove.

The docstring and the new test claim more than CometExecRule guarantees. Two files outside the diff also need changes.

• docs/source/contributor-guide/plugin_overview.md still says the extension "registers two physical plan optimization rules with Spark: CometScanRule and CometExecRule" (plugin_overview.md#L50-L52). Could you update that paragraph in this PR to describe CometRule and keep the two subsections as the two phases it runs?
• RevertNativeForTransitionHeavyStagesSuite.applyFullColumnarPipeline chains CometScanRule(spark).apply and CometExecRule(spark).apply by hand (RevertNativeForTransitionHeavyStagesSuite.scala#L46-L51). That's the "caller that needs the whole conversion" the PR describes. Switching it to CometRule(spark).apply would also give CometRule its only direct test coverage outside the plan stability suites.

[mbutrovich]

"Operator conversion against unconverted Spark scans converts nothing" isn't accurate. CometExecRule wraps an unconverted FileSourceScanExec in CometSparkToColumnarExec when spark.comet.convert.parquet.enabled=true (CometExecRule.scala#L411-L412), and then converts the operators above it. It also converts local table scans, empty relations and shuffles with no scan conversion at all. The suggestion below states the narrower property.

The docstring's last sentence is true for spark.comet.explain.transformations, but Spark's own plan change log now sees one rule. AdaptiveSparkPlanExec.applyPhysicalRules calls logger.logRule(rule.ruleName, ...) once per registered rule (Spark AdaptiveSparkPlanExec.scala#L856-L860). Anyone filtering spark.sql.planChangeLog.rules on org.apache.comet.rules.CometScanRule or CometExecRule will stop seeing output and must switch to org.apache.comet.rules.CometRule. The PR description should mention this, and the docstring should say how the rule appears in that log.

Most of this docstring is archaeology. "Previously registered as separate rules, adjacently", "nothing ever ran between them" and "keep their own classes, files and tests" describe what this PR changed, not the code. None of it is needed to understand the current code, so it should go. It belongs in the PR description and commit message.

Suggested change

	* The two were previously registered as separate rules, adjacently, in both the columnar and the
	* query-stage-prep paths. Nothing ever ran between them, and neither is useful on its own -
	* [[CometExecRule]] seeds its native chain only from the nodes [[CometScanRule]] produces
	* (`CometScanExec`, `CometBatchScanExec`, `CometContribScanMarker`), so operator conversion
	* against unconverted Spark scans converts nothing. Composing them here makes that ordering an
	* invariant of the code rather than of the registration order, and gives callers that need the
	* whole conversion - rather than half of it - a single entry point.
	*
	* The two rules keep their own classes, files and tests; this only fixes how they are sequenced.
	* `ruleName` in `spark.comet.explain.transformations` output is still each inner rule's own,
	* since this delegates to their `apply`.
	* Native scans come only from the nodes [[CometScanRule]] produces (`CometScanExec`,
	* `CometBatchScanExec`, `CometContribScanMarker`), so [[CometExecRule]] must run after it.
	* Running [[CometExecRule]] alone leaves scans on Spark's readers. Composing the two here makes
	* that ordering part of the code instead of depending on the order the rules are registered, and
	* gives callers that need the whole conversion a single entry point.
	*
	* `spark.comet.explain.transformations` logs each inner rule under its own `ruleName`. Spark's
	* `spark.sql.planChangeLog.*` output logs this rule as a single step named `CometRule`.

[andygrove]

Fixed, thanks — you're right on all three. Took your suggestion nearly verbatim for the first paragraph, and dropped the archaeology.

On the plan change log: I checked, and it's only the query-stage-prep registration that gets a per-rule entry. In the non-AQE columnar path ApplyColumnarRulesAndInsertTransitions is itself the logged rule and the columnar rules inside it are never named, so the old names never appeared there to begin with. The docstring now says query-stage preparation logs this pass as org.apache.comet.rules.CometRule, and I've added a section to the PR description calling out that anyone filtering spark.sql.planChangeLog.rules on either old name has to switch.

[mbutrovich]

This test passes only because spark.comet.convert.parquet.enabled defaults to false (CometConf.scala#L180-L186). With it on, the same plan comes back with a CometSparkToColumnarExec over the scan and native operators above it, so the name and the comment describe an invariant that doesn't hold. It also tests CometExecRule in isolation and never runs CometRule, which is the class this PR adds.

Could it test the property CometRule guarantees instead? Set COMET_CONVERT_FROM_PARQUET_ENABLED to false explicitly so the test doesn't lean on a default, assert that CometExecRule alone leaves the FileSourceScanExec in place, and assert that CometRule on the same plan produces a CometNativeScanExec. That fails if someone reorders or drops the scan phase. The current test doesn't.

test("scan conversion must run before operator conversion") {
  ...
  withSQLConf(
    CometConf.COMET_ENABLED.key -> "true",
    CometConf.COMET_EXEC_ENABLED.key -> "true",
    CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "false") {
    val plan = stripAQEPlan(sparkPlan)
    assert(countOperators(CometExecRule(spark).apply(plan), classOf[FileSourceScanExec]) == 1)
    assert(countOperators(CometRule(spark).apply(plan), classOf[CometNativeScanExec]) == 1)
  }
}

[andygrove]

Good catch, the test was leaning on the default and never touched the class the PR adds. Rewrote it along the lines you suggested.

One wrinkle worth recording: your version reuses a single plan for both assertions, and that doesn't work. Fallback reasons are tags on the Spark nodes, and CometNativeScan.isSupported returns false immediately when the scan already carries one, so the CometExecRule-alone pass poisons the FileSourceScanExec and CometRule on the same tree then declines the scan too — the second assertion gets 0. Each direction now gets its own plan from createSparkPlan.

I also checked it isn't vacuous by mutating CometRule: it fails when the scan phase is dropped, and it fails when the two phases are swapped.

[andygrove]

Thanks @mbutrovich, all four addressed in 5ee8d8f.

plugin_overview.md now says the extension registers one rule and describes CometRule as two phases, with the existing subsections retitled "Phase 1: CometScanRule" and "Phase 2: CometExecRule" so the ordering is visible from the headings. RevertNativeForTransitionHeavyStagesSuite.applyFullColumnarPipeline calls CometRule now — that suite was the only place in the repo chaining the two by hand; the other tests that instantiate one rule are deliberately testing that rule in isolation, so I left them.

The docstring and test changes are in the two threads above. Re-ran CometExecRuleSuite, CometScanRuleSuite and RevertNativeForTransitionHeavyStagesSuite green, plus -Pspark-3.4 test-compile.

[mbutrovich]

Thanks @andygrove!

[sunchao]

Correctness

Reviewed 5ee8d8f4 against 5ca14992. I found no new or remaining P1/P2 issue in this refactor.

Previously, the extension registered scan conversion and operator conversion as adjacent rules in both the columnar and AQE preparation paths. CometRule now calls the same two public apply methods in the same order. The inner rule implementations are unchanged. Scan fallback tags reach operator conversion as before, and a rejected V1 native scan still returns to its wrapped Spark scan. The enabled, streaming, schema and operator-support guards retain their existing behavior, including their type, null and error handling.

I checked registration and execution against the maintained Spark 3.5 and 4.0 sources. Both preserve builder order, run columnar pre-rules before inserting transitions, and run post-rules in reverse order. The removed scan columnar wrapper supplied only Spark's default identity post-rule. The remaining post-transition logic is unchanged, as is the placement of the Spark 3.4 DPP fallback shim before conversion. AQE preparation still reuses the phase instances, while the columnar hook still constructs fresh ones per invocation. The wrapper adds no traversal into subqueries or materialized stages.

The earlier documentation and test requests are addressed. The new regression explicitly disables the Parquet bridge, checks the exec-only control, and calls CometRule for native-scan conversion. Separate fresh plans avoid carrying fallback tags from the control into the wrapper case. The contributor overview describes one registered rule with two phases, and the full-pipeline helper now calls that rule.

One visible diagnostic change is documented: AQE's spark.sql.planChangeLog.rules filter must match org.apache.comet.rules.CometRule. Comet's spark.comet.explain.transformations logging still uses each inner rule's name.

Validation

The local source checks passed, including 54 composition, registration and provenance assertions plus 46 checked source ranges. I did not run a local Scala build or Spark/JNI tests. The author's suite and mutation-test results remain author-reported. Maintained Spark 3.4 and 4.1 sources were unavailable for this review.

At the review cutoff, Comet CI remained incomplete, with build/test work queued or running and several SQL-test jobs skipped. The completed Scala syntactic lint and general lint jobs checked out merge 535ed565, which combines this head with 09b44ad6. Its reviewed files match the head, but its whole tree differs from the head. Those lint results do not establish runtime correctness.

Performance

The two existing conversion phases still run once per invocation, with the same internal traversals. The wrapper adds a constant-sized rule object and direct calls. It does not clone plans, introduce a retry loop, or add work per input row. Combining the registrations also removes a default identity post-rule and changes Spark's logging boundary. I found no material new cost requiring a benchmark, and this review makes no query-time or planning-time improvement claim.

Design

A single entry point makes the scan-before-operator dependency explicit for both extension registration and callers that need complete conversion. Keeping the phase classes separate preserves focused tests and makes each phase's responsibilities clear. The existing version shims and post-transition rules stay in their established positions. I found no simpler change that would preserve these benefits with less coupling.

Abstraction & complexity

The abstraction is small: two private delegates and one ordered apply method. It centralizes the required order without duplicating conversion logic or introducing a general pipeline framework. Updating the full-pipeline test helper and contributor documentation keeps the surrounding code consistent with that contract.