Skip to content

fix: support AQE table-cache stages for Comet in-memory scans - #5733

Open
peterxcli wants to merge 2 commits into
apache:mainfrom
peterxcli:fix/aqe-in-memory-cache
Open

fix: support AQE table-cache stages for Comet in-memory scans#5733
peterxcli wants to merge 2 commits into
apache:mainfrom
peterxcli:fix/aqe-in-memory-cache

Conversation

@peterxcli

@peterxcli peterxcli commented Sep 6, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5245.

Rationale for this change

AQE does not recognize CometInMemoryTableScanExec as a table-cache stage. On first access to a cached join, this prevents AQE from using the materialized cache's partitioning to remove redundant shuffles and sorts. Recognizing the cache stage also requires Comet's native input traversal to stop at that stage.

What changes are included in this PR?

  • Implement InMemoryTableScanLike on Spark 3.5+ through version-specific inheritance shims, with shared implementations exposing cache materialization, the underlying cache RDD, and runtime statistics. Retain a LeafExecNode shim for Spark 3.4, which has no table-cache stage API.
  • Treat QueryStageExec as a native input boundary so table-cache stages are handled alongside shuffle and broadcast stages.
  • Port three AQE regressions from Spark's AdaptiveQueryExecSuite, with source links: cold/warm cache materialization and shuffle/sort elimination, partition preservation beside a table-cache stage, and cache statistics used for join selection.

How are these changes tested?

  • Spark 4.1: all 36 tests in CometInMemoryCacheSuite passed; the three new AQE regressions also passed after consolidating the implementation into the shared class.
  • Spark 3.5: all three new AQE regressions passed.
  • Spark 3.4 / JDK 17: the cache-statistics regression passed; the two table-cache-stage tests are explicitly skipped because Spark 3.4 lacks that API.
  • spotless:check, scalastyle:check, and git diff --check.

benchmark

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

Ratios compare observed medians; CPU contention and the debug native build prevent a reliable production speedup claim.

The before blocks' cold medians varied from 1391.5 ms to 14207.9 ms, illustrating the contention problem.

  • All measured cycles returned 250,000 groups and a string-length checksum of 1,388,890.
  • Cold outer-query shuffles decreased from 2 to 1, with a TableCacheQueryStageExec in the fixed plan. The remaining shuffle computes the final scalar checksum.
  • Warm outer-query shuffles remained 1 in both variants; cache decoding is unchanged.
  • Cold task counts increased from 34 to 49 because the cache stage explicitly materializes the cached relation. Removing a shuffle alone therefore does not guarantee lower latency.

A 100,000-row pilot and an interrupted 1,000,000-row attempt are excluded from these aggregates; the latter had no matched after-run. A release-native rerun on a quiet machine is needed to quantify the latency benefit.

@peterxcli
peterxcli force-pushed the fix/aqe-in-memory-cache branch from 311f463 to 2c44fec Compare September 6, 2026 06:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add AQE test coverage for CometInMemoryTableScanExec

1 participant