andygrove opened a new pull request, #6208: URL: https://github.com/apache/datafusion-comet/pull/6208
## Which issue does this PR close? Closes #6202. ## Rationale for this change On Spark 3.5+, AQE wraps each cache scan in a `TableCacheQueryStageExec` and, once the stage materializes, plans the operators above it again. `CometExecRule` had no case for the stage, so those operators stayed on Spark. For `SELECT k, count(*) FROM t GROUP BY k` over a cached `t`, the initial plan is fully native, but the final plan runs both aggregates and the exchange on Spark over a `CometColumnarToRow`. Results are correct, so the cost is performance, and it becomes more visible once #5634 turns the native cache scan on by default. ## What changes are included in this PR? Two cases in `CometExecRule`, one per cache format: - A query stage whose plan is a `CometInMemoryTableScanExec` is a native input through `CometExchangeSink`, as a Comet shuffle or broadcast stage already is. The scan produces Arrow batches, and `foreachUntilCometInput` already accepts any `QueryStageExec` as an input. - Spark's cache scan does not come back from the re-plan as a bare stage. When the `CometScanWrapper` around a `CometSparkToColumnarExec` is removed, the `CometSparkToColumnarExec` inherits the scan's logical link, so AQE reuses it over the stage as the relation's physical plan. It is not a `CometNativeExec`, so the operators the re-plan puts above it could not convert. An existing `CometSparkToColumnarExec` is now wrapped in a `CometScanWrapper` again, the way `isCometScan` nodes already are on every pass. Neither case applies with default settings, since the native cache scan and `spark.comet.sparkToColumnar.enabled` are both off. ## How are these changes tested? A new test in `CometInMemoryCacheSuite` runs the three queries from the issue (an aggregate, a broadcast join and a filtered aggregate) with the native cache on and off, over a cold and a warm cache, and checks that the executed final plan has no Spark operators. It inspects the plan after `checkAnswer` rather than using `checkSparkAnswerAndOperator`, which inspects an adaptive plan that has not run, so it only ever sees the initial plan. That suite registers Comet's extensions twice (through `spark.plugins` and `withExtensions`), so there the re-plan brings back two stacked conversions. A second test in `CometExecSuite` covers a single registration with Spark's default cache format. Both fail without the change. Locally, the cache suites pass on Spark 3.4 (where the new tests are skipped), 3.5, 4.0, 4.1 and 4.2, and `CometExecSuite`, `CometAggregateSuite`, `CometJoinSuite`, the shuffle suites and `CometExecRuleSuite` pass on 4.1. -- 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]
