andygrove opened a new pull request, #7065: URL: https://github.com/apache/datafusion-comet/pull/7065
## Which issue does this PR close? Part of #5905 (item J3). ## Rationale for this change Above `spark.shuffle.sort.bypassMergeThreshold` partitions (200 by default), the JVM columnar shuffle uses the sort-based writer. It wrote each spill, and the final file, with one `writeSortedFileNative` JNI call per partition holding rows. Each call paid a fixed cost before it converted a row. It decoded every column's protobuf data type, copied the row address and size arrays that Java had built row by row in a boxed `ArrayBuffer`, and opened the spill file. It also allocated batch-size builders, including a dictionary builder per string column, pre-encoded the IPC schema, created a compression context, and issued its own `write(2)`. A map task pays that `partitions x (spills + 1)` times, so at high partition counts it dominates the write. In the native benchmark below it comes to about 20 µs per partition before any JNI or Java overhead, and this PR removes about three quarters of it. ## What changes are included in this PR? - `SpillSorter.writeSortedFileNative` makes one native call per spill. The new `Native.writeSortedPartitionsNative` takes the address of the sorter's sorted packed record pointers and a page table holding the base address of each data page by page number. Native code resolves each record where it lies, writes every partition in partition order, and returns each partition's length through `SpillInfo.partitionLengths` and its checksum through the sorter's checksum array. Java no longer walks the records or builds a `RowPartition` on this path. - In `row.rs`, the state that does not depend on the rows moves into a `SortedRowEncoder` that one call reuses for every partition: the builders, the block writer with its pre-encoded schema, the compression context and the batch schema. Encoded blocks are buffered up to 1 MiB before they are written, so a spill of many small partitions is no longer one write per partition. `process_sorted_partitions` is the new per-spill entry point. `process_sorted_row_partition`, which the bypass merge writer still calls once per partition, runs on the same encoder. - `Checksum::update` takes the bytes of one block, since the buffer can now hold several. - Benchmarks: a `sorted_spill` group in `native/shuffle/benches/row_columnar.rs` writes one sorted spill both ways, and `CometColumnarShuffleWriteBenchmark` runs the shuffle end to end at 1,000 and 10,000 partitions, with and without spills. - `jvm_shuffle.md` describes the two native write calls. This replaces the loop in `SpillSorter` that #6361 fixes for `spark.shuffle.checksum.enabled=false`. With checksums disabled the new call gets an empty checksum array and algorithm -1, and never indexes the array. Whichever of the two lands second has a conflict in `SpillSorter.java` to resolve. The output is byte-identical to `main`: the same blocks in the same order, with the same partition lengths and checksums. I checked this locally against `main`'s per-partition calls for lz4, zstd, snappy and no compression, with dictionary-encoded and plain strings, partitions spanning several batches, and up to 2,000 partitions. The bypass merge writer's output is also unchanged. ## How are these changes tested? - New Rust tests next to the existing `write_sorted_file` tests in `jni_api.rs`. One compares the per-spill output byte for byte with writing each partition on its own, for every codec and checksum algorithm and without checksums, over records in two pages. Another checks that checksums continued across two spills match CRC32 and Adler-32 computed independently over each partition's bytes in both files. The rest cover an empty spill and the error cases: a page the sorter does not hold, a partition id beyond the partition count, unsorted pointers, and an unsupported record length size. - A unit test of the record pointer decoding in `spark_unsafe::row`, which the nightly Miri job also runs. - A new `SpillSorterSuite` test drives a `SpillSorter` through a spill and a final write over 1 KiB pages, so the records span pages. It runs with both allocators, whose page tables differ, and with checksums disabled. It checks the partition lengths against the files, the Adler-32 checksums continued across both files using `java.util.zip.Adler32`, and the write and spill metrics. - Existing suites: `CometShuffleSuite` and `DisableAQECometShuffleSuite`, which run 201 partitions with a spill threshold of 10 so the sort path spills constantly, plus `SpillSorterSuite`, `CometDiskBlockWriterSuite`, `CometShuffle4_0Suite`, `CometShuffleEncryptionSuite`, `CometShuffleManagerSuite` and the two allocator suites. ### Benchmarks Run locally on an Apple M3 Ultra. Other work was running on the machine, so cases were interleaved and repeated. Native, one spill of 100,000 rows (Int64, Int32, Float64, Utf8, Date32), lz4, Adler-32 checksums, `sorted_spill`. Median of three interleaved runs: | Partitions | Rows per partition | One call per partition (`main`) | One call per spill (this PR) | Change | |---:|---:|---:|---:|---:| | 100 | 1,000 | 10.6 ms | 9.0 ms | -15% | | 1,000 | 100 | 31.8 ms | 15.8 ms | -50% | | 10,000 | 10 | 213 ms | 62 ms | -71% | `CometColumnarShuffleWriteBenchmark`, 1M rows, `local[1]`, best time of each run, two runs per side: | Partitions | Case | `main` | This PR | Change | |---:|---|---:|---:|---:| | 1,000 | Spark shuffle, for reference | 471 / 440 ms | 451 / 441 ms | | | 1,000 | Comet JVM shuffle | 385 / 343 ms | 277 / 271 ms | -21% to -28% | | 1,000 | Comet JVM shuffle, spill every 100,000 rows | 1,032 / 1,096 ms | 566 / 557 ms | -45% to -49% | | 10,000 | Spark shuffle, for reference | 733 / 731 ms | 761 / 736 ms | | | 10,000 | Comet JVM shuffle | 1,132 / 1,123 ms | 688 / 683 ms | -39% | | 10,000 | Comet JVM shuffle, spill every 100,000 rows | 6,591 / 7,890 ms | 3,186 / 3,122 ms | -52% to -60% | What is left per partition is the cost of one block per partition rather than of the call. With 10,000 partitions the per-spill path still spends about 5 µs per partition with lz4 and about 4 µs without compression. That goes to the per-block lz4 frame encoder (W1 in #5905), and to building and framing a small batch, including turning a string column's dictionary back into plain strings (W5 and J6). -- 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]
