fix: support AQE table-cache stages for Comet in-memory scans - #5733
Open
peterxcli wants to merge 2 commits into
Open
fix: support AQE table-cache stages for Comet in-memory scans#5733peterxcli wants to merge 2 commits into
peterxcli wants to merge 2 commits into
Conversation
peterxcli
force-pushed
the
fix/aqe-in-memory-cache
branch
from
September 6, 2026 06:20
311f463 to
2c44fec
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #5245.
Rationale for this change
AQE does not recognize
CometInMemoryTableScanExecas 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?
InMemoryTableScanLikeon Spark 3.5+ through version-specific inheritance shims, with shared implementations exposing cache materialization, the underlying cache RDD, and runtime statistics. Retain aLeafExecNodeshim for Spark 3.4, which has no table-cache stage API.QueryStageExecas a native input boundary so table-cache stages are handled alongside shuffle and broadcast stages.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?
CometInMemoryCacheSuitepassed; the three new AQE regressions also passed after consolidating the implementation into the shared class.spotless:check,scalastyle:check, andgit diff --check.benchmark
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.
TableCacheQueryStageExecin the fixed plan. The remaining shuffle computes the final scalar checksum.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.