andygrove opened a new issue, #7058: URL: https://github.com/apache/datafusion-comet/issues/7058
### Describe the bug When the JVM columnar shuffle picks its bypass-merge writer and the shuffle has more partitions than a task has room for full memory pages, it writes Arrow IPC blocks of about one row each. The shuffle then gets 12x larger and around 190x slower than Spark's. `SELECT /*+ REPARTITION(64) */ * FROM orders` (TPC-H SF2, 3M rows, noop sink, `local[8]`, 6 GB off-heap): | Configuration | Time | Shuffle bytes written | |---|---:|---:| | Comet, default local mode (bypass writer) | 266 s | 2.48 GB (~830 bytes per row) | | Comet, `spark.executor.cores=8` (sort-based writer instead) | 1.7 s | 156 MB | | Comet, `spark.buffer.pageSize=1m` (bypass writer, 1 MB pages) | 2.3 s | 156 MB | | Spark | 1.4 s | 205 MB | The same query on `lineitem` (12M rows) took 15.5 minutes. The reduce side is slow too, because it decodes the one-row blocks one JNI call at a time. ### Cause The bypass writer creates one `CometDiskBlockWriter` per partition. Each buffers its rows in pages from a shared allocator whose page size is Spark's `pageSizeBytes`, capped at 128 MB ([CometBypassMergeSortShuffleWriter.java#L173-L177](https://github.com/apache/datafusion-comet/blob/c230bf4e986ecfaf317c0621b13cd84931c948ae/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java#L173-L177)). `SpillWriter.acquireNewPageIfNecessary` asks for `required` bytes, but Spark's `MemoryConsumer.allocatePage` hands out `max(pageSize, required)`, so every partition holds a whole page (32 MB in the run above) as soon as it has one row. Spark sizes `pageSizeBytes` as `nextPowerOf2(executionPool / cores / 16)`, capped at 64 MB, so a task's share holds only 16 to 32 pages unless the cap binds. With more partitions than that, a writer that needs a page finds the pool full. `ArrowIPCWriter.spill` then spills the largest sibling, just enough to cover `required` ([CometDiskBlockWriter.java#L365-L402](https://github.com/apache/datafusion-comet/blob/c230bf4e986ecfaf317c0621b13cd84931c948ae/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java#L365-L402)). Under round-robin that sibling has buffered about one row since it got its page, so nearly every row becomes its own block. Each block costs a `writeSortedFileNative` call (builders, file open and close, codec setup) and about 800 bytes of IPC framing. ### Who hits it `CometShuffleManager.shouldBypassMergeSort` picks the bypass writer when `numPartitions <= spark.shuffle.sort.bypassMergeThreshold` (200) and `numPartitions * spark.executor.cores <= spark.comet.shuffle.jvm.maxWritersPerExecutor` (100) ([CometShuffleManager.scala#L302-L325](https://github.com/apache/datafusion-comet/blob/c230bf4e986ecfaf317c0621b13cd84931c948ae/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala#L302-L325)). `spark.executor.cores` reads as 1 when it is not set, so this covers: - local mode and standalone clusters without explicit executor cores: any JVM columnar shuffle with up to 100 partitions; - 1-core executors (the YARN and Kubernetes default): up to 100 partitions; - 2-core executors: up to 50 partitions. The JVM columnar shuffle is what `df.repartition(n)` uses by default, since `spark.comet.shuffle.native.partitioning.roundrobin.enabled` is false. It is also used under a Spark operator Comet does not convert. A 2-core executor, emulated with `local[2]`, 2 GB off-heap and `spark.executor.cores=2` on 300K `orders` rows, falls off the cliff where the page math predicts (32 pages per task here): | `REPARTITION(n)` | Comet | Shuffle bytes written | |---:|---:|---:| | 24 | 1.6 s (first run, includes JIT warmup) | 15.6 MB | | 32 | 0.46 s | 15.8 MB | | 40 | 8.5 s | 125 MB | | 48 | 16.7 s | 230 MB | | Spark, 48 | 0.9 s | 20.7 MB | Below the cliff each partition still pins a full page, which the task's other consumers can't use. The writer's allocator can't be asked to spill (`CometUnifiedShuffleMemoryAllocator.spill` returns 0). ### Steps to reproduce ```scala // local[8], spark.memory.offHeap.enabled=true, spark.memory.offHeap.size=6g, // spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager, // spark.executor.cores not set spark.read.parquet("orders.parquet").createOrReplaceTempView("orders") // TPC-H SF2 spark.sql("SELECT /*+ REPARTITION(64) */ * FROM orders").write.format("noop").mode("overwrite").save() ``` Compare the stage's shuffle bytes written with the same query under `spark.comet.enabled=false`. ### Expected behavior The bypass writer's shuffle size and time stay close to the sort-based writer's for any partition count it accepts, and a partition's buffer costs memory in proportion to the rows it holds. ### Additional context Measured on `main` at c230bf4e98 with a release build on an Apple M3 Max, Spark 4.1.3, JDK 17. The 190x includes macOS-specific costs (malloc's `madvise` on free, the W^X toggle on every JNI call). The 12x shuffle size and the one-row blocks don't depend on the platform. The behavior is old, not a 1.1 regression: the bypass writer and its executor-cores check date from #30, and the page-sized unified allocator from #1485. A fix that keeps the bypass writer: start each partition's pages small and grow them geometrically up to the current page size. One page per partition then costs kilobytes, and the page count per writer stays logarithmic. That matters because Spark fails the task, rather than spilling, once a task holds 8,192 pages. -- 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]
