andygrove opened a new pull request, #6607: URL: https://github.com/apache/datafusion-comet/pull/6607
## Which issue does this PR close? Closes #6598. ## Rationale for this change When a shuffle's child is a Spark row-based operator, Comet uses its JVM columnar shuffle. #6598 found that converting the child's rows with `CometSparkToColumnarExec` and using native shuffle instead is faster. This adds that behind `spark.comet.convert.shuffleInput.enabled`, off by default, next to the other `spark.comet.convert.*` settings (#6593). It reuses the `ArrowWriter` that #6566 sped up. ## What changes are included in this PR? - `CometShuffleExchangeExec.convertsInputForNativeShuffle` decides whether a shuffle takes the conversion. The config has to be on, and the shuffle's child must be a Spark row-based plan that would otherwise go to the JVM columnar shuffle. `CometSparkToColumnarExec` has to support the child's columns, and native shuffle has to support the partitioning and the columns. The checks are `shuffleSupported`'s own, so the rule converts only a shuffle that will become native. - A shuffle that hashes a decimal wider than 18 digits into more than one partition stays on the JVM columnar shuffle. Native shuffle doesn't hash those as Spark does (#5994), and the other input of a join may still be partitioned by Spark. This is the same guard as in #6564. - The conversion leaves Celeborn alone for now, since it has no JVM columnar shuffle to replace. - `CometExecRule` puts `CometSparkToColumnarExec` over the shuffle's child. As in #6564, it first applies Spark's own transition rule to the child's subtree, because Spark inserts no transitions below a `RowToColumnarTransition`. The generic fallback case no longer reports those transitions as unsupported operators when AQE runs the rule over the plan again. - `revertRedundantColumnarShuffle` (#4004) also reverts a native shuffle over converted rows when both its parent and the converted operator are Spark aggregates. That's the same row → Arrow → shuffle → Arrow → row pattern. - The shuffle tuning guide describes the setting. The new `CometShuffleInputConversionBenchmark` hash-partitions 4M rows into 200 partitions and aggregates them. It checks every arm's plan and result before timing. Best times on an M3 Max, `local[1]`, release build: | Case | Spark | Comet | Comet, converted | Comet (repeat) | | ----------------------------------------------------- | ------- | ------- | ---------------- | -------------- | | RDD rows: int, long, double | 2145 ms | 1973 ms | 1546 ms | 2014 ms | | RDD rows: int, long, string, double, decimal(18,2) | 3079 ms | 3505 ms | 2627 ms | 3167 ms | | `map` over a Parquet scan | 1936 ms | 1846 ms | 1145 ms | 1771 ms | "Comet" is the JVM columnar shuffle. The repeat row shows the noise: other jobs were running on the machine. This overlaps with #6564, which converts a typed Dataset operation's output the same way. The two share the transition handling and the fallback case for transitions, so whichever lands second has a small conflict in `CometExecRule`. ## How are these changes tested? `CometShuffleInputConversionSuite` covers: - a shuffle of RDD rows, with AQE on and off - a shuffle over a typed `map` of a Comet scan, which checks that the Spark operators under the conversion read Comet operators through Comet's transitions and that no later pass reports those transitions as fallbacks - aggregates whose partial aggregate runs on Spark - the revert between two Spark aggregates, with the revert on and off - a join with one input on native shuffle and one on the JVM columnar shuffle - a join on wide decimal keys - columns the conversion does not support - range partitioning - `spark.comet.shuffle.mode=jvm` - the default The suite turns off the conversion of leaf operators, which `CometTestBase` turns on and production leaves off. I checked that the suite fails when each piece is removed: the conversion, the wide decimal guard, the transitions, the fallback case for transitions, and the revert of the new shape. On Spark 4.1, the new suite passes along with `CometExecRuleSuite`, `CometShuffleSuite` and `DisableAQECometShuffleSuite` (192 tests). On Spark 3.5, the new suite and `CometExecRuleSuite` pass. The strict warnings build and the semantic scalafix check pass on Spark 3.5. -- 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]
