andygrove opened a new pull request, #7059: URL: https://github.com/apache/datafusion-comet/pull/7059
## Which issue does this PR close? Closes #7058. ## Rationale for this change The bypass-merge JVM columnar shuffle writer gives every partition its own `CometDiskBlockWriter`. Each one held a page of Spark's page size, 32 to 64 MB, as soon as it buffered a row. Spark sizes that page so a task's memory holds only 16 to 32 of them. With more partitions than that, every new page spilled a sibling writer that held about one row, so the shuffle came out as one-row Arrow IPC blocks. That writer is what `df.repartition(n)` uses when `n × spark.executor.cores ≤ 100`. `spark.executor.cores` reads as 1 when unset, so this covers local mode, standalone clusters without explicit cores, and 1- and 2-core executors. #7058 has the details. ## What changes are included in this PR? - `SpillWriter` asks a new `nextPageRequest(required)` hook what size to request for a page. The default is `required`, which the allocator rounds up to its page size as before, so the sort-based `SpillSorter` is unchanged. - `CometDiskBlockWriter` takes an initial and a maximum page size. Its writer starts each batch on a 64 KiB page and doubles each further page up to the maximum, based on the pages the batch already holds. Writing a batch out frees its pages, so the next batch starts small again. A partition then costs kilobytes until its rows need more. A batch takes a logarithmic number of pages, which stays well below the 8,192 pages Spark allows a task, where small fixed pages would not. - `CometBypassMergeSortShuffleWriter` creates its allocator with the 64 KiB initial page size, and passes the old page size, `min(128 MB, pageSizeBytes)`, as the maximum. - The JVM shuffle contributor guide describes the page sizing. Its selection rule also said "partitions < threshold" where the code uses `≤`, which is fixed here. ## How are these changes tested? Two new tests in `CometDiskBlockWriterSuite`: - "bypass writers with more partitions than full pages fit write without spilling" runs the full bypass writer over 32 partitions of 100 rows each, against a pool with room for eight 1 MiB pages. It asserts that nothing spilled. On `main` it reports 3.3 GB spilled: about 3,200 spills, each freeing a 1 MiB page that held one row. - "a partition writer's pages double from a small first page up to the page size" checks that a writer's memory goes 64, 192, 448, 960, 1984, 3008 and 4032 KiB as it takes pages, then drops back to 64 KiB once a batch is written out. Three mutations each fail one of these tests: the bypass writer back on full-size pages, pages that don't double, and a page count that never resets between batches. The JVM shuffle suites pass, 181 tests: `CometShuffleSuite`, `DisableAQECometShuffleSuite`, `CometShuffleEncryptionSuite`, `CometShuffleManagerSuite`, `CometShuffleInputConversionSuite`, `SpillSorterSuite`, both shuffle memory allocator suites, `CometDiskBlockWriterSuite` and `CometTaskMetricsSuite`. End to end, on TPC-H SF2 with a noop sink, Spark 4.1.3, `local[8]`, 6 GB off-heap and `spark.executor.cores` unset (so the bypass writer is used), release build on an Apple M3 Max. The machine was shared with other work, so the times are rough, and the shuffle sizes are exact. | Query | `main` | This PR | Spark | |---|---:|---:|---:| | `orders`, `REPARTITION(64)` | 266 s, 2.48 GB | 1.1 s, 156 MB | 1.0 s, 205 MB | | `lineitem`, `REPARTITION(64)` | 15.5 min | 3.3 s, 702 MB | 3.8 s, 1.13 GB | A 2-core executor, emulated with `local[2]`, 2 GB off-heap and `spark.executor.cores=2`, over 300K `orders` rows: | `REPARTITION(n)` | `main` | This PR | |---:|---:|---:| | 24 | 1.6 s (first run, includes JIT warmup), 15.6 MB | 0.41 s, 15.6 MB | | 32 | 0.46 s, 15.8 MB | 0.42 s, 15.8 MB | | 40 | 8.5 s, 125 MB | 0.41 s, 15.8 MB | | 48 | 16.7 s, 230 MB | 0.41 s, 15.9 MB | Spark takes 0.42 s and writes 20.7 MB at 48 partitions. -- 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]
