[EPIC] In-memory cache audit: correctness and performance follow-ups for InMemoryTableScanExec support #6834

What is the problem the feature request solves?

Comet's support for InMemoryTableScanExec was audited on main (bfacc43, the day #5634 made spark.comet.exec.inMemoryCache.enabled default to true on Spark 3.5+). This epic collects the correctness and performance findings so they can be worked and tracked in one place. Items are listed by priority.

Correctness

• Wrong results: a sorted cache under a native hash join loses its order with default settings. Fix: fix: keep a streamed side sorted outside the native plan in order through outer hash joins #6835. CometInMemoryTableScanExec inherits outputOrdering from the Spark scan it replaces (CometExec.outputOrdering = originalPlan.outputOrdering), so a relation cached after sortWithinPartitions reports its sort. Spark removes the sort above a left outer CometHashJoin or CometBroadcastHashJoin (BuildRight) as redundant, but the native ScanExec declares no ordering, so DataFusion appends the unmatched probe rows after the matched ones in each batch. This is the root cause described in Native hash joins lose the streamed side's order when it arrives sorted from outside the native plan #6787, which was filed against the opt-in spark.comet.convert.inMemoryCache.enabled path and left the native cache scan unchecked. Reproduced on Spark 4.1 with Native hash joins lose the streamed side's order when it arrives sorted from outside the native plan #6787's shape (big as (i % 100, i) for 10000 rows, small as (i * 10, i) for 10 rows, AQE off, 2 shuffle partitions, big.repartition(2, $"_1").sortWithinPartitions("_1").cache(), left outer join on _1, then sortWithinPartitions("k")):

  Reader	Sort in plan	Partitions sorted
  CometHashJoin over CometInMemoryTableScan	no	false, false
  CometBroadcastHashJoin over CometInMemoryTableScan	no	false, false
  Spark ShuffledHashJoin over the same Comet-format cache	no	true, true

  Fix in the joins per Native hash joins lose the streamed side's order when it arrives sorted from outside the native plan #6787, not by dropping the cache scan's ordering, which would put sorts back above every sort-merge join, window and sort aggregate over a sorted cache. The existing test "sort aggregate over a sorted cache keeps its grouping-key order" in CometInMemoryCacheSuite documents the same hazard for aggregates.

• Cached AQE queries fail after transition reversion on Spark 4.0+ #6687: transition reversion breaks a cached AQE query on Spark 4.0+ ("FilterExec has column support mismatch"). Fix: fix: preserve columnar output when reverting a cached AQE result stage #6852. Only reachable with spark.comet.exec.transitionRevert.enabled=true, which defaults to false.

• Spark's PartitionBatchPruningSuite no longer verifies pruning under Comet. Fix: test: count the partitions and batches the native cache scan reads #6836. CometInMemoryTableScanExec has none of Spark's test-only readPartitions and readBatches accumulators, so the Spark SQL diffs assert only the answer when a Comet scan is present. Spark's forty-odd pruning cases therefore no longer prove that any batch was skipped; CometInMemoryCacheSuite covers pruning in a handful of tests. Add the two accumulators to doExecuteColumnar, gated on spark.sql.inMemoryTableScanStatistics.enable as Spark's scan is, and restore the original assertions in the diffs.

Verified sound during the audit, with no action needed: canonicalization and exchange reuse across scans; statistics bounds (all-null batches, NaN, collated strings, the StartsWith exclusion); the write-side layout check and read-side node/buffer validation; dictionary decoding; the zero-column count(*) path; Kryo registrations and the plugin's install gate; delegation of unsupported types; observed metrics; leak handling on the compress and decompress failure paths. Predicates that contain subqueries sit inside the kept Spark scan, out of reach of subquery planning and reuse (#6577 covers the subqueries exposure), but the stats filter only extracts literals, so they are never evaluated there.

Performance

• Spark 3.4 and 3.5 build the cache through rows whenever AQE is on. Spark marks a cached AdaptiveSparkPlanExec columnar only from 4.0 (CachedBatchSerializer.convertToColumnarPlanIfPossible), so on 3.x InMemoryRelation takes the row path: native plan, columnar-to-row, UnsafeRow, then back to Arrow before compression. The materialization table in docs/source/user-guide/latest/in-memory-cache.md (944 ms against Spark's 3285 ms) was measured on 4.1 and does not apply to 3.x. Comet cannot change this without a Spark hook; document it under Limitations.

• Statistics gathering is a per-cell JVM loop. ArrowCachedBatchSerializer.gatherColumnStats calls isNullAt and a typed getter per row per column through the ColumnVector interface, including for nested columns that only record a null count, where an Arrow-backed vector answers from its validity bitmap. The benchmark has a case for the collector but the docs record no number for it. Take null counts from getNullCount when Arrow-backed as a first step.

• The in-memory cache stores small input batches as they arrive, so a cached repartition over many map tasks holds thousands of small batches #6950: a cache built from a shuffle read keeps one small batch per map output block. The columnar write stores each input batch as it arrives. Native filters already emit full 8192-row batches, but the shuffle reader returns each map task's block as its own batch, so repartition(50, a).cache() over 5M rows read from 200 map tasks holds 10000 batches of about 500 rows instead of 633 full ones: 34% more payload, 56% more memory and about 60% longer to decode. Coalesce small batches in encodeBatches up to spark.sql.inMemoryColumnarStorage.batchSize, where the row path cuts them.

• Every cached batch allocates a 1 MiB heap chunk, however small its payload #6951: every cached batch allocates a 1 MiB heap chunk, however small its payload. CachedBatchIpc.serialize writes into a ChunkedByteBufferOutputStream with 1 MiB chunks, which allocates a whole zeroed chunk on the first write and then copies the payload into an exactly sized buffer. Batches of 100 to 2000 rows take 13% to 57% longer to write, and the 10000-batch cache in The in-memory cache stores small input batches as they arrive, so a cached repartition over many map tasks holds thousands of small batches #6950 allocates about 10 GiB of chunks to hold 37 MiB. Size the chunk from the payload, capped at 1 MiB.

• On Spark 4.0+, an adaptive cached plan that ends in a Spark row operator is converted into Spark column vectors and then copied into Arrow #6952: on Spark 4.0+, an adaptive cached plan that ends in a Spark row operator is copied twice. AQE plans the final stage for columnar output, so Spark puts a RowToColumnarExec at its root and encodeBatches copies the resulting OnHeapColumnVectors into Arrow, where the row path writes Arrow once. This happens even when Comet runs none of the query, and the build takes 11% longer than with AQE off. Replace that conversion with CometSparkToColumnarExec.

• A cached aggregate stores each output slice with the bytes of the slices before it, inflating its size statistics and memory #6995: a cached aggregate stores each output slice with the bytes of the slices before it. CachedBatchIpc.serialize writes a sliced batch's data buffers from byte 0, so a partition whose aggregate emits K slices stores and reports about (K + 1) / 2 times its string data. 10M distinct keys at Spark's defaults report 505 MiB against 144 MiB, and a 200K-key cache in one partition reports 28.7 MiB against 2.9 MiB, which turns its broadcast join into a sort-merge join. Rebase the offsets before unloading, as fix: keep the values of sliced batches when coalescing a broadcast #6900 does for broadcasts.

• Decode and encode are JVM-only. A read copies the body from heap to off-heap, inflates zstd via zstd-jni, runs VectorLoader, wraps CometVectors and only then exports over FFI. Handing the IPC bytes to native and decoding with arrow-rs would remove that JVM work and allocation; the same holds for statistics and compression on the write side. Fits alongside the prefetch item in Adopt techniques from Spark's ArrowCachedBatchSerializer (SPARK-57268) in Comet's cache format #5487.

• Dictionary-encoded columns are stored plain, so low-cardinality strings pay full inflate and decode cost on every read. Storing dictionaries needs a per-batch schema or dictionary message.

• The serializer returns no vectorTypes, so Spark's generated columnar-to-row code reads through the ColumnVector interface rather than a concrete class.

• Cached reads that feed Spark operators are slower than Spark's own cache format #5485: Spark consumers of Comet's format are slower than Spark's own format, up to 1.6x and 2.5x on wide relations. Open and accurately documented.

Related

• Native hash joins lose the streamed side's order when it arrives sorted from outside the native plan #6787 native hash joins lose the streamed side's order
• Cached AQE queries fail after transition reversion on Spark 4.0+ #6687 cached AQE queries fail after transition reversion on Spark 4.0+
• Cached reads that feed Spark operators are slower than Spark's own cache format #5485 cached reads that feed Spark operators are slower than Spark's own cache format
• Adopt techniques from Spark's ArrowCachedBatchSerializer (SPARK-57268) in Comet's cache format #5487 techniques from Spark's ArrowCachedBatchSerializer
• feat: enable Comet's in-memory cache by default #5634 enabled the cache by default
• fix: draw the cached plan below CometInMemoryTableScan in the SQL tab and event log #6577 CometInMemoryTableScanExec.subqueries exposes the Spark scan