Operators above a cached relation fall back to Spark after AQE materializes the table-cache stage

[andygrove]

Describe the bug

On Spark 3.5+ AQE wraps a cache scan in a TableCacheQueryStageExec, materializes it, and then re-plans the rest of the query. CometExecRule has no case for that stage, so on the re-plan the stage is a non-Comet leaf and the operators directly above it stay on Spark. The initial plan is fully native and the final plan is not. Results are correct, so the only symptom is performance.

It happens with Comet's native cache scan (spark.comet.exec.inMemoryCache.enabled=true) and with the default path, where CometSparkColumnarToColumnar converts Spark's cache scan. Both conversions are lost on the re-plan.

A filter directly on the scan survives: the filter and the scan share a logical node, so AQE reuses the physical CometFilter over the stage when it re-plans. That is why AQE SPARK-42101: cold and warm Comet cache materialization does not catch this: it asserts that the scan is native, not the aggregate above it. Extended explain does not explain it either. The partial aggregate records no fallback reason, because CometExecRule skips its handler when the child is not a CometNativeExec, and the final aggregate's reason only says that its child aggregate is not Comet.

Steps to reproduce

On main at 67803a7, default Spark 4.1 profile, AQE on, spark.comet.exec.inMemoryCache.enabled=true:

spark.range(0, 10000, 1, 4).selectExpr("id", "id % 10 AS k").createOrReplaceTempView("t")
spark.catalog.cacheTable("t")
spark.table("t").count()
spark.sql("SELECT k, count(*) FROM t GROUP BY k").collect()

Initial plan:

CometHashAggregate [k#1L, count#91L], [Final], [k#1L], [count(1)]
+- CometColumnarExchange hashpartitioning(k#1L, 4), ENSURE_REQUIREMENTS, CometColumnarShuffle
   +- CometHashAggregate [k#1L], [Partial], [k#1L], [partial_count(1)]
      +- CometInMemoryTableScan Scan In-memory table t [k#1L]

Final plan:

HashAggregate(keys=[k#1L], functions=[count(1)])
+- AQEShuffleRead coalesced
   +- ShuffleQueryStage 1
      +- Exchange hashpartitioning(k#1L, 4), ENSURE_REQUIREMENTS
         +- HashAggregate(keys=[k#1L], functions=[partial_count(1)])
            +- CometColumnarToRow
               +- TableCacheQueryStage 0
                  +- CometInMemoryTableScan Scan In-memory table t [k#1L]

Operators left on Spark in the final plan, with both tables cached (d is a 10-row table with columns k2 and name). The cold and warm cache give the same result, and with AQE off every query is fully native:

Query	Non-Comet operators after the re-plan
SELECT k, count(*) FROM t GROUP BY k	HashAggregate, Exchange, HashAggregate
SELECT name, sum(id) FROM t JOIN d ON k = k2 GROUP BY name	HashAggregate, Project, BroadcastHashJoin, BroadcastExchange
SELECT sum(id) FROM t WHERE k = 3	none

With spark.comet.exec.inMemoryCache.enabled=false the first query ends the same way, with Spark's InMemoryTableScan and a ColumnarToRow under the stage in place of the Comet nodes.

Expected behavior

The operators above the cache scan stay native after the re-plan, as they were in the initial plan.

Additional context

CometExecRule already treats Comet shuffle and broadcast stages as native inputs, for example ShuffleQueryStageExec(_, _: CometShuffleExchangeExec) goes to CometExchangeSink. A TableCacheQueryStageExec whose plan is a CometInMemoryTableScanExec could be handled the same way. The class only exists from Spark 3.5, so it needs a shim or a match by name. Spark's own cache scan would need a CometSparkToColumnarExec over the stage. I reproduced this on the default Spark 4.1 profile only. Spark 3.4 has no table-cache stages, so it should not be affected.

CometInMemoryCacheBenchmark sets spark.sql.adaptive.enabled=false, so the speedups in the in-memory cache user guide are measured without this re-plan. This matters for #5634, which turns the native cache scan on by default: under AQE, a plain aggregate or join over a cached table ends up reading Comet's format through Spark operators.

Found while auditing the in-memory cache.