fix: let AQE optimize queries over Comet caches

[peterxcli]

Which issue does this PR close?

Closes #5245.

Rationale for this change

A pipeline that caches a join and then aggregates it on the same join key can perform a redundant shuffle on its first action. The join's output may already satisfy the aggregation's partitioning requirements, but that partitioning is only known after the cache is materialized.

For example, with AQE and Comet cache scans enabled and broadcast joins disabled:

from pyspark.sql import functions as F

left = spark.range(250_000).selectExpr("cast(id as string) c1")
right = spark.range(250_000).selectExpr("cast(id as string) c2")
q1 = left.join(right, left.c1 == right.c2).cache()

# First action: the cache has not been populated yet.
q2 = q1.groupBy("c1").agg(F.max("c2"))
result = q2.collect()

Before this patch, AQE does not recognize the Comet cache scan as a table cache stage, so the outer query can retain a shuffle and sorts planned before the cache was materialized. This patch lets AQE materialize the cache as a stage and replan the aggregation using its resulting partitioning. The regression test verifies that the redundant outer shuffle and sorts are removed.

What changes are included in this PR?

CometInMemoryTableScanExec implements Spark's InMemoryTableScanLike interface on Spark 3.5 and later. Shared methods expose whether the cache is materialized, its underlying RDD, and its runtime statistics. Small inheritance shims preserve compatibility with Spark 3.4, which lacks the interface.

Comet's native input traversal treats QueryStageExec as an input boundary, allowing it to consume table cache stages. Three regressions adapted from Spark's AdaptiveQueryExecSuite cover cold and warm cache access, partition preservation beside a cache stage, and statistics used for join selection. The statistics assertions use computeStats(), which also works on Spark 4.2.

How are these changes tested?

Benchmark

An exploratory run used Spark 3.5.9 on an Apple M4 with local[4]. The workload cached a join of two 250,000-row inputs, then grouped and aggregated the result. Cold timing includes query planning and cache materialization; warm timing uses the populated cache. Four JVMs ran in before/after/after/before order, with three warm-up cycles per JVM and 10 measured samples per variant and phase. All measured samples are included.

The timing comparison is inconclusive: competing builds caused substantial variation, and both variants used a debug native library. The ratios below compare observed medians.

Phase	Before median (ms)	After median (ms)	Observed speedup ratio (before / after)	Before range (ms)	After range (ms)
Cold cache	4201.2	770.4	5.45×	1231.9–17868.9	696.3–8435.3
Warm cache	127.0	100.8	1.26×	80.2–467.6	84.2–140.4

Every measured cycle returned the expected 250,000 groups and checksum of 1,388,890. Cold outer-query shuffles decreased from two to one. Cold task counts increased from 34 to 49 because the cache stage explicitly materializes the cached relation. Warm queries had one outer shuffle in both variants.

A small pilot and an interrupted million-row attempt are excluded from the table; the latter lacked a matching after-run. A release build on a quiet machine is needed to quantify the latency benefit.

Manually Testing

for:

import org.apache.spark.sql.functions._

val left = spark.range(0, 250000, 1, 16)
  .selectExpr("cast(id as string) c1")
val right = spark.range(0, 250000, 1, 16)
  .selectExpr("cast(id as string) c2")

val q1 = left.join(right, left("c1") === right("c2")).cache()
val q2 = q1.groupBy("c1").agg(max("c2"))

val result = q2.collect()

before:

Details

SortAggregate
  WholeStageCodegen (2)
    CometColumnarToRow
      InputAdapter
        CometSort
          AQEShuffleRead
            ShuffleQueryStage
              CometColumnarExchange hashpartitioning(c1#2, 16)
                SortAggregate
                  WholeStageCodegen (1)
                    CometColumnarToRow
                      InputAdapter
                        CometSort
                          CometInMemoryTableScan 

after:

Details

SortAggregate
  SortAggregate
    WholeStageCodegen (1)
      CometColumnarToRow
        InputAdapter
          TableCacheQueryStage
            CometInMemoryTableScan

[sunchao]

Correctness

Reviewed head 2c44fecfb80a3986f55ff11afdeddae379d6bc2d against base 7e1984399eb887cd13109698ee55cf2ce150f849. I found no verified P1/P2 issue in the six changed files.

Previously, replacing Spark's cache scan with CometInMemoryTableScanExec hid the InMemoryTableScanLike interface from AQE. A cold cached join therefore could not expose its final partitioning soon enough to remove the outer redundant shuffle/sorts. The PR restores that interface and makes the resulting table-cache stage an input boundary for native execution.

The three methods match the maintained Spark 3.5 and 4.0 implementations: materialization readiness comes from the cache builder, baseCacheRDD() returns the unfiltered stored batches, and runtime statistics come from the same relation. Spark retains responsibility for waiting for cache materialization and publishing complete row/size statistics. Comet continues to delegate ordering and partitioning to the original Spark scan, including its attribute remapping and final adaptive cached plan. The change preserves the existing serializer, projection, pruning, null/type handling, unsupported-schema fallback and non-AQE execution path.

The Spark 3.4 shim remains a LeafExecNode, while 3.5/4.x implement the cache interface. Canonical source comparison covered the maintained 3.5 and 4.0 branches; the required maintained 3.4 and 4.1 branches were unavailable, so those remain source-coverage limits.

Validation and remaining CI qualification

All three new regressions passed in the Linux Spark 3.5 job, Linux Spark 4.0 job, and macOS Spark 4.0 job. The cache suite had 35 passing cases plus one version-gated cancellation on 3.5, and 36 passing cases on both 4.0 jobs. The new cases check exact results, cold/warm materialization, partition alignment and nonzero compressed-cache statistics for join selection. Existing planning-laziness, empty/projection, pruning, storage and fallback cases also passed.

These jobs checked out c79fb9e487da9d77a2701922a79b926a98aa6133: I verified its parents are this base/head and its entire source tree equals the reviewed head. The native CI build succeeded. I did not run a separate local build or benchmark.

At the 2026-09-08 01:54:54 UTC snapshot, checks were 63 successful, 23 skipped and one failed. The macOS Spark 4.0 scan job reports a JVM SIGBUS during Parquet scan suites; it does not execute the cache suite. Its cause is not established, and the later skipped workflow does not resolve it. This code review does not qualify that failure as preexisting or mark the full CI run green; it still needs triage before merge.

Performance

The structural benefit is supported by the cold/warm plan assertions: AQE can use the materialized cache's partitioning to eliminate redundant work. The new interface adds no per-row decoding, copying or serialization. Spark's table-cache stage materializes the full base cache when cold and skips that materialization job when already loaded; that can increase first-touch task count even when the outer query shuffles less.

The author's matched-result experiment reports cold medians of 4,201.2 to 770.4 ms and warm medians of 127.0 to 100.8 ms, but uses a debug native build under CPU contention. Those timings support further measurement, not a qualified speedup claim. The existing cache benchmark disables AQE and pre-materializes the cache, so it cannot measure this change's cold path. The available measurements do not establish a release-build, quiet-host cold/warm AQE speedup.

Design

The two production changes form a coherent fix: expose Spark's existing cache-stage contract, then supply that stage's columnar output to the native plan. Adding only the interface would leave native input discovery unable to account for the new leaf stage. Reusing Spark's materialization and statistics machinery avoids a separate cache scheduler or duplicated completion tracking.

The regression tests exercise the important decisions directly. In particular, the join beside a table-cache stage checks that AQE does not coalesce its shuffle into an incompatible partition count, and the statistics test keeps the large compressed cache above the broadcast threshold while allowing the genuinely small relation to broadcast. I do not see a simpler design that preserves both responsibilities.

Abstraction & complexity

The small version shims isolate the API difference while keeping the implementation in one place. Treating QueryStageExec as an input boundary is appropriate because the stage owns execution of its wrapped plan; broadcast handling and shuffle direct-read detection retain their specialized paths. The change does not add a new framework or alter the cache format.

The test-only class-name check for the table-cache stage keeps the shared suite compilable on Spark 3.4, and the two stage-specific tests are explicitly gated to 3.5+. I found no actionable unnecessary abstraction or complexity in this scope.

[peterxcli]

Added spark shell and spark ui manually testing result at the bottom of PR description.

[sunchao]

LGTM, thanks @peterxcli !