andygrove opened a new pull request, #6991: URL: https://github.com/apache/datafusion-comet/pull/6991
## Which issue does this PR close? Closes #6952. ## Rationale for this change 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. When that stage ends in a row operator, Spark puts a `RowToColumnarExec` on top. It fills Spark's own `OnHeapColumnVector`s, and Comet's serializer then copies those into Arrow. That's two copies where the row path, which Spark uses without AQE and on 3.x, makes one, and it happens whether or not Comet runs any part of the query. ## What changes are included in this PR? - A new post-columnar rule, `CometCacheBuildRule`, replaces a `RowToColumnarExec` at the root of the plan with `CometSparkToColumnarExec`, which writes Arrow straight from the rows. The cache serializer then stores those batches as they are. - Only a transition at the root is replaced. That is where Spark puts the one it adds for a stage's columnar output, and no other stage that Spark plans for columnar output has 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. - The rule has the same guards as `CometCacheColumnarRule`. It runs when Comet is loaded and `spark.comet.exec.inMemoryCache.enabled` is on, and not in plan-only mode, so it also covers a session with Comet exec turned off. It also requires two more things. Comet's serializer has to be the installed `spark.sql.cache.serializer`, so that another serializer still gets the batches Spark planned for it. And the schema has to be one `CometSparkToColumnarExec` can convert. Otherwise Spark's transition stays. - It runs after `EliminateRedundantTransitions` and before `CometCacheColumnarRule`. ## How are these changes tested? Two new tests in `CometInMemoryCacheSuite`: - "AQE builds a cache whose final stage ends in a Spark operator straight into Arrow" caches an AQE query whose final stage ends in a Spark `Project` over a native scan and shuffle. It checks that the cached final plan has `CometSparkToColumnarExec` at the root, with output batches, and no `RowToColumnarExec`. It also covers the cases that keep Spark's transition: Comet's cache turned off, plan-only mode, and an `array<int>` column, which `CometSparkToColumnarExec` does not convert. In every case the relation is stored in Comet's format and the rows match Spark's. The test runs on Spark 4.0+, because 3.x builds such a cache through rows. - "CometCacheBuildRule converts a root transition only for Comet's cache serializer" applies the rule directly to a transition it builds. The transition is converted at the root, left alone below it, and left alone under a conf that names Spark's default serializer. The serializer is a static conf, and this suite pins it to Comet's, so a direct test is the only way to cover that guard. Six mutants each fail one of the tests: not registering the rule, dropping the cache, plan-only, serializer or schema guard, and matching transitions below the root. `CometInMemoryCacheSuite`, `CometInMemoryCachePruningSuite`, the three Kryo suites and `CachedBatchRowIteratorSuite` pass locally on Spark 3.4, 3.5, 4.0, 4.1 and 4.2. Scalafix and the strict-warnings build pass on 3.5. I measured the cache build of #6952's query with a probe that isn't part of this PR. The query reads 5M rows of bigint, int, double and string columns and runs `repartition(8, $"a")` followed by a projection. Every arm ran in the same JVM, interleaved, nine builds each, on an Apple M3 Max with JDK 17, Spark 4.1 and a debug native library. The `main` arm turns `spark.comet.exec.inMemoryCache.enabled` off while the cache is built. That skips this rule and leaves the rest of the build the same. | Final stage | `main` | This PR | Without AQE (one copy) | | --- | ---: | ---: | ---: | | Spark operators only (Comet exec and scan off) | 1018 ms | 868 ms | 864 ms | | A Spark `Project` over a native scan and shuffle | 2808 ms | 2690 ms | | With AQE on, a cache built by Spark operators now costs what the row path does. -- 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]
