This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6082-ccdff0a4bfeee377ffa8654a1ef7843e63b3b8b8 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 57ef3275d3d4e6e3c4d5f3de27b9b71daf736164 Author: Andy Grove <[email protected]> AuthorDate: Mon Sep 21 22:11:28 2026 +0000 refactor: compose CometScanRule and CometExecRule into a single CometRule (#6082) * 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. * test: pin the invariant that motivates composing the rules * 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. --- docs/source/contributor-guide/plugin_overview.md | 14 ++++--- .../apache/comet/CometSparkSessionExtensions.scala | 24 +++++------ .../scala/org/apache/comet/rules/CometRule.scala | 46 ++++++++++++++++++++++ .../apache/comet/rules/CometExecRuleSuite.scala | 39 ++++++++++++++++++ ...RevertNativeForTransitionHeavyStagesSuite.scala | 6 +-- 5 files changed, 106 insertions(+), 23 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/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..9b9639ac18 --- /dev/null +++ b/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)) +} 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..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,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) + } + } + } + } + } 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) } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
