viirya commented on code in PR #6991: URL: https://github.com/apache/datafusion-comet/pull/6991#discussion_r4304233291
########## spark/src/main/scala/org/apache/comet/rules/CometCacheBuildRule.scala: ########## @@ -0,0 +1,71 @@ +/* + * 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 scala.collection.mutable.ListBuffer + +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.comet.CometSparkToColumnarExec +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.{RowToColumnarExec, SparkPlan} +import org.apache.spark.sql.internal.StaticSQLConf + +import org.apache.comet.CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED +import org.apache.comet.CometSparkSessionExtensions.isCometLoaded + +/** + * Builds a cached adaptive plan that ends in a Spark row operator straight into Arrow. + * + * From Spark 4.0, when the cache serializer accepts columnar input, as Comet's does, Spark marks + * a cached `AdaptiveSparkPlanExec` columnar and AQE plans its final stage for columnar output. A + * final stage that ends in a row operator then gets a `RowToColumnarExec` on top, which fills + * Spark's own column vectors, and Comet's serializer copies those into Arrow before it writes + * them. `CometSparkToColumnarExec` writes Arrow straight from the rows, so the serializer stores + * its batches as they are, one copy instead of two. + * + * That `RowToColumnarExec` is the root of the plan these rules are handed. Spark plans no other + * stage for columnar output with a row operator at its root: every other query is planned for row + * output, an exchange stage is rooted at its exchange, and Comet's dynamic pruning subquery is + * rooted at a broadcast exchange. A transition below the root is left alone. + * + * Runs only when Comet's cache is enabled and its serializer is installed, so another serializer + * gets the batches Spark planned for it, and for schemas `CometSparkToColumnarExec` can convert. + * + * @param preview + * true in the plan-only preview, which shows the plan Comet would execute. Otherwise the rule + * leaves plans alone in plan-only mode, where Spark executes each query unchanged. + */ +case class CometCacheBuildRule(preview: Boolean = false) extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan match { + case transition @ RowToColumnarExec(child) if appliesTo(transition) => + val conversion = CometSparkToColumnarExec(child) Review Comment: Spark's `RowToColumnarExec` and Comet's own row path (`convertInternalRowToCachedBatch`) both size cached batches by `spark.sql.inMemoryColumnarStorage.batchSize`. `CometSparkToColumnarExec` sizes them by `spark.comet.batchSize` instead. After this change the same `df.cache()` gets cut into different batches depending on whether AQE planned it. On this path a user's `inMemoryColumnarStorage.batchSize` is also silently ignored. Batch size sets how fine-grained the cache's min/max pruning stats are, and how well each batch compresses. That's why `CometInMemoryCacheBenchmark` pins it to 10000 in every conf. Could the rule pass `conf.columnBatchSize` through to the conversion? One way is an optional batch-size parameter on `CometSparkToColumnarExec`. Then all three ways of building the cache would agree. ########## spark/src/main/scala/org/apache/comet/rules/CometCacheBuildRule.scala: ########## @@ -0,0 +1,71 @@ +/* + * 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 scala.collection.mutable.ListBuffer + +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.comet.CometSparkToColumnarExec +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.{RowToColumnarExec, SparkPlan} +import org.apache.spark.sql.internal.StaticSQLConf + +import org.apache.comet.CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED +import org.apache.comet.CometSparkSessionExtensions.isCometLoaded + +/** + * Builds a cached adaptive plan that ends in a Spark row operator straight into Arrow. + * + * From Spark 4.0, when the cache serializer accepts columnar input, as Comet's does, Spark marks + * a cached `AdaptiveSparkPlanExec` columnar and AQE plans its final stage for columnar output. A + * final stage that ends in a row operator then gets a `RowToColumnarExec` on top, which fills + * Spark's own column vectors, and Comet's serializer copies those into Arrow before it writes + * them. `CometSparkToColumnarExec` writes Arrow straight from the rows, so the serializer stores + * its batches as they are, one copy instead of two. + * + * That `RowToColumnarExec` is the root of the plan these rules are handed. Spark plans no other + * stage for columnar output with a row operator at its root: every other query is planned for row + * output, an exchange stage is rooted at its exchange, and Comet's dynamic pruning subquery is + * rooted at a broadcast exchange. A transition below the root is left alone. + * + * Runs only when Comet's cache is enabled and its serializer is installed, so another serializer + * gets the batches Spark planned for it, and for schemas `CometSparkToColumnarExec` can convert. + * + * @param preview + * true in the plan-only preview, which shows the plan Comet would execute. Otherwise the rule + * leaves plans alone in plan-only mode, where Spark executes each query unchanged. + */ +case class CometCacheBuildRule(preview: Boolean = false) extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan match { + case transition @ RowToColumnarExec(child) if appliesTo(transition) => + val conversion = CometSparkToColumnarExec(child) + transition.logicalLink.foreach(conversion.setLogicalLink) + conversion + case _ => plan + } + + private def appliesTo(transition: RowToColumnarExec): Boolean = + isCometLoaded(conf) && + COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf) && + (preview || !CometRule.planOnlyApplies(conf, transition)) && Review Comment: I don't think the `preview` branch can ever match. `CometRule.buildPreview` inserts transitions with `ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false)`, and with `outputsColumnar = false` Spark never puts a `RowToColumnarExec` at the root. So the preview plan never reaches the `case` in `apply`. Could we drop the `preview` parameter and its `@param` scaladoc, and keep only the `planOnlyApplies` guard? `preview` does matter in `CometCacheColumnarRule`, which transforms the whole tree. Keeping a copy here that has no effect suggests otherwise to the next reader. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
