[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.

• 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