andygrove opened a new pull request, #6852: URL: https://github.com/apache/datafusion-comet/pull/6852
## Which issue does this PR close? Closes #6687. ## Rationale for this change `RevertNativeForTransitionHeavyStages` assumed a result stage is always collected as rows and passed `outputColumnar = false` for it. That does not hold for a cached query on Spark 4.0+. When the cache serializer accepts columnar input, as Comet's `ArrowCachedBatchSerializer` does, Spark marks the cached `AdaptiveSparkPlanExec` columnar, AQE plans the final stage with `outputsColumnar = true`, and the cache builder calls `executeColumnar()` on it. Once the rule reverted that stage to a row-based root such as `FilterExec`, nothing restored the columnar output, and building the cache failed with `FilterExec has column support mismatch`. Transition reversion is off by default, so this only affects users who enable `spark.comet.exec.transitionRevert.enabled`. ## What changes are included in this PR? - Both result-stage call sites go through a new `revertResultStageIfNeeded`, which reads the requested format off the stage root with `outputColumnar = plan.supportsColumnar && !plan.supportsRowBased`, as proposed in the issue. Spark's transition insertion leaves a columnar-only root only when columnar output was requested; when rows are wanted it puts a columnar-to-row transition on top. When the format is columnar, the existing logic in `revertStageIfNeeded` wraps a row-based reverted root in `RowToColumnarExec`, which is what Spark itself inserts in that case. - Without AQE, Spark always plans for row output, so that path behaves as before. ## How are these changes tested? A new test in `RevertNativeForTransitionHeavyStagesSuite` reproduces the issue. With AQE on, `transitionRevert.enabled=true`, `maxTransitions=0` and `spark.comet.exec.filter.enabled=false`, it caches `SELECT _2, s FROM (SELECT _2, sum(_1) AS s FROM tbl GROUP BY _2) WHERE s > 10` and checks the answer against Spark's. It then checks that the cache was built from columnar output, that the result stage was reverted, and that the reverted stage is still columnar. On `main` the test fails with `FilterExec has column support mismatch`. So that the cache is built from columnar output, the suite now installs `ArrowCachedBatchSerializer` the way `CometInMemoryCacheSuite` does, clearing Spark's memoized serializer before and after the suite. The test runs on Spark 4.0+ and is skipped on 3.x, where AQE builds a cache through rows. `RevertNativeForTransitionHeavyStagesSuite` passes locally on Spark 3.4, 3.5, 4.0, 4.1 and 4.2. On Spark 4.1, `CometShuffleInputConversionSuite`, `CometExecRuleSuite`, `CometCelebornShufflePlanningSuite` and `CometParquetWriterSuite`, which also exercise the rule, pass as well. -- 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]
