perf: the bypass JVM columnar shuffle writer writes one-row blocks when its partitions' pages don't fit in a task's memory #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):

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). 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). 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). 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):

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

// 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 c230bf4 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.