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 Comet and codegen switches (including NO_CODEGEN with whole-stage enabled), 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.

[sunchao]

Re-reviewed 091eb002 → eafdba5e. The Comet/cache enable-switch guards and runtime toggle test address the earlier feedback. The cache documentation and benchmark configuration now reflect those switches.

The existing NO_CODEGEN gate, reference registration, and positive reader assertion follow-ups remain nonblocking (P3): the unfused conversion still returns valid UnsafeRows, the current BoundReference-only projection adds no external references, and the benchmark logs its actual reader. No new or remaining verified P1/P2 issue found.

The cache and iterator suites passed in the Spark 3.5, 4.0, and 4.1 exec jobs. CI tested merge 975c72c1; all eight authored Scala files match this head. The current merge preview is distinct. No local benchmark rerun; maintained Spark 3.4/4.1 source branches remain unavailable.

[peterxcli]

@andygrove addressed review! ptal again, thanks!

[sunchao]

Re-reviewed eafdba5e → ee0e24f3. The three earlier follow-ups are addressed: the cache rule now honors NO_CODEGEN, the generated iterator uses Spark's reference registry, and the benchmark requires the fused columnar reader. The expanded runtime test covers disabling and re-enabling each codegen setting on the same cache with AQE on and off. No new or remaining verified P1/P2 issue found.

The expanded cache test and all six iterator cases passed in the Spark 3.4, 3.5, 4.0, and 4.1 exec jobs. Their checkout f21e504a has the entire tree of this head. The final commit only retries CI. That earlier run's red jobs failed during dependency downloads or never acquired runners.

At 2026-09-15 15:09 UTC, current-head CI still has 15 pending checks, alongside 40 successful and 9 skipped checks. No local test or benchmark rerun. Performance numbers remain historical author measurements. Direct source comparison used maintained Spark 3.5/4.0 branches; maintained 3.4/4.1 sources remain unavailable.

[sunchao]

CI follow-up for ee0e24f3 against fad62309: the reviewed source is unchanged, and I found no new P1/P2 issue. My existing approval stands. The current run has 16 canceled jobs. Required Checks failed because its upstream groups were canceled. The cancellation cause is not established.

The Spark 4.1 exec log reports 897 passed, 0 failed and 5 ignored tests, including the cache runtime-toggle cases and all six iterator cases, before the job was canceled during later Maven packaging. Its actual checkout 3fe8418e has the assigned base/head as parents and the same entire tree as the reviewed head. Its downloaded native artifact matches the producer's archive digest. This is completed test-phase evidence, while the job and required CI remain incomplete. Earlier identical-tree exec results remain prior evidence. No local product test or benchmark was added. Direct source comparison remains limited to the available maintained Spark 3.5/4.0 branches.

[sunchao]

Re-reviewed ee0e24f3 → 14b5d7b6 against 8c229a70. The main-branch merge preserves all eight authored Scala files and the same iterator-suite registrations. The codegen gate, reference registration and positive benchmark reader assertion remain addressed. No new or remaining verified P1/P2 findings. My existing approval stands.

The native build and Spark 4.1 build passed on a5c6e99e, whose entire tree equals the reviewed head. The Spark build skipped tests. At 09:30 UTC on September 16, CI had 19 successful checks, 13 skipped and three still running, including the Spark 4.1 exec suite. The previous canceled run's 897 passing tests remain historical evidence. No local runtime or benchmark was rerun. Performance numbers remain author measurements, and maintained Spark 3.4/4.1 source gaps remain recorded.

[sunchao]

@peterxcli could you rebase this PR?

[sunchao]

Re-reviewed cea2caae after the merge from main. The authored cache-reader changes are unchanged: the rule, iterator, benchmark, and iterator tests are byte-identical to the prior review, and the remaining differences are inherited from main. I checked the integration with the new duplicate-field schema fallback; the rule and serializer continue to decline unsupported cache schemas.

The earlier NO_CODEGEN, generated-reference initialization, and benchmark-plan assertion findings remain fixed. No new or remaining verified P1/P2 findings.

Current Comet CI is successful. The Spark 4.1 exec job ran 982 tests successfully, including the cache consumer regressions and all six CODEGEN_ONLY/NO_CODEGEN iterator cases. Its actual checkout was merge commit 4af18e46, whose full tree matches this head. Five ignored tests and skipped jobs are not counted as executed.

No local runtime suite or benchmark was run; the PR's timing claims remain historical. Source semantics were checked against the maintained Spark 3.5/4.0 branches; the maintained 3.4/4.1 branches were unavailable.

[sunchao]

@peterxcli sorry can you rebase it again :(

[sunchao]

Summary

• Prior state and problem: Spark consumers materialized intermediate rows when reading Comet caches. The existing row reader also copied every projected row.
• Design approach: Insert ColumnarToRowExec at eligible generated-consumer boundaries and use an indexed, reusable row iterator elsewhere.
• Correctness / compatibility analysis: Found one introduced P2 issue affecting wide cache projections, detailed below. Source comparisons covered Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and experimental 4.2.0. The enable switches, AQE stage preservation and row ownership otherwise follow the relevant Spark contracts. Earlier review concerns are addressed.
• Key design decisions: Reusing Spark's transitions, vector accessors and row writers keeps the implementation focused. However, supplying ctx.currentVars disables Spark's projection splitting and introduces repeated compilation failures for wide schemas.
• Implementation sketch: One post-transition planning rule, one iterator factory, serializer wiring, configuration documentation, integration tests and a benchmark. Reviewed all 11 files in the full base-relative diff.
• Behavioral changes worth calling out: Eligible consumers read vectors directly. Other consumers receive reusable rows whose retained values require copying. The reported benchmark improvements remain author measurements and were not independently rerun.
• Suggested improvements: Preserve generated-reader scalability for wide projections and add coverage beyond the current 150-column test, as described in the finding.

Reviewed dbfb182d3e2dd426b9bc65764cac7732847b834e against a86c9672a63c909f0fd7b752c86b5a679df4e84b. The PR is not a draft. Routed skills: review-comet-pr, review-comet-expression-pr and review-comet-ffi-pr.

Exact-head CI: 23 successful checks and 14 skipped, with no failures. The Spark 4.1 exec job passed 1,077 tests, including both new cache tests and all six iterator cases. Its checkout, fe7a310c61635709c1b4fa2a5e51d8ad530867fd, has the requested base/head as parents and the same complete tree as the reviewed head.

Validation limits: Compiled the exact iterator source locally against Spark 4.1.3 on JDK 21 and ran bounded comparison probes. No full local Comet build or Spark SQL compatibility suite was run. Spark SQL suites and macOS CI were skipped, and runtime validation of other Spark profiles was not repeated. The local probes isolate reader generation and execution rather than exercising an entire cached query.

[sunchao]

[P2] Preserve code splitting for wide cache projections. Reading 1,500 selected INT columns with alternating nullable/non-nullable attributes makes the generated next() exceed the JVM's 64 KB method limit. Setting ctx.currentVars above causes GenerateUnsafeProjection to skip its normal splitExpressions path, whereas the previous UnsafeProjection.create(...) reader compiles successfully for the same schema. Under default FALLBACK, each partition retries the failed compilation and then uses the interpreted reader. The local probe measured later constructions at 400–515 ms versus 43–63 ms for the previous path, with correct results from both. CODEGEN_ONLY fails outright. Could this split the generated writer into bounded methods, or select the existing generated UnsafeProjection path for wide schemas before attempting oversized compilation? A 1,500-column regression case would cover this.

Evidence: Compiled the unchanged exact-head CachedBatchRowIterator.scala against Spark 4.1.3 with Scala 2.13.17/JDK 21. A package-local harness created attributes using (0 until 1500).map(i => AttributeReference(s"c$i", IntegerType, nullable = i % 2 == 0)()) and matching OnHeapColumnVectors. With CODEGEN_ONLY, the previous UnsafeProjection.create(attrs, attrs) path passed, while new CachedBatchRowIterator(attrs).createObject(Iterator.single(batch)) failed with InternalCompilerException: Code grows beyond 64 KB while compiling next(). Widths 150, 500 and 1,000 passed. In a separate default-FALLBACK probe with 4,096 rows and logging disabled, five alternating old/new runs returned identical checksums. The final three construction times were 62.95/45.58/43.28 ms for the previous path and 399.82/484.91/514.83 ms for the new interpreted fallback. Spark's GenerateUnsafeProjection.writeExpressionsToBuffer explicitly bypasses splitting when ctx.currentVars != null. Harnesses and logs are under /tmp/comet-5859-dbfb-probe/.

[andygrove]

This starts well below the 64 KB limit. Under it the method still compiles, but once next() passes HotSpot's 8000-byte HugeMethodLimit the JIT never compiles it and it stays interpreted. I timed this iterator against main's UnsafeProjection plus copy() over nullable bigint columns, 8 batches of 4096 rows each. It is 0.39x at 6 columns and about the same as main at 50. It is 1.7x slower at 100 columns and 13 to 15x slower from 120 to 500. String columns reach 4x at 150. With -XX:-DontCompileHugeMethods the 120 and 150 column cases drop back to about 1.2x, which points at the method size. End to end, reading a 150-column Comet cache took 320 ms on this PR against 66 ms on the merge base with Comet on and exec off, and 323 ms against 64 ms with Comet off. Spark reads of relations wider than spark.sql.codegen.maxFields always take this path, because InMemoryTableScanExec.supportsColumnar is false for them.

So falling back only when compilation fails would not be enough. CodeGenerator.compile returns the ByteCodeStats that line 110 discards, and that is what WholeStageCodegenExec checks before it backs off. A width bound would work too. For the fallback itself, UnsafeProjection.create(fields) over batch.getRow(i) in an indexed loop, without the copy(), came in at 0.64 to 0.93x of main at every width I tried. Could the benchmark and the tests cover 100 and 200 columns as well as 1,500?

[andygrove]

I re-read this against current main, which now includes #5394, #5543, #6082 and #6208. CometCacheColumnarRule holds up. Its eligibility check matches CollapseCodegenStages.supportCodegen on 3.5 and 4.1, and the table-cache stage path fuses on both cold and warm reads. The comments below are about wide reads through the row iterator, where I've added numbers to sunchao's thread, two tests that check less than they claim, the benchmark, and the user guide.

[andygrove]

This test still sets spark.comet.enabled=false, and since the rule now checks isCometLoaded, none of the four combinations take the fused path. The vectorized && CODEGEN_ONLY case runs the row iterator like the others, and the test only asserts that ColumnarToRowExec is absent, so nothing notices. That leaves sum(key) and sum(length(s)) in the next test as the only fused-path coverage. Could this run with Comet on and exec and shuffle off, like the next test, and assert that the transition is present for vectorized && CODEGEN_ONLY? A consumer that reads every column through codegen, such as df.filter($"key" >= 0), would cover the scalar types as well. I tried that locally and the whole type matrix passes through the fused path with AQE on and off.

[andygrove]

checkAnswer first runs df.materializedRdd.count() on a separate query execution and only then df.collect(). So the cold iteration's cache is already loaded when the plan asserted here runs, and the path where the table-cache stage materializes and AQE re-plans is never the one checked. The #6208 test earlier in this file uses QueryTest.checkAnswer(df, expected, checkToRDD = false) for this reason. Could you do the same here? I tried it locally, and the cold run is then really cold and still gets one fused transition.

[andygrove]

This session turns Comet on but never enables on-heap or off-heap memory for it, so isCometLoaded returns false and the rule never fires. Without ENABLE_COMET_ONHEAP=true in the environment, which make benchmark-... does not set, the comet arm fails its own Expected the fused columnar cache reader assertion. CometBenchmarkBase sets spark.comet.exec.onHeap.enabled for this reason. Could you set it here too?

[andygrove]

#5543 added a Spark-operator arm to CometInMemoryCacheBenchmark after this PR was opened. Would it make sense to add the fused case there instead of a second harness? It already extends CometBenchmarkBase, and it is the benchmark the user guide's Limitations table cites. Its Spark-operator arm runs with Comet off, so after this PR it only measures the row iterator. The numbers in the description also predate #5543's storage format change, so could you re-run them on current main?

[andygrove]

This describes the new split, but the Limitations section of docs/source/user-guide/latest/in-memory-cache.md still says Spark-operator reads are slower than Spark's own format and that the cause is not yet established. Its table was measured with Comet off, and those reads now go through CachedBatchRowIterator. With Comet on and exec off, eligible consumers take the fused path instead. Could you update that section in this PR, including when the fused path applies? With exec on, a Spark operator above the cache already reads through CometColumnarToRow over the native scan, so that case does not change.

[andygrove]

Since #5394, CometColumnar.postColumnarTransitions runs these rules in plan-only mode too, and this rule rewrites plans that contain no Comet operators. With spark.comet.explain.planOnly.enabled=true, a read of a Comet-format cache executes a ColumnarToRowExec that Spark would not have planned, while the plan-only doc says Spark executes the query unchanged. Should this also return early when plan-only is on? A plan-only case in the runtime-settings test would pin it down.

[sunchao]

Summary

• Prior state and problem: Spark consumers materialized intermediate rows when reading Comet caches, and the existing row reader copied each projected row.
• Design approach: Insert Spark’s ColumnarToRowExec at eligible generated-consumer boundaries and use an indexed, reusable row iterator elsewhere.
• Correctness / compatibility analysis: Reviewed all 11 files in the full base-relative diff and compared relevant Spark sources across 3.4.3, 3.5.9, 4.0.4, 4.1.3 and experimental 4.2.0. No additional introduced P1/P2 issues found within this review beyond the existing unresolved wide-projection P2. Independently reproduced it: the previous reader succeeds with 1,500 integer columns, while this head’s generated next() exceeds the JVM’s 64 KB method limit under CODEGEN_ONLY.
• Key design decisions: Reusing Spark’s transition, vector accessors and unsafe-row writers keeps the implementation focused. However, setting ctx.currentVars bypasses Spark’s projection splitting. Under default FALLBACK, the local probe confirmed repeated failed compilation followed by interpreted execution, with substantially higher construction overhead.
• Implementation sketch: One post-transition planning rule, one iterator factory, serializer wiring, configuration documentation, integration tests, a benchmark and registration of the iterator suite in both platform workflows.
• Behavioral changes worth calling out: Eligible consumers read vectors directly. Remaining row consumers receive reusable rows and must copy retained rows. The generated writer owns variable-width values across batch release. The PR’s application-level speedup claims remain historical author measurements.
• Suggested improvements: Address the existing wide-reader thread by preserving bounded generated methods or choosing a scalable projection path before oversized compilation. No additional review comments are proposed.

Reviewed dbfb182d3e2dd426b9bc65764cac7732847b834e against a86c9672a63c909f0fd7b752c86b5a679df4e84b. The PR is not a draft. Routed skills: review-comet-pr, review-comet-expression-pr and review-comet-ffi-pr.

Exact-head CI: 23 successful checks and 14 skipped, with no failed checks. The Spark 4.1 exec job passed 1,077 tests, including both added cache tests and all six iterator cases. Its checkout, fe7a310c61635709c1b4fa2a5e51d8ad530867fd, has the requested base/head as parents and the same complete tree as the reviewed head.

Validation limits: Compiled the unchanged iterator locally against Spark 4.1.3, Scala 2.13.17 and JDK 21, then ran bounded old/new comparisons in CODEGEN_ONLY and FALLBACK. These isolate reader behavior rather than entire cached queries. No full local Comet build, Spark SQL compatibility suite or application benchmark was run. Exact-head Spark SQL suites and macOS CI were skipped, and local runtime validation covered only Spark 4.1.3. Project code was unchanged and nothing was published.