andygrove opened a new issue, #6202: URL: https://github.com/apache/datafusion-comet/issues/6202
### 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`](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala#L123-L154) 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 67803a7a4, default Spark 4.1 profile, AQE on, `spark.comet.exec.inMemoryCache.enabled=true`: ```scala 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)`](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala#L495-L497) 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`](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala#L102), 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
