Skip to content

Add AQE test coverage for CometInMemoryTableScanExec #5245

Description

@andygrove

What is the problem the feature request solves?

#5051 adds CometInMemoryTableScanExec and a Comet cache serializer (CometCachedBatch) behind spark.comet.exec.inMemoryCache.enabled. CometInMemoryCacheSuite covers the cache format, stats pruning, storage levels, and the fallback paths, but all of its tests run with AQE disabled.

AQE interacts with cached relations through machinery that the suite never exercises: TableCacheQueryStageExec materializes a cached plan as a query stage, and AQE re-optimizes using the cache's runtime statistics. Comet's scan replaces the node those mechanisms wrap, so a regression there (wrong stage materialization, wrong row-count/size statistics driving a join strategy) would not be caught today.

Describe the potential solution

Port the AQE + table-cache scenarios from Spark's AdaptiveQueryExecSuite and run them with spark.comet.exec.inMemoryCache.enabled=true, asserting the cached scan is CometInMemoryTableScanExec:

  1. SPARK-42101 — leaves a cached join cold, first-touches it through an AQE aggregation, and checks cold/warm materialization and the resulting plan rewrites.
  2. Table-cache stage in an AQE join — verifies TableCacheQueryStageExec and the surrounding shuffle behavior.
  3. SPARK-37742 — verifies AQE does not choose a join strategy from invalid cache runtime statistics.

Scenario 3 is worth particular attention: the Comet serializer reports its own sizes for CometCachedBatch, so the statistics AQE reads come from Comet code rather than Spark's DefaultCachedBatch accounting.

Additional context

Deferred from review of #5051 (#5051 (comment)). The cache path itself is off by default, so this is test coverage for an experimental feature rather than a fix for known-broken behavior.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions