perf: fuse Comet cache vector reads into Spark codegen

[peterxcli]

Which issue does this PR close?

Related to #5485.

Numeric cache encoding and decoding is split into #5869, stacked on this PR.

Rationale for this change

Spark can read Comet's cached Arrow vectors through a row iterator even when the consumer supports code generation. This materializes an intermediate UnsafeRow for every row before the consumer reads its fields.

What changes are included in this PR?

Feed eligible cache scans through Spark's ColumnarToRowExec, which fuses vector reads into the generated consumer. Honor the Comet enable and cache enable switches, and preserve AQE cache stages and existing columnar boundaries. For remaining row readers, generate an indexed iterator using Spark's reusable UnsafeRow writer, removing adapters and the extra copy while preserving owned variable-width values and interpreted fallback.

How are these changes tested?

Cache and iterator suites pass on Spark 3.4.3, 3.5.9, and 4.1.3, covering cold/warm AQE caches, runtime disable/re-enable of both Comet switches, codegen and interpreted paths, row ownership, nulls, nested values, sorting, joins, and batch boundaries.

Benchmark

The patch reduces cached-read time by 42–78% versus Comet main. Reading all six mixed columns takes 28% less time than vanilla Spark; reading six numeric columns still takes 15% more time.

Spark 4.1.3, JDK 21, Apple M4, 6 GiB heap, one local worker; 5M rows and six columns. Mixed uses three longs and three strings; numeric uses six longs. All queries use Spark operators with Comet native execution disabled. Vectorized cache reading is enabled for all three cases: vanilla Spark and Comet main choose row readers; the patch uses the fused columnar path.

Medians of 30 actions per cell across two fresh JVMs, with five warm-ups per query and reversed run order. Cache creation and planning are outside timing. Main: 8320ae481; measured patch: cd80194cd (reader code unchanged at 091eb0020). Timings predate the enable-switch guards. The updated harness enables Comet and its cache reader with native execution and shuffle disabled; a smoke run confirms the same reader plans and correct answers. Patch / Spark is the elapsed-time ratio; lower is better.

Schema	Columns read	Vanilla Spark (ms)	Comet main (ms)	Comet patch (ms)	Patch / Spark
Mixed	count(*)	52.85	148.42	42.64	0.81×
Mixed	1 long	64.22	198.10	66.38	1.03×
Mixed	1 string	173.64	310.53	120.80	0.70×
Mixed	3 columns	326.77	425.84	219.41	0.67×
Mixed	6 columns	503.43	624.37	363.18	0.72×
Numeric	count(*)	37.22	135.37	29.30	0.79×
Numeric	1 long	60.21	183.93	53.07	0.88×
Numeric	3 columns	94.06	244.05	101.92	1.08×
Numeric	6 columns	157.03	313.40	180.37	1.15×

A separate forced-row control takes 541 ms for six mixed columns and 226 ms for six numeric columns, versus 363 and 180 ms with the columnar path. All 1,080 measured actions, including this control, matched uncached answers. Spark six-column medians varied from 488–505 ms for mixed and 151–168 ms for numeric between JVMs. These are cached aggregate reads on one machine, not whole-application speedups.

[sunchao]

Reviewed 091eb002 against base 8b818b53. No verified P1/P2 findings.

Correctness

Spark's cache scan supports both row and columnar output, so an eligible generated consumer could previously receive an intermediate UnsafeRow for every cached row. The new rule inserts Spark's ColumnarToRowExec at that consumer edge only when the child can expose a supported Comet Arrow cache. It wraps an AQE cache stage without replacing its plan, preserving stage materialization and the scan's output, partitioning and ordering. Existing transitions and columnar consumers are left intact.

The remaining row path reads vectors into Spark's generated UnsafeRow writer, with InterpretedUnsafeProjection as the fallback. Removing the extra copy follows Spark's existing cache-reader contract: successive next() calls reuse the row, while callers retaining rows must copy them. The writer owns variable-width values, so advancing or exhausting the upstream iterator can release its Arrow batch without invalidating the last returned row. The tests exercise that boundary with real Arrow allocation, repeated hasNext(), nulls and long UTF-8 values, and separately cover empty batches, zero-column rows and wide projections in both factory modes.

I verified the new cache and iterator tests passing in the Spark 3.5, Spark 4.0 and Spark 4.1 exec jobs. Those jobs checked out 46327db9, whose parents are exactly the assigned base and head and whose entire tree equals the reviewed head. Their suites report 805, 851 and 852 passed tests with zero failures; respectively 7, 3 and 0 tests were canceled, and each job ignored 5. At 2026-09-12 20:40:44 UTC, the head has 73 successful checks and 8 skipped checks. I did not run a separate local product build. Direct semantic comparison used the available maintained Spark 3.5/4.0 sources; maintained 3.4/4.1 sources were unavailable.

Performance

The main improvement removes intermediate row materialization when Spark can fuse the cache-vector reads into its generated consumer. The remaining row path also avoids the per-row adapter chain and final copy, while retaining Spark's projection and null handling. Compiler setup is per iterator construction, not per row; batch binding and indexed reads remain straightforward.

The benchmark validates uncached answers, cache residency, format, selected columns and the actual reader plan before timing. Its fresh-JVM comparison and forced-row control support separating the fused path from the row-iterator improvement. I verified that the reported measured commit cd80194c has identical changed reader code at this head apart from the added diagram comment. The reported 42–78% improvement over Comet main is still author-measured: I did not independently reproduce the timings or obtain the raw samples. The six-column numeric case remains about 15% slower than vanilla Spark in that report, so these results support cached-read improvements rather than a universal Spark speedup.

Design

Applying the cache rule after Comet's other post-transition rewrites lets it inspect the final consumer boundary. Eligibility follows Spark's whole-stage support, field-count and expression checks and uses the cache serializer and supported schema to identify the physical data. This also handles a Comet cache read after native execution is disabled. Existing serializer fallback for unsupported schemas remains in place.

The regression coverage checks both planning and execution: planning a cold cache must not materialize it, execution must populate it, AQE must retain the cache stage, and existing row/column boundaries must remain stable. Row-consumer tests cover reordered attributes, nested and nullable data, sorting, joins, count and limit across small batches. Registering the iterator suite in both platform workflows keeps the ownership and code-generation checks in normal CI.

Abstraction & complexity

The implementation uses one focused planning rule and one iterator factory. It reuses Spark's ColumnarToRowExec, vector-access generation, UnsafeRow projection and interpreted fallback instead of adding another expression evaluator or buffer format. The upstream decoder continues to own and close the batches, keeping allocation responsibility in one place.

The two paths have clear roles: fused vector reads for eligible generated consumers, and a reusable row projection for other consumers. Their contracts and the diagram explain the boundary without adding configuration or a new extension interface. I found no blocking simplification or additional abstraction needed.

[andygrove]

CometCacheColumnarRule rewrites the plan without checking whether Comet is enabled. CometScanRule and CometExecRule both start with if (!isCometLoaded(conf)) return plan, so setting spark.comet.enabled=false normally takes Comet out of physical planning completely. This rule only checks conf.wholeStageEnabled and then fires on any relation whose serializer is ArrowCachedBatchSerializer. Since spark.sql.cache.serializer is static and installed at startup by CometDriverPlugin, that stays true for the life of the application. The new test in CometInMemoryCacheSuite runs with spark.comet.enabled=false and still expects the rewrite, so this looks deliberate, but it does mean Comet changes the plan after a user has switched it off.

The same goes for spark.comet.exec.inMemoryCache.enabled, whose doc string says that disabling it at runtime only sends cached scans back to Spark's execution path. A user who hits a problem on the fused path has no way to turn it off short of spark.sql.inMemoryColumnarStorage.enableVectorizedReader=false, which changes a lot more than this. Would it make sense to gate the rule on isCometLoaded(conf) and on COMET_EXEC_IN_MEMORY_CACHE_ENABLED? The CachedBatchRowIterator improvement applies either way, so falling back to it when Comet is disabled still leaves those users better off than today.

While you are in there, the doc string for COMET_EXEC_IN_MEMORY_CACHE_ENABLED says that reads feeding Spark operators still pay a row conversion the default format avoids. That is what this PR removes for eligible codegen consumers, and configs.md is generated from that string, so the published guidance goes stale the moment this merges. Could you update it to describe the current split, including the numeric case where Comet's cache is still behind Spark's?

The rest of it holds up well. The premise is right that InMemoryTableScanExec.supportsRowBased is true so Spark deliberately does not insert the transition itself, getValueFromVector and ColumnarBatchRow bottom out in the same accessors so there is no codegen versus interpreted divergence, and dropping the .copy() matches what DefaultCachedBatchSerializer already does.

[andygrove]

Thanks for adding the enable-switch guards and the runtime toggle test, that addresses my earlier comment. One more gate question. CollapseCodegenStages.apply only inserts whole-stage codegen when spark.sql.codegen.factoryMode is not NO_CODEGEN as well as wholeStageEnabled. Should this rule check the same thing? Otherwise under NO_CODEGEN with whole-stage on we insert a ColumnarToRowExec that never fuses and runs its plain doExecute. The existing tests always pair NO_CODEGEN with whole-stage off, so it might be worth adding that combination once the gate matches.

[andygrove]

The generated code hard-codes references[0] as the batch iterator while ctx.references starts empty. That holds today because nothing in GenerateUnsafeProjection.createCode for BoundReferences adds a reference, but if that ever changes index 0 would silently become something else. Would you consider registering the iterator with ctx.addReferenceObj("batches", batches) and passing ctx.references.toArray to generate, the way Spark's own generators do?

[andygrove]

This asserts the row reader for spark and comet-row but never asserts the columnar reader for comet. If the rule stops firing for any reason, the comet column of the results table would quietly measure the row path. Could you add assert(columnar) for the comet format so the benchmark fails loudly instead?