[andygrove]

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 OnHeapColumnVectors, 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.

[sunchao]

Reviewed the entire four-file diff at e5ddf6e427e92a3c679c644587a670ac158b6781 against 7d8d9b1544b734c85a93cca2cc9d5380ad70accd. The PR remains non-draft. Existing reviews, conversation comments, inline comments and threads were empty.

Routed skills: review-comet-pr and review-comet-ffi-pr, with the contributor guidance on Arrow ownership, memory accounting and timestamps.

Summary

• Prior state and problem: On Spark 4.0+, an AQE cache build ending in a row operator first populated Spark column vectors, which Comet then copied into Arrow. This added an unnecessary conversion even when native execution was disabled.
• Design approach: Replace a root RowToColumnarExec with the existing CometSparkToColumnarExec when Comet is loaded, its cache is enabled, its serializer is configured and the schema is supported.
• Correctness: Checked Spark’s cache planning and transition insertion against upstream sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. Traced Arrow layout compatibility, null handling, timestamp representation and task cleanup. The replacement preserves output attributes, ordering, partitioning and the logical link. No introduced P1/P2 correctness issue was identified.
• Compatibility analysis: Spark 3.x retains its existing adaptive-cache row path. Spark 4.x supplies the root transition this rule targets. Other serializers and unsupported schemas retain Spark’s transition. Comparison with branch-1.1 confirms the intended cache-build optimization, with no additional unintended compatibility change identified.
• Key design decisions: Matching only the root avoids changing transitions serving other columnar consumers. The serializer and schema checks constrain the replacement, while the plan-only guard preserves execution behavior.
• Implementation sketch: CometRule.postColumnarRules registers the new rule after EliminateRedundantTransitions and before CometCacheColumnarRule. The extension’s ordering documentation is updated. Two tests exercise the AQE build and the rule’s guards.
• Performance: The new path removes the intermediate Spark-vector conversion and lets the serializer consume compatible Arrow vectors directly. The rule adds no recursive plan traversal. The author’s benchmark timings were not independently reproduced. No evidence-backed P1/P2 performance regression was found.
• Design: A small post-columnar rule fits Spark’s planning lifecycle and keeps cache-build routing separate from cache-read optimization. The implementation is straightforward to reason about.
• Abstraction & complexity: Reusing CometSparkToColumnarExec, its schema predicate and existing Arrow readers avoids another conversion implementation or ownership mechanism. No unnecessary abstraction meeting the reporting bar was identified.
• Behavioral changes worth calling out: Eligible cache builds can use the optimization with Comet execution disabled, provided Comet’s serializer remains installed. Converted batches follow spark.comet.batchSize instead of Spark’s columnBatchSize. Cached format and query results remain unchanged.
• Suggested improvements: No additional change meeting the P1/P2 reporting bar is requested.

Validation: The exact-head JVM reactor completed test-compile on Spark 4.1.3 with JDK 17. CometInMemoryCacheSuite passed, including both new tests. ScalaTest reported 74 successful test executions across the full suite and an additional selected AQE test run, with zero failures or aborts. Tests used the native artifact from the exact base SHA, whose native sources are unchanged by this PR. The working tree remains clean.

Exact-head CI: Preflight, lint, Java lint across the configured Spark profiles, CodeQL and Celeborn compatibility checks passed. No failures were reported at the final inspection. Native build, Rust tests, strict Scala warnings and the labeled Spark 4.1 SQL build were still running. Full Spark SQL results remain pending. Other Spark versions were checked against source but were not run locally.

No introduced P1/P2 issues found within this review. No substantiated existing blockers remain.

[viirya]

Thanks for this, Andy. Keeping the change to the transition Spark adds at the root is a nice way to keep it small. I checked that against Spark's planning: ensureOutputsRowBased never puts a RowToColumnarExec at the root, an exchange stage is rooted at its exchange, and only convertToColumnarPlanIfPossible marks an AdaptiveSparkPlanExec columnar. So the root match only catches the cache's final stage. The batches CometSparkToColumnarExec produces use the same Arrow schema encodeBatches reads with, so they go down the direct write path as the description says.

A few things I'd like to see addressed before merge. Two are inline. The third is on a line outside the diff.

spark.comet.exec.inMemoryCache.enabled now also decides how an adaptive cache is built on Spark 4.0+, not only how cached tables are scanned and read. Its doc string in CometConf.scala only talks about the read side. Could you add a sentence about the build? configs.md is generated from that string, so that's where users will look.

[viirya]

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.

[viirya]

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.