Repository navigation
Conversation
The sort-based JVM columnar shuffle writer wrote each spill, and the final file, with one writeSortedFileNative JNI call per partition. Every call decoded the schema, copied the row address and size arrays, opened the spill file, allocated batch-size builders, pre-encoded the IPC schema and wrote its own blocks, so a map task paid that partitions x (spills + 1) times. SpillSorter now makes one writeSortedPartitionsNative call per spill. Native code reads the sorter's sorted packed record pointers through a page table, converts every partition with one reused encoder, buffers the encoded blocks, and returns each partition's length and checksum. The output is byte-identical to the per-partition calls, which the bypass merge writer still makes. Adds a sorted_spill group to the row_columnar bench and CometColumnarShuffleWriteBenchmark. Part of apache#5905 (J3).
andygrove
added
the
run-benchmark-check
github-actions
Bot
added
enhancement
The per-spill test compared the output byte for byte only for an Int64 column, so the batch schema and block writer that the encoder now keeps across partitions were never exercised when a string column switches between a dictionary and a plain array. A new test writes such a spill, compares it with writing each partition on its own, and decodes it. The SpillSorterSuite test checked partition lengths and checksums against the files, which a page table mistake that read valid rows from the wrong page would still pass. It now decodes every partition back to its rows. Also fixes the strict Scala warnings build, where the new benchmark widened an Int row count to Long, and rewords the JVM shuffle guide's write path, which read as one native call per partition.
With zstd, a block writer built for a plain string column stores its offsets as lengths and one built for a dictionary column does not, so the writer the encoder keeps across partitions changes the bytes of a block. The test now runs that codec as well as lz4.
SortedRowEncoder kept the batch schema twice, once to reuse it and once next to the block writer to notice that it changed. It now keeps the schema only with the block writer, compares the column types against it, and builds both when they change. The encoder also owns the output file: encode_batch writes the buffered blocks out once they reach 1 MiB, and finish writes the rest and returns the encode time, so neither caller has to remember the final write. Its builders are sized for the rows the call encodes, at most one batch. process_sorted_partitions splits the record pointers into partitions with chunk_by, and the benchmark groups its records the same way. SpillSorter no longer routes a spill's metrics through a throwaway ShuffleWriteMetrics, which only existed for the per-partition calls. It reports the same values to the task's metrics directly. bytes is only used by tests now, so it moves to dev-dependencies. The tests reuse long_values and createSpillSorter, and SpillSorterSuite reads each spill file once.
andygrove
added
the
run-spark-4.1-tests
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
Thanks for this. The per-spill call removes a cost that grew with partitions x (spills + 1), and the numbers line up with that, from -15% at 100 partitions to -71% at 10,000. I traced the pointer decoding against both allocators. CometUnboundedShuffleMemoryAllocator and the TaskMemoryManager path both encode the offset relative to the page base, and every record goes into a page in allocatedPages, so the page table covers what getOffsetInPage used to resolve. Moving the shared state into SortedRowEncoder also removes the duplication between the two native paths. The dictionary-change test with zstd offsets as lengths covers the one new risk that keeping the block writer across partitions brings, and the SpillSorterSuite case decodes the rows back, so a wrong page lookup would fail it.
One thing about the description. Rewriting the loop in SpillSorter also fixes the ArrayIndexOutOfBoundsException that sort-based JVM shuffle hits on branch-1.1 with spark.shuffle.checksum.enabled=false (#6256), and the new SpillSorterSuite case covers it. Could the description say this fixes the sort-writer half of #6256? The bypass path, where doSpilling overwrites checksum with Long.MIN_VALUE so the checksum != -1 check no longer works, is still open and left to #6361. This PR is a performance change and not something to backport, so #6361 (or a minimal cherry-pick of its SpillSorter fix) is what would bring the fix to branch-1.1. Saying that here would make it clear what each PR covers once this one lands.
A few inline comments below.
| * #setSpillInfo}, every partition in one native call, recording each partition's length and | ||
| * checksum. | ||
| */ | ||
| public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) throws IOException { |
There was a problem hiding this comment.
This method no longer calls Native.writeSortedFileNative. It calls writeSortedPartitionsNative, and the bypass writer, through SpillWriter.doSpilling, is now the only caller of the old native method. Two methods with the same name doing different things is why jvm_shuffle.md needs a paragraph to tell them apart. Would you consider renaming this one, for example to writeSortedPartitions? Then a grep for writeSortedFileNative would only find the per-partition path.
|
|
||
| /// Encoded blocks are buffered until they reach this many bytes, so that a spill of many small | ||
| /// partitions is not written one partition at a time. | ||
| const OUTPUT_BUFFER_BYTES: usize = 1 << 20; |
There was a problem hiding this comment.
The description says this buffer, like the builders, is native memory that no pool reserves, and that it holds at most 1 MiB plus one block. Could the doc comment here say the same? This is new buffering on this path, and someone reading the code later will want to know whether it is reserved without going back to the PR.
| * }}} | ||
| * Results will be written to "spark/benchmarks/CometColumnarShuffleWriteBenchmark-**results.txt". | ||
| */ | ||
| object CometColumnarShuffleWriteBenchmark extends CometBenchmarkBase { |
There was a problem hiding this comment.
CometShuffleBenchmark already targets the sort-based writer on purpose, through manyPartitions = 201 and spark.comet.shuffle.jvm.spillThreshold=30000. Is there a reason the 1,000 and 10,000 partition cases need their own benchmark object rather than a group in CometShuffleBenchmark? If it is the different session setup (CometShuffleBenchmark runs with spark.comet.exec.onHeap.enabled=true), a sentence in the scaladoc would explain it. Otherwise, keeping all JVM shuffle write benchmarks in one place would make them easier to maintain.
2 participants
Which issue does this PR close?
Closes #7061.
Part of #5905 (item J3).
Rationale for this change
Above
spark.shuffle.sort.bypassMergeThresholdpartitions (200 by default), the JVM columnar shuffle uses the sort-based writer. It wrote each spill, and the final file, with onewriteSortedFileNativeJNI 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 boxedArrayBuffer, 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 ownwrite(2). A map task pays thatpartitions 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.writeSortedFileNativemakes one native call per spill. The newNative.writeSortedPartitionsNativetakes 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 throughSpillInfo.partitionLengthsand its checksum through the sorter's checksum array. Java no longer walks the records or builds aRowPartitionon this path.row.rs, the state that does not depend on the rows moves into aSortedRowEncoderthat one call reuses for every partition: the builders, the block writer with its pre-encoded schema, the compression context and the open file. 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. Like the builders, the buffer is native memory that no memory pool reserves. It holds at most 1 MiB plus one block. The builders are sized for the rows a call encodes, at most one batch, so a call with a few rows no longer allocates batch-size builders.process_sorted_partitionsis 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::updatetakes the bytes of one block, since the buffer can now hold several.sorted_spillgroup innative/shuffle/benches/row_columnar.rswrites one sorted spill both ways, andCometColumnarShuffleWriteBenchmarkruns the shuffle end to end at 1,000 and 10,000 partitions, with and without spills.jvm_shuffle.mddescribes the two native write calls.This replaces the loop in
SpillSorterthat #6361 fixes forspark.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 inSpillSorter.javato 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 againstmain'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?
write_sorted_filetests injni_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. A second does the same for a string column whose batches switch between dictionary-encoded and plain, within a partition and from one partition to the next, with lz4 and with zstd storing offsets as lengths, and decodes the strings back. 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.spark_unsafe::row, which the nightly Miri job also runs.SpillSorterSuitetest drives aSpillSorterthrough 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 decodes every partition of both files back to its rows, and checks the partition lengths against the files, the Adler-32 checksums continued across both files usingjava.util.zip.Adler32, and the write and spill metrics.CometShuffleSuiteandDisableAQECometShuffleSuite, which run 201 partitions with a spill threshold of 10 so the sort path spills constantly, plusSpillSorterSuite,CometDiskBlockWriterSuite,CometShuffle4_0Suite,CometShuffleEncryptionSuite,CometShuffleManagerSuiteand 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:main)CometColumnarShuffleWriteBenchmark, 1M rows,local[1], best time of each run, two runs per side:mainThe repro from #7061: TPC-H SF2
orders(3M rows),SELECT /*+ REPARTITION(n) */ * FROM orderswith a noop sink,local[8], 6 GB off-heap andspark.executor.cores=8, on Spark 4.1.1. Median of 5 warm runs, two sessions per side:mainAt 2,000 partitions the map stage's total task time drops from 6.8 s to 3.0 s, against 3.7 to 4.8 s for Spark. The reduce stage is unchanged.
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 (#6916), and to building and framing a small batch, including turning a string column's dictionary back into plain strings (#6597).