From 7fa90eed4a61fd5a2fa7fe5de8b0f099e2065762 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 09:51:53 -0600 Subject: [PATCH 1/3] refactor: compose CometScanRule and CometExecRule into one CometRule The two were registered as separate rules, adjacently, in both the columnar and the query-stage-prep paths. Nothing ever ran between them, and neither is useful alone: CometExecRule seeds its native chain only from the nodes CometScanRule produces, so operator conversion over unconverted Spark scans converts nothing. Compose them so the ordering is an invariant of the code rather than of the registration order, and so callers that need the whole conversion have one entry point. Both rules keep their classes, files and tests. --- .../comet/CometSparkSessionExtensions.scala | 24 ++++------ .../org/apache/comet/rules/CometRule.scala | 47 +++++++++++++++++++ 2 files changed, 56 insertions(+), 15 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/rules/CometRule.scala diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 4b37d7b61a..749a3d6ecb 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/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 = diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala new file mode 100644 index 0000000000..ac61a461d3 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -0,0 +1,47 @@ +/* + * 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. + * + * 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`. + */ +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)) +} From 8dd66d41adda0afa290e83a271fdf5c1881c10b0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 10:37:02 -0600 Subject: [PATCH 2/3] test: pin the invariant that motivates composing the rules --- .../comet/rules/CometExecRuleSuite.scala | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 9d65dec834..a746ce7c8e 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -1086,4 +1086,30 @@ class CometExecRuleSuite extends CometTestBase { } } + test("operator conversion alone converts nothing without scan conversion") { + withTempPath { path => + createTestDataFrame.write.parquet(path.toString) + withTempView("test_data") { + spark.read.parquet(path.toString).createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT id, id * 2 as doubled FROM test_data WHERE id % 2 == 0") + assert(countOperators(sparkPlan, classOf[FileSourceScanExec]) == 1) + + // CometExecRule seeds its native chain only from the nodes CometScanRule produces, so on + // its own it converts nothing: the scan is untouched and every operator above it is + // refused for want of Arrow input. This is why the two are composed into `CometRule` + // rather than registered as independent rules that happen to run in the right order. + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true") { + val execOnly = CometExecRule(spark).apply(stripAQEPlan(sparkPlan)) + assert(countOperators(execOnly, classOf[FileSourceScanExec]) == 1) + assert( + stripAQEPlan(execOnly).collect { case p: CometNativeExec => p }.isEmpty, + s"operator conversion alone should convert nothing, got:\n$execOnly") + } + } + } + } + } From 5ee8d8f40d1e2326a6db03c775ffaa8d6000654e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 13:28:22 -0600 Subject: [PATCH 3/3] review: tighten the CometRule docstring, test and docs The docstring claimed operator conversion against unconverted Spark scans converts nothing. That is wrong: with spark.comet.convert.parquet.enabled=true CometExecRule bridges the scan with a CometSparkToColumnarExec and converts the operators above it, and it converts local table scans, empty relations and shuffles with no scan conversion at all. State the narrower property instead, drop the archaeology about what this PR changed, and record how the pass is named in Spark's plan change log. The test named an invariant that only held because spark.comet.convert.parquet.enabled defaults to false, and it never ran CometRule. Pin the conf explicitly and assert both directions: CometExecRule alone leaves the FileSourceScanExec in place, CometRule on the same query produces a CometNativeScanExec. Each direction needs its own plan, because fallback reasons are node tags and CometNativeScan.isSupported declines a scan already carrying one. Also switch RevertNativeForTransitionHeavyStagesSuite.applyFullColumnarPipeline, which chained the two rules by hand, to CometRule, and update plugin_overview.md to describe the single registered rule and its two phases. --- .../contributor-guide/plugin_overview.md | 14 ++++--- .../org/apache/comet/rules/CometRule.scala | 19 +++++---- .../comet/rules/CometExecRuleSuite.scala | 39 ++++++++++++------- ...tNativeForTransitionHeavyStagesSuite.scala | 6 +-- 4 files changed, 47 insertions(+), 31 deletions(-) diff --git a/docs/source/contributor-guide/plugin_overview.md b/docs/source/contributor-guide/plugin_overview.md index c3330a65d6..6dfb0bdf3c 100644 --- a/docs/source/contributor-guide/plugin_overview.md +++ b/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`. diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala index ac61a461d3..9b9639ac18 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -26,17 +26,16 @@ import org.apache.spark.sql.execution.SparkPlan /** * Comet's plan conversion pass: scan conversion followed by operator conversion. * - * 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. + * 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. * - * 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`. + * `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] { diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index a746ce7c8e..8bf6097754 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -1086,27 +1086,40 @@ class CometExecRuleSuite extends CometTestBase { } } - test("operator conversion alone converts nothing without scan conversion") { + 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 sparkPlan = - createSparkPlan(spark, "SELECT id, id * 2 as doubled FROM test_data WHERE id % 2 == 0") - assert(countOperators(sparkPlan, classOf[FileSourceScanExec]) == 1) + 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) - // CometExecRule seeds its native chain only from the nodes CometScanRule produces, so on - // its own it converts nothing: the scan is untouched and every operator above it is - // refused for want of Arrow input. This is why the two are composed into `CometRule` - // rather than registered as independent rules that happen to run in the right order. withSQLConf( CometConf.COMET_ENABLED.key -> "true", - CometConf.COMET_EXEC_ENABLED.key -> "true") { - val execOnly = CometExecRule(spark).apply(stripAQEPlan(sparkPlan)) - assert(countOperators(execOnly, classOf[FileSourceScanExec]) == 1) + 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( - stripAQEPlan(execOnly).collect { case p: CometNativeExec => p }.isEmpty, - s"operator conversion alone should convert nothing, got:\n$execOnly") + countOperators( + CometRule(spark).apply(forCometRule), + classOf[CometNativeScanExec]) == 1) } } } diff --git a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala index 31fac6f459..0db2351c72 100644 --- a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala +++ b/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) }