Repository navigation
Use native shuffle for row input by converting it with CometSparkToColumnarExec #6598
Description
What is the problem the feature request solves?
When a shuffle's child is a Spark row-based operator, Comet uses its JVM columnar shuffle. It serializes each UnsafeRow into off-heap pages, sorts the rows by partition, and converts them to Arrow natively, one partition at a time. Converting those rows with CometSparkToColumnarExec instead and using native shuffle is faster.
I measured both on the same query: an RDD scan of 4M rows, repartition(200, k), then an aggregate, at local[1] with AQE off. The two arms alternated over 7 runs on a 1.1.0-SNAPSHOT release build.
| Columns | JVM columnar shuffle | CometSparkToColumnarExec + native shuffle |
|---|---|---|
| int, long, double | 1602 ms | 1354 ms (-15%) |
| int, long, string, double, decimal(18,2) | 2160-2175 ms | 1687-1759 ms (-19% to -22%) |
The second arm set spark.comet.sparkToColumnar.enabled, whose default operator list includes RDDScan. That only converts leaf operators. The build predates #6566, whose faster decimal path should widen the second gap.
CometSparkToColumnarExec converts rows more slowly than the native converter does (#6596 has the per-type numbers). So the columnar shuffle's extra 60-120 ns per row is spent elsewhere in its write path. I haven't profiled where.
Describe the potential solution
When a shuffle would otherwise use the JVM columnar shuffle, insert CometSparkToColumnarExec over its child and use native shuffle. Put it behind a config that is off by default, and keep the JVM columnar shuffle as the fallback for whatever native shuffle or CometSparkToColumnarExec can't take:
- Native shuffle must hash as Spark does, because the other side of a join may still use Spark's partitioner. It doesn't for decimals wider than 18 digits (Support Spark-compatible native hash and xxhash64 for decimals with precision >18 #5994), so those shuffles stay on the columnar shuffle, as in feat: run the operators above a typed Dataset operation natively #6564.
CometSparkToColumnarExecdeclines arrays and maps other thanarray<string>andmap<string,string>.- Spark inserts no transitions below a
RowToColumnarTransition, so the child's subtree needs its own, as in feat: run the operators above a typed Dataset operation natively #6564. - A shuffle between two Spark aggregates goes back to Spark's shuffle under Avoid JVM shuffle when parent stage will just convert back to rows #4004, and the new shape needs the same.
Additional context
This reuses the ArrowWriter that CometSparkToColumnarExec shares with the cache serializer and CometLocalTableScanExec, so #6566 applies to these shuffles too.
One difference to weigh before turning it on by default: CometArrowAllocator allocates the converted batches outside any memory pool, while the JVM columnar shuffle accounts for its pages.
#6593 proposes a spark.comet.convert.* config for each Spark-to-Arrow conversion. This would be another one.
Activity