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) }