refactor: compose CometScanRule and CometExecRule into a single CometRule

docs/source/contributor-guide/plugin_overview.md

@@ -47,11 +47,15 @@ The plugin also registers `CometSparkSessionExtensions` with Spark's extension A
 
 ## CometSparkSessionExtensions
 
-On initialization, this class registers two physical plan optimization rules with Spark: `CometScanRule`
-and `CometExecRule`. These rules run whenever a query stage is being planned during Adaptive Query Execution, and
-run once for the entire plan when Adaptive Query Execution is disabled.
+On initialization, this class registers one physical plan optimization rule with Spark: `CometRule`. It runs whenever
+a query stage is being planned during Adaptive Query Execution, and runs once for the entire plan when Adaptive Query
+Execution is disabled.
 
-### CometScanRule
+`CometRule` is two phases, applied in order: scan conversion (`CometScanRule`), then operator conversion
+(`CometExecRule`). The order matters, because operator conversion builds its native plan up from the nodes that scan
+conversion produces. Each phase is described below.
+
+### Phase 1: CometScanRule
 
 `CometScanRule` replaces any Parquet scans with Comet operators. There are different paths for Spark v1 and v2 data sources.
 
@@ -68,7 +72,7 @@ convert the output from Spark's scan to Arrow arrays. Note that both `spark.come
 Refer to the [Supported Spark Data Types](https://datafusion.apache.org/comet/user-guide/datatypes.html) section
 in the contributor guide to see a list of currently supported data types.
 
-### CometExecRule
+### Phase 2: CometExecRule
 
 This rule traverses bottom-up from the original Spark plan and attempts to replace each operator with a Comet equivalent.
 For example, a `ProjectExec` will be replaced by `CometProjectExec`.

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -34,7 +34,7 @@ import org.apache.spark.sql.internal.SQLConf
 
 import org.apache.comet.CometConf._
 import org.apache.comet.iceberg.IcebergWriteStrategy
-import org.apache.comet.rules.{CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages}
+import org.apache.comet.rules.{CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages}
 import org.apache.comet.shims.ShimCometSparkSessionExtensions
 
 /**
@@ -49,7 +49,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions
  *   2. PlanSubqueries               -- Spark creates SubqueryExec for scalar subqueries
  *   3. EnsureRequirements            -- Spark inserts shuffles/sorts
  *   4. ApplyColumnarRulesAndInsertTransitions:
- *      a. preColumnarTransitions:   CometScanRule, CometExecRule
+ *      a. preColumnarTransitions:   CometRule (CometScanRule then CometExecRule)
  *         - CometExecRule.convertSubqueryBroadcasts converts SubqueryBroadcastExec to
  *           CometSubqueryBroadcastExec for exchange reuse with Comet broadcasts
  *      b. insertTransitions:        ColumnarToRow/RowToColumnar added
@@ -62,7 +62,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions
  * {{{
  *   Initial plan:
  *     PlanAdaptiveSubqueries:       creates SubqueryAdaptiveBroadcastExec (SAB) for AQE DPP
- *     queryStagePreparationRules:   CometScanRule, CometExecRule
+ *     queryStagePreparationRules:   CometRule (CometScanRule then CometExecRule)
  *       - CometExecRule.convertSubqueryBroadcasts wraps SABs in
  *         CometSubqueryAdaptiveBroadcastExec to prevent Spark's
  *         PlanAdaptiveDynamicPruningFilters from replacing DPP with Literal.TrueLiteral
@@ -75,7 +75,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions
  *           CometSubqueryBroadcastExec with BroadcastQueryStageExec for broadcast reuse
  *        d. CometReuseSubquery                       -- deduplicates converted subqueries
  *     2. postStageCreationRules -> ApplyColumnarRulesAndInsertTransitions:
- *        a. preColumnarTransitions: CometScanRule, CometExecRule (no-ops, already converted)
+ *        a. preColumnarTransitions: CometRule (no-op, already converted)
  *        b. insertTransitions
  *        c. postColumnarTransitions: RevertNativeForTransitionHeavyStages,
  *                                    EliminateRedundantTransitions
@@ -91,25 +91,19 @@ class CometSparkSessionExtensions
     with Logging
     with ShimCometSparkSessionExtensions {
   override def apply(extensions: SparkSessionExtensions): Unit = {
-    extensions.injectColumnar { session => CometScanColumnar(session) }
-    extensions.injectColumnar { session => CometExecColumnar(session) }
+    extensions.injectColumnar { session => CometColumnar(session) }
     // Pre-3.5 only: tag AQE DPP regions so the conversion rules below leave them Spark-native.
-    // Registered before CometScanRule/CometExecRule so tags are in place when conversion runs.
+    // Registered before CometRule so tags are in place when conversion runs.
     // No-op on Spark 3.5+; see CometSpark34AqeDppFallbackRule's class docstring.
     injectPreSpark35QueryStagePrepRuleShim(extensions, CometSpark34AqeDppFallbackRule)
-    extensions.injectQueryStagePrepRule { session => CometScanRule(session) }
-    extensions.injectQueryStagePrepRule { session => CometExecRule(session) }
+    extensions.injectQueryStagePrepRule { session => CometRule(session) }
     injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters)
     injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery)
     extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) }
   }
 
-  case class CometScanColumnar(session: SparkSession) extends ColumnarRule {
-    override def preColumnarTransitions: Rule[SparkPlan] = CometScanRule(session)
-  }
-
-  case class CometExecColumnar(session: SparkSession) extends ColumnarRule {
-    override def preColumnarTransitions: Rule[SparkPlan] = CometExecRule(session)
+  case class CometColumnar(session: SparkSession) extends ColumnarRule {
+    override def preColumnarTransitions: Rule[SparkPlan] = CometRule(session)
 
     override def postColumnarTransitions: Rule[SparkPlan] = {
       val rules =

spark/src/main/scala/org/apache/comet/rules/CometRule.scala

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.SparkPlan
+
+/**
+ * Comet's plan conversion pass: scan conversion followed by operator conversion.
+ *
+ * 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 the order the rules are registered in, and gives
+ * callers that need the whole conversion a single entry point.
+ *
+ * `spark.comet.explain.transformations` logs each inner rule under its own `ruleName`, since this
+ * delegates to their `apply`. Spark's own plan change log sees one rule: query-stage preparation
+ * logs this pass as `org.apache.comet.rules.CometRule`, which is the name
+ * `spark.sql.planChangeLog.rules` has to match.
+ */
+case class CometRule(session: SparkSession) extends Rule[SparkPlan] {
+
+  private val scanRule = CometScanRule(session)
+  private val execRule = CometExecRule(session)
+
+  override def apply(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan))
+}

spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala

@@ -1086,4 +1086,43 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
+  test("scan conversion must run before operator conversion") {
+    withTempPath { path =>
+      createTestDataFrame.write.parquet(path.toString)
+      withTempView("test_data") {
+        spark.read.parquet(path.toString).createOrReplaceTempView("test_data")
+        val query = "SELECT id, id * 2 as doubled FROM test_data WHERE id % 2 == 0"
+
+        // One plan per rule application. Fallback reasons are recorded as tags on the Spark
+        // nodes, and CometNativeScan.isSupported declines a scan already carrying one, so
+        // reusing the plan the exec rule just refused would hold the second case down.
+        val forExecRule = stripAQEPlan(createSparkPlan(spark, query))
+        val forCometRule = stripAQEPlan(createSparkPlan(spark, query))
+        assert(countOperators(forExecRule, classOf[FileSourceScanExec]) == 1)
+        assert(countOperators(forCometRule, classOf[FileSourceScanExec]) == 1)
+
+        withSQLConf(
+          CometConf.COMET_ENABLED.key -> "true",
+          CometConf.COMET_EXEC_ENABLED.key -> "true",
+          // Off by default, but pinned here: with it on, CometExecRule bridges the unconverted
+          // scan with a CometSparkToColumnarExec and converts the operators above it, which is a
+          // different path from the one under test.
+          CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "false") {
+          // CometExecRule builds its native plan up from the nodes CometScanRule produces, so on
+          // its own it leaves the scan on Spark's reader.
+          assert(
+            countOperators(
+              CometExecRule(spark).apply(forExecRule),
+              classOf[FileSourceScanExec]) == 1)
+          // CometRule runs both phases, in that order. This fails if the scan phase is ever
+          // reordered or dropped.
+          assert(
+            countOperators(
+              CometRule(spark).apply(forCometRule),
+              classOf[CometNativeScanExec]) == 1)
+        }
+      }
+    }
+  }
+
 }

spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala

@@ -44,9 +44,9 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
   }
 
   private def applyFullColumnarPipeline(plan: SparkPlan): SparkPlan = {
-    val cometPlan = CometScanRule(spark).apply(plan)
-    val execPlan = CometExecRule(spark).apply(cometPlan)
-    val withTransitions = ApplyColumnarRulesAndInsertTransitions(Seq.empty, false).apply(execPlan)
+    val cometPlan = CometRule(spark).apply(plan)
+    val withTransitions =
+      ApplyColumnarRulesAndInsertTransitions(Seq.empty, false).apply(cometPlan)
     EliminateRedundantTransitions(spark).apply(withTransitions)
   }