From 5c53def8865cb66d0bf165b507bf81742b4ff045 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 13:56:27 -0600 Subject: [PATCH 1/5] perf: write every partition of a JVM shuffle spill in one native call 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 #5905 (J3). --- docs/source/contributor-guide/jvm_shuffle.md | 26 +- native/core/src/execution/jni_api.rs | 419 +++++++++++++++- native/shuffle/benches/row_columnar.rs | 185 ++++++- native/shuffle/src/spark_unsafe/row.rs | 473 +++++++++++++++--- native/shuffle/src/writers/checksum.rs | 38 +- .../CometUnboundedShuffleMemoryAllocator.java | 7 +- .../spark/shuffle/sort/SpillSorter.java | 152 +++--- .../comet/execution/shuffle/SpillWriter.java | 26 +- .../main/scala/org/apache/comet/Native.scala | 58 ++- .../spark/shuffle/sort/SpillSorterSuite.scala | 131 +++++ .../CometColumnarShuffleWriteBenchmark.scala | 119 +++++ 11 files changed, 1412 insertions(+), 222 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index f41877b08dc..e64e714d5ce 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -152,13 +152,14 @@ Selection logic in `CometShuffleManager.shouldBypassMergeSort()`: ### Writers -| Class | Location | Description | -| ----------------------------------- | ---------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------- | -| `CometBypassMergeSortShuffleWriter` | `.../shuffle/CometBypassMergeSortShuffleWriter.java` | Hash-based writer. Creates one `CometDiskBlockWriter` per partition. Supports async writes. | -| `CometUnsafeShuffleWriter` | `.../shuffle/CometUnsafeShuffleWriter.java` | Sort-based writer. Uses `CometShuffleExternalSorter` to buffer and sort records, then merges spill files. | -| `CometDiskBlockWriter` | `.../shuffle/CometDiskBlockWriter.java` | Buffers rows in memory pages for a single partition. Spills to disk via native encoding. Used by bypass writer. | -| `CometShuffleExternalSorter` | `.../shuffle/sort/CometShuffleExternalSorter.java` | Buffers records across all partitions, sorts by partition ID, spills sorted data. Used by unsafe writer. | -| `SpillWriter` | `.../shuffle/SpillWriter.java` | Base class for spill logic. Manages memory pages and calls `Native.writeSortedFileNative()` for Arrow IPC encoding. | +| Class | Location | Description | +| ----------------------------------- | ---------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | +| `CometBypassMergeSortShuffleWriter` | `.../shuffle/CometBypassMergeSortShuffleWriter.java` | Hash-based writer. Creates one `CometDiskBlockWriter` per partition. Supports async writes. | +| `CometUnsafeShuffleWriter` | `.../shuffle/CometUnsafeShuffleWriter.java` | Sort-based writer. Uses `CometShuffleExternalSorter` to buffer and sort records, then merges spill files. | +| `CometDiskBlockWriter` | `.../shuffle/CometDiskBlockWriter.java` | Buffers rows in memory pages for a single partition. Spills to disk via native encoding. Used by bypass writer. | +| `CometShuffleExternalSorter` | `.../shuffle/sort/CometShuffleExternalSorter.java` | Buffers records across all partitions, sorts by partition ID, spills sorted data. Used by unsafe writer. | +| `SpillSorter` | `.../shuffle/sort/SpillSorter.java` | Holds one spill's records for `CometShuffleExternalSorter`. Writes every partition of the spill with one `Native.writeSortedPartitionsNative()` call. | +| `SpillWriter` | `.../shuffle/SpillWriter.java` | Base class for spill logic. Manages memory pages. `doSpilling()` encodes one partition's rows with `Native.writeSortedFileNative()`. | ### Reader @@ -174,8 +175,15 @@ Selection logic in `CometShuffleManager.shouldBypassMergeSort()`: 1. `ShuffleWriteProcessor` calls `CometShuffleManager.getWriter()` 2. Writer receives `Iterator[Product2[K, V]]` where V is `UnsafeRow` 3. Rows are serialized and buffered in off-heap memory pages -4. When memory threshold or batch size is reached, `SpillWriter.doSpilling()` is called -5. Native code (`Native.writeSortedFileNative()`) converts rows to Arrow arrays and writes IPC format +4. When memory threshold or batch size is reached, the buffered rows are written: `CometDiskBlockWriter` + calls `SpillWriter.doSpilling()` for its partition, and `CometShuffleExternalSorter` calls + `SpillSorter.writeSortedFileNative()` for every partition +5. Native code converts the rows to Arrow arrays and writes them as Arrow IPC blocks. `doSpilling()` + passes one partition's row addresses and sizes to `Native.writeSortedFileNative()`. `SpillSorter` + sorts its record pointers by partition with `Native.sortRowPartitionsNative()` and passes them, with + the address of each data page, to `Native.writeSortedPartitionsNative()`, which reads the rows in + place, writes every partition of the spill in one call, and returns each partition's length and + checksum 6. For bypass writer: partition files are concatenated into final output 7. For sort writer: spill files are merged diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index cabbb58d2bc..4def25abc6e 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -33,8 +33,11 @@ use super::{serde, utils::SparkArrowConvert}; use crate::{ errors::{try_unwrap_or_throw, CometError, CometResult}, execution::{ - metrics::utils::update_comet_metric, planner::PhysicalPlanner, serde::to_arrow_datatype, - shuffle::spark_unsafe::row::process_sorted_row_partition, sort::RdxSort, + metrics::utils::update_comet_metric, + planner::PhysicalPlanner, + serde::to_arrow_datatype, + shuffle::spark_unsafe::row::{process_sorted_partitions, process_sorted_row_partition}, + sort::RdxSort, }, jvm_bridge::{JVMClasses, JavaShufflePartitionPusher, ShufflePartitionPusher}, }; @@ -1755,7 +1758,7 @@ fn shuffle_partition_offsets( Ok(offsets) } -/// Used by Comet shuffle external sorter to write sorted records to disk. +/// Used by the bypass merge sort shuffle writer to write one partition's records to disk. /// # Safety /// This function is inherently unsafe since it deals with raw pointers passed from JNI. #[no_mangle] @@ -1848,12 +1851,7 @@ unsafe fn write_sorted_file( Some(current_checksum as u32) }; - let compression_codec = match compression_codec { - "zstd" => CompressionCodec::Zstd(compression_level), - "lz4" => CompressionCodec::Lz4Frame, - "snappy" => CompressionCodec::Snappy, - _ => CompressionCodec::Lz4Frame, - }; + let compression_codec = jvm_shuffle_codec(compression_codec, compression_level); // `process_sorted_row_partition` only reads through the pointers. let (written_bytes, checksum, encode_nanos) = process_sorted_row_partition( @@ -1875,6 +1873,143 @@ unsafe fn write_sorted_file( Ok([written_bytes, checksum, encode_nanos]) } +/// The codec `spark.comet.shuffle.compression.codec` names. Unknown names fall back to lz4. +fn jvm_shuffle_codec(name: &str, level: i32) -> CompressionCodec { + match name { + "zstd" => CompressionCodec::Zstd(level), + "lz4" => CompressionCodec::Lz4Frame, + "snappy" => CompressionCodec::Snappy, + _ => CompressionCodec::Lz4Frame, + } +} + +/// Used by Comet shuffle external sorter to write every partition of one sorted spill to disk. +/// # Safety +/// This function is inherently unsafe since it deals with raw pointers passed from JNI. +#[no_mangle] +pub unsafe extern "system" fn Java_org_apache_comet_Native_writeSortedPartitionsNative( + e: EnvUnowned, + _class: JClass, + record_pointers_addr: jlong, + num_records: jint, + page_addresses: JLongArray, + uao_size: jint, + serialized_datatypes: JObjectArray, + file_path: JString, + prefer_dictionary_ratio: jdouble, + batch_size: jint, + checksum_algo: jint, + partition_checksums: JLongArray, + partition_lengths: JLongArray, + compression_codec: JString, + compression_level: jint, + offsets_as_lengths: jboolean, + tracing_enabled: jboolean, +) -> jlong { + try_unwrap_or_throw(&e, |env| unsafe { + with_trace( + "writeSortedPartitionsNative", + tracing_enabled != JNI_FALSE, + || { + let num_records = usize::try_from(num_records).map_err(|_| { + CometError::Internal(format!("Negative record count {num_records}")) + })?; + let record_pointers: &[i64] = if num_records == 0 { + &[] + } else { + // SAFETY: JVM unsafe memory allocation is aligned with long. + debug_assert!( + record_pointers_addr != 0, + "writeSortedPartitionsNative: null record pointers" + ); + debug_assert_eq!( + (record_pointers_addr as usize) % std::mem::align_of::(), + 0, + "writeSortedPartitionsNative: record pointers not aligned to i64" + ); + std::slice::from_raw_parts(record_pointers_addr as *const i64, num_records) + }; + let data_types = convert_datatype_arrays(env, serialized_datatypes)?; + let mut pages = vec![0_i64; page_addresses.len(env)?]; + page_addresses.get_region(env, 0, &mut pages)?; + let mut checksums = vec![0_i64; partition_checksums.len(env)?]; + partition_checksums.get_region(env, 0, &mut checksums)?; + let mut lengths = vec![0_i64; partition_lengths.len(env)?]; + let output_path: String = file_path.try_to_string(env)?; + let compression_codec: String = compression_codec.try_to_string(env)?; + + let encode_nanos = write_sorted_partitions( + record_pointers, + &pages, + uao_size, + &data_types, + &output_path, + prefer_dictionary_ratio, + batch_size, + checksum_algo, + &mut checksums, + &mut lengths, + &compression_codec, + compression_level, + offsets_as_lengths != JNI_FALSE, + )?; + + partition_lengths.set_region(env, 0, &lengths)?; + partition_checksums.set_region(env, 0, &checksums)?; + Ok(encode_nanos) + }, + ) + }) +} + +/// Converts the rows of one sorted spill, every partition of it, into Arrow batches of +/// `data_types` and appends them to `output_path` in partition order. `record_pointers` are the +/// sorter's sorted Spark `PackedRecordPointer`s, whose page numbers index `page_addresses`. +/// Fills `partition_lengths` and, when `checksum_algo` is not -1, continues each written +/// partition's checksum in `partition_checksums`, where `i64::MIN` is a checksum not yet started. +/// Returns the encode nanos. Core of `Native.writeSortedPartitionsNative`. +/// +/// # Safety +/// Each pointer must resolve, through `page_addresses`, to a readable record. +#[allow(clippy::too_many_arguments)] +unsafe fn write_sorted_partitions( + record_pointers: &[i64], + page_addresses: &[i64], + uao_size: i32, + data_types: &[ArrowDataType], + output_path: &str, + prefer_dictionary_ratio: f64, + batch_size: i32, + checksum_algo: i32, + partition_checksums: &mut [i64], + partition_lengths: &mut [i64], + compression_codec: &str, + compression_level: i32, + offsets_as_lengths: bool, +) -> CometResult { + let uao_size = usize::try_from(uao_size) + .map_err(|_| CometError::Internal(format!("Unsupported record length size {uao_size}")))?; + let batch_size = usize::try_from(batch_size) + .map_err(|_| CometError::Internal(format!("Unsupported batch size {batch_size}")))?; + let checksum_algo = (checksum_algo != -1).then_some(checksum_algo); + let compression_codec = jvm_shuffle_codec(compression_codec, compression_level); + + process_sorted_partitions( + record_pointers, + page_addresses, + uao_size, + data_types, + output_path, + prefer_dictionary_ratio, + batch_size, + checksum_algo, + partition_checksums, + partition_lengths, + &compression_codec, + offsets_as_lengths, + ) +} + #[no_mangle] /// Used by Comet shuffle external sorter to sort in-memory row partition ids. pub extern "system" fn Java_org_apache_comet_Native_sortRowPartitionsNative( @@ -3471,6 +3606,272 @@ mod tests { ); } + /// The page numbers `SorterRecords` stores rows under. They are not contiguous, so that a + /// lookup that ignored the page number would read the wrong page. + const SORTER_PAGE_NUMBERS: [usize; 2] = [3, 1]; + + /// 16-byte rows stored the way Spark's sorter stores them: in pages, each record its length + /// in 4 bytes, then the row's 4-byte big-endian size from `UnsafeRowSerializer`, then the + /// row. Rows alternate between two pages, and `pointers` holds the packed record pointer of + /// each row, sorted by partition as `sort_row_partitions` sorts them. + struct SorterRecords { + _pages: Vec>, + page_addresses: Vec, + pointers: Vec, + } + + impl SorterRecords { + fn new(rows: &[[u8; 16]], partitions: &[usize]) -> Self { + let mut page_bytes: Vec> = vec![Vec::new(); SORTER_PAGE_NUMBERS.len()]; + let mut pointers = Vec::new(); + for (i, (row, &partition)) in rows.iter().zip(partitions).enumerate() { + let page = i % SORTER_PAGE_NUMBERS.len(); + let offset = page_bytes[page].len(); + page_bytes[page].extend_from_slice(&(4 + row.len() as i32).to_ne_bytes()); + page_bytes[page].extend_from_slice(&(row.len() as i32).to_be_bytes()); + page_bytes[page].extend_from_slice(row); + pointers.push( + (partition as i64) << 40 + | (SORTER_PAGE_NUMBERS[page] as i64) << 27 + | offset as i64, + ); + } + sort_row_partitions(&mut pointers); + + // `u64` elements keep the rows 8-byte aligned, as in a Spark page. + let pages: Vec> = page_bytes + .iter() + .map(|bytes| { + let mut page = vec![0_u64; bytes.len().div_ceil(8)]; + unsafe { + std::ptr::copy_nonoverlapping( + bytes.as_ptr(), + page.as_mut_ptr() as *mut u8, + bytes.len(), + ) + }; + page + }) + .collect(); + let mut page_addresses = vec![0_i64; SORTER_PAGE_NUMBERS.iter().max().unwrap() + 1]; + for (page, &number) in pages.iter().zip(&SORTER_PAGE_NUMBERS) { + page_addresses[number] = page.as_ptr() as i64; + } + Self { + _pages: pages, + page_addresses, + pointers, + } + } + + /// Writes the records with `write_sorted_partitions` in batches of 2 rows. + fn write( + &self, + path: &std::path::Path, + checksum_algo: i32, + partition_checksums: &mut [i64], + partition_lengths: &mut [i64], + codec: &str, + ) -> CometResult { + unsafe { + write_sorted_partitions( + &self.pointers, + &self.page_addresses, + 4, + &[DataType::Int64], + path.to_str().unwrap(), + 1.0, + 2, + checksum_algo, + partition_checksums, + partition_lengths, + codec, + 1, + false, + ) + } + } + } + + /// The partition of row `i` of `long_rows(0..23)`: partitions 1, 3 and 6 of 7 get no rows, + /// and every other one more than one 2-row batch. + fn partition_of(i: usize) -> usize { + [0, 2, 4, 5, 2, 0, 4][i % 7] + } + + #[test] + fn write_sorted_partitions_matches_writing_each_partition() { + let values: Vec = (0..23).collect(); + let rows = long_rows(&values); + let partitions: Vec = (0..rows.len()).map(partition_of).collect(); + let num_partitions = 7; + let records = SorterRecords::new(&rows, &partitions); + let dir = tempfile::tempdir().unwrap(); + + // Every codec and checksum algorithm, and no checksum (-1) with no checksum array. + for (codec, algo) in [("lz4", 1), ("zstd", 0), ("snappy", 2), ("lz4", -1)] { + let name = format!("{codec}-{algo}"); + let path = dir.path().join(&name); + let mut checksums = vec![i64::MIN; if algo == -1 { 0 } else { num_partitions }]; + let mut lengths = vec![0_i64; num_partitions]; + records + .write(&path, algo, &mut checksums, &mut lengths, codec) + .unwrap(); + + // Writing each partition's rows, in the order the sorter keeps them, on its own. + let mut expected = Vec::new(); + for partition in 0..num_partitions { + let partition_rows: Vec<[u8; 16]> = rows + .iter() + .zip(&partitions) + .filter(|(_, &p)| p == partition) + .map(|(row, _)| *row) + .collect(); + if partition_rows.is_empty() { + assert_eq!(lengths[partition], 0, "{name}: partition {partition}"); + if algo != -1 { + assert_eq!(checksums[partition], i64::MIN, "{name}: {partition}"); + } + continue; + } + let sizes = vec![16_i32; partition_rows.len()]; + let ([written, checksum, _], data) = write_long_rows( + &partition_rows, + &sizes, + &dir.path().join(format!("{name}-{partition}")), + (algo != -1).then_some(algo), + i64::MIN, + codec, + ) + .unwrap(); + assert_eq!(lengths[partition], written, "{name}: partition {partition}"); + if algo != -1 { + assert_eq!(checksums[partition], checksum, "{name}: {partition}"); + } + expected.extend(data); + } + let data = std::fs::read(&path).unwrap(); + assert_eq!(data, expected, "{name}"); + + // The blocks hold every row, partition by partition. + let mut decoded = Vec::new(); + for block in shuffle_blocks(&data) { + let batch = read_ipc_compressed(block).unwrap(); + let column = arrow::array::AsArray::as_primitive::( + batch.column(0), + ); + decoded.extend(column.values().iter().copied()); + } + let mut expected_values = values.clone(); + expected_values.sort_by_key(|&v| partition_of(v as usize)); + assert_eq!(decoded, expected_values, "{name}"); + } + } + + #[test] + fn write_sorted_partitions_continues_partition_checksums() { + let rows = long_rows(&(0..23).collect::>()); + let partitions: Vec = (0..rows.len()).map(partition_of).collect(); + let num_partitions = 7; + let records = SorterRecords::new(&rows, &partitions); + let dir = tempfile::tempdir().unwrap(); + + // Two spills share the checksums, so after the second each partition's checksum covers + // its bytes in both files, as Spark checksums the merged partition. + let algorithms = [(0, crc32fast::hash as fn(&[u8]) -> u32), (1, adler32)]; + for (algo, checksum_of) in algorithms { + let mut checksums = vec![i64::MIN; num_partitions]; + let mut spills = Vec::new(); + for spill in 0..2 { + let path = dir.path().join(format!("{algo}-{spill}")); + let mut lengths = vec![0_i64; num_partitions]; + records + .write(&path, algo, &mut checksums, &mut lengths, "lz4") + .unwrap(); + spills.push((std::fs::read(&path).unwrap(), lengths)); + } + + for partition in [0, 2, 4, 5] { + let mut bytes = Vec::new(); + for (data, lengths) in &spills { + let start = lengths[..partition].iter().sum::() as usize; + bytes.extend_from_slice(&data[start..start + lengths[partition] as usize]); + } + assert_eq!( + checksums[partition], + checksum_of(&bytes) as i64, + "algorithm {algo}, partition {partition}" + ); + } + } + } + + #[test] + fn write_sorted_partitions_without_records() { + let records = SorterRecords::new(&[], &[]); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("empty"); + let mut checksums = vec![i64::MIN; 3]; + let mut lengths = vec![0_i64; 3]; + records + .write(&path, 1, &mut checksums, &mut lengths, "lz4") + .unwrap(); + assert_eq!(lengths, [0, 0, 0]); + assert_eq!(checksums, [i64::MIN; 3]); + assert!(std::fs::read(&path).unwrap().is_empty()); + } + + #[test] + fn write_sorted_partitions_rejects_bad_record_pointers() { + let rows = long_rows(&[1, 2, 3]); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("bad"); + let write = |records: &SorterRecords, num_partitions: usize| { + records + .write(&path, -1, &mut [], &mut vec![0; num_partitions], "lz4") + .unwrap_err() + .to_string() + }; + + let records = SorterRecords::new(&rows, &[0, 1, 2]); + let error = write(&records, 2); + assert!(error.contains("partition 2 of 2"), "{error}"); + + let mut unsorted = SorterRecords::new(&rows, &[0, 1, 2]); + unsorted.pointers.reverse(); + let error = write(&unsorted, 3); + assert!(error.contains("not sorted by partition"), "{error}"); + + let mut missing_page = SorterRecords::new(&rows, &[0, 1, 2]); + missing_page.page_addresses[SORTER_PAGE_NUMBERS[1]] = 0; + let error = write(&missing_page, 3); + assert!( + error.contains(&format!("page {}", SORTER_PAGE_NUMBERS[1])), + "{error}" + ); + + let error = unsafe { + write_sorted_partitions( + &records.pointers, + &records.page_addresses, + 2, + &[DataType::Int64], + path.to_str().unwrap(), + 1.0, + 2, + -1, + &mut [], + &mut [0; 3], + "lz4", + 1, + false, + ) + } + .unwrap_err() + .to_string(); + assert!(error.contains("record length size 2"), "{error}"); + } + /// Exports `columns` as Arrow C Data structs and converts them with `columnar_to_row_convert`, /// returning the offsets, the lengths, and the row bytes. fn convert_exported( diff --git a/native/shuffle/benches/row_columnar.rs b/native/shuffle/benches/row_columnar.rs index e6a04b822cf..5687281bc96 100644 --- a/native/shuffle/benches/row_columnar.rs +++ b/native/shuffle/benches/row_columnar.rs @@ -23,7 +23,9 @@ use arrow::datatypes::{DataType as ArrowDataType, Field, Fields}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; -use datafusion_comet_shuffle::spark_unsafe::row::{process_sorted_row_partition, SparkUnsafeRow}; +use datafusion_comet_shuffle::spark_unsafe::row::{ + process_sorted_partitions, process_sorted_row_partition, SparkUnsafeRow, +}; use datafusion_comet_shuffle::spark_unsafe::unsafe_object::SparkUnsafeObject; use datafusion_comet_shuffle::CompressionCodec; use std::sync::Arc; @@ -378,6 +380,184 @@ fn benchmark_map_conversion(c: &mut Criterion) { group.finish(); } +// ─── Sorted spill ─────────────────────────────────────────────────────────── + +/// Rows in one spill of the sort-based JVM shuffle writer. +const SPILL_ROWS: usize = 100_000; + +/// Defaults of `spark.comet.shuffle.jvm.batchSize` and +/// `spark.comet.shuffle.jvm.preferDictionary.ratio`. +const JVM_SHUFFLE_BATCH_SIZE: usize = 8192; +const JVM_SHUFFLE_PREFER_DICTIONARY_RATIO: f64 = 10.0; + +/// Spark's default shuffle checksum, `spark.shuffle.checksum.algorithm=ADLER32`. +const ADLER32: i32 = 1; + +/// Bytes Spark's sorter stores before each record for its length (`UnsafeAlignedOffset`). +const UAO_SIZE: usize = 4; + +/// Bytes `UnsafeRowSerializer` writes before each row for its size. +const SERIALIZED_ROW_SIZE_BYTES: usize = 4; + +fn spill_schema() -> Vec { + vec![ + ArrowDataType::Int64, + ArrowDataType::Int32, + ArrowDataType::Float64, + ArrowDataType::Utf8, + ArrowDataType::Date32, + ] +} + +/// One spill of `SPILL_ROWS` rows of `spill_schema`, laid out in a single page the way +/// `SpillSorter.insertRecord` stores records, and sorted by partition. +struct SortedSpill { + /// Keeps the rows alive: `u64` elements keep every row 8-byte aligned, as in a Spark page. + page: Vec, + /// The packed record pointers, sorted by partition, as the sorter holds them. + record_pointers: Vec, + /// Address and size of each row, grouped by partition: the arrays the JVM builds for each + /// call when it writes the spill one partition at a time. + partitions: Vec<(Vec, Vec)>, +} + +fn build_sorted_spill(num_partitions: usize) -> SortedSpill { + let schema = spill_schema(); + let bitset = SparkUnsafeRow::get_row_bitset_width(schema.len()); + let fixed = bitset + schema.len() * UNSAFE_ROW_POINTER_SIZE; + + let mut bytes: Vec = Vec::new(); + let mut records: Vec<(usize, usize, usize)> = Vec::with_capacity(SPILL_ROWS); + for i in 0..SPILL_ROWS { + // A low-cardinality string column, padded to a word as UnsafeRow pads variable data. + let s = format!("status-{:03}", i % 100); + let row_size = fixed + s.len().div_ceil(8) * 8; + let mut row = vec![0u8; row_size]; + let slot = |field: usize| bitset + field * UNSAFE_ROW_POINTER_SIZE; + row[slot(0)..slot(0) + 8].copy_from_slice(&(i as i64).to_le_bytes()); + row[slot(1)..slot(1) + 4].copy_from_slice(&((i % 1000) as i32).to_le_bytes()); + row[slot(2)..slot(2) + 8].copy_from_slice(&(i as f64 * 0.5).to_le_bytes()); + write_pointer(&mut row, slot(3), fixed, s.len()); + row[fixed..fixed + s.len()].copy_from_slice(s.as_bytes()); + row[slot(4)..slot(4) + 4].copy_from_slice(&((i % 20000) as i32).to_le_bytes()); + + let record_offset = bytes.len(); + let record_len = SERIALIZED_ROW_SIZE_BYTES + row_size; + bytes.extend_from_slice(&(record_len as u32).to_ne_bytes()); + bytes.extend_from_slice(&(row_size as u32).to_be_bytes()); + bytes.extend_from_slice(&row); + + // Spread rows over the partitions the way a hash partitioner would. + let partition = + ((i as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 32) as usize % num_partitions; + records.push((partition, record_offset, row_size)); + } + // A stable sort keeps each partition's rows in page order, like sorting packed pointers. + records.sort_by_key(|&(partition, _, _)| partition); + + let mut page = vec![0u64; bytes.len().div_ceil(8)]; + // SAFETY: `page` holds at least `bytes.len()` bytes. + unsafe { + std::ptr::copy_nonoverlapping(bytes.as_ptr(), page.as_mut_ptr() as *mut u8, bytes.len()); + } + let page_addr = page.as_ptr() as i64; + + let mut partitions: Vec<(Vec, Vec)> = Vec::new(); + let mut current = usize::MAX; + for &(partition, record_offset, row_size) in &records { + if partition != current { + partitions.push((Vec::new(), Vec::new())); + current = partition; + } + let (addrs, sizes) = partitions.last_mut().unwrap(); + addrs.push(page_addr + (record_offset + UAO_SIZE + SERIALIZED_ROW_SIZE_BYTES) as i64); + sizes.push(row_size as i32); + } + // Every record is in page 0, so a pointer is the partition id and the offset in the page. + let record_pointers = records + .iter() + .map(|&(partition, record_offset, _)| ((partition as i64) << 40) | record_offset as i64) + .collect(); + + SortedSpill { + page, + record_pointers, + partitions, + } +} + +/// Writes one sorted spill the two ways the JVM shuffle writers call native code, at partition +/// counts where a partition gets 1000, 100 and 10 of the spill's rows: one call per partition, +/// with the partition's row addresses and sizes, and one call for the whole spill, with the +/// sorter's packed record pointers. +fn benchmark_sorted_spill(c: &mut Criterion) { + let mut group = c.benchmark_group("sorted_spill"); + let schema = spill_schema(); + let codec = CompressionCodec::Lz4Frame; + + for num_partitions in [100, 1000, 10000] { + let spill = build_sorted_spill(num_partitions); + + group.bench_function( + BenchmarkId::new("per_partition", format!("partitions_{num_partitions}")), + |b| { + b.iter(|| { + let tmp = Builder::new().tempfile().unwrap(); + let path = tmp.path().to_str().unwrap().to_string(); + for (addrs, sizes) in &spill.partitions { + process_sorted_row_partition( + addrs.len(), + JVM_SHUFFLE_BATCH_SIZE, + addrs.as_ptr() as *mut i64, + sizes.as_ptr() as *mut i32, + &schema, + path.clone(), + JVM_SHUFFLE_PREFER_DICTIONARY_RATIO, + true, + ADLER32, + None, + &codec, + true, + ) + .unwrap(); + } + }); + }, + ); + + let page_addresses = [spill.page.as_ptr() as i64]; + group.bench_function( + BenchmarkId::new("per_spill", format!("partitions_{num_partitions}")), + |b| { + b.iter(|| { + let tmp = Builder::new().tempfile().unwrap(); + let mut checksums = vec![i64::MIN; num_partitions]; + let mut lengths = vec![0_i64; num_partitions]; + unsafe { + process_sorted_partitions( + &spill.record_pointers, + &page_addresses, + UAO_SIZE, + &schema, + tmp.path().to_str().unwrap(), + JVM_SHUFFLE_PREFER_DICTIONARY_RATIO, + JVM_SHUFFLE_BATCH_SIZE, + Some(ADLER32), + &mut checksums, + &mut lengths, + &codec, + true, + ) + } + .unwrap(); + }); + }, + ); + } + + group.finish(); +} + fn config() -> Criterion { Criterion::default() } @@ -388,6 +568,7 @@ criterion_group! { targets = benchmark_primitive_columns, benchmark_struct_conversion, benchmark_list_conversion, - benchmark_map_conversion + benchmark_map_conversion, + benchmark_sorted_spill } criterion_main!(benches); diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index 260dca4a2ea..8bdbb66d87a 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -1357,6 +1357,186 @@ fn make_builders( Ok(builder) } +/// Bits of Spark's `PackedRecordPointer`: the partition id in the upper 24 bits, then the page +/// number in 13 bits and the offset in the page in the lower 27. +const PACKED_PARTITION_ID_SHIFT: u32 = 40; +const PACKED_PAGE_NUMBER_SHIFT: u32 = 27; +const PACKED_PAGE_NUMBER_MASK: u64 = (1 << 13) - 1; +const PACKED_OFFSET_MASK: u64 = (1 << 27) - 1; + +/// Bytes `UnsafeRowSerializer` writes before each row for the row's size. +const SERIALIZED_ROW_SIZE_BYTES: i64 = 4; + +/// 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; + +/// Converts batches of Spark `UnsafeRow`s to Arrow and encodes each batch as a shuffle block. +/// It holds the state that does not depend on the rows, so one encoder serves every partition of +/// a spill: the array builders, the block writer with its pre-encoded schema, and the +/// compression context. +struct SortedRowEncoder<'a> { + schema: &'a [DataType], + prefer_dictionary_ratio: f64, + codec: &'a CompressionCodec, + offsets_as_lengths: bool, + /// `builder_to_array` finishes every builder, which resets it for the next batch. + data_builders: Vec>, + /// The schema of the last batch, reused while the columns keep their types. + batch_schema: Option, + /// The block writer pre-encodes the batch schema, so it is kept until the schema changes. + block_writer: Option<(SchemaRef, ShuffleBlockWriter)>, + codec_context: ShuffleCodecContext, + /// Encoded blocks not yet written to the output. + buffer: Vec, + /// Encode and compression time across all batches. + ipc_time: Time, +} + +impl<'a> SortedRowEncoder<'a> { + fn try_new( + schema: &'a [DataType], + batch_size: usize, + prefer_dictionary_ratio: f64, + codec: &'a CompressionCodec, + offsets_as_lengths: bool, + ) -> Result { + let data_builders = schema + .iter() + .map(|dt| make_builders(dt, batch_size, prefer_dictionary_ratio)) + .collect::, _>>()?; + Ok(Self { + schema, + prefer_dictionary_ratio, + codec, + offsets_as_lengths, + data_builders, + batch_schema: None, + block_writer: None, + codec_context: ShuffleCodecContext::default(), + buffer: Vec::new(), + ipc_time: Time::default(), + }) + } + + /// Converts rows `row_start..row_end` of the rows at `row_addresses_ptr` and `row_sizes_ptr` + /// to one batch and buffers its shuffle block, adding the block to `checksum`. Returns the + /// block's size. + /// + /// # Safety + /// + /// As for `append_columns`: both pointers must point to at least `row_end` values, and each + /// address to a readable row of the corresponding size. + fn encode_batch( + &mut self, + row_addresses_ptr: *mut jlong, + row_sizes_ptr: *mut jint, + row_start: usize, + row_end: usize, + checksum: Option<&mut Checksum>, + ) -> Result { + for (idx, builder) in self.data_builders.iter_mut().enumerate() { + append_columns( + row_addresses_ptr, + row_sizes_ptr, + row_start, + row_end, + self.schema, + idx, + builder, + self.prefer_dictionary_ratio, + )?; + } + let arrays = self + .data_builders + .iter_mut() + .zip(self.schema.iter()) + .map(|(builder, datatype)| { + builder_to_array(builder, datatype, self.prefer_dictionary_ratio) + }) + .collect::, _>>()?; + let batch = self.make_batch(arrays, row_end - row_start)?; + + // The schema changes between batches when a dictionary-encoded column falls back to a + // plain array (see `builder_to_array`). + if self + .block_writer + .as_ref() + .is_none_or(|(cached, _)| !Arc::ptr_eq(cached, batch.schema_ref())) + { + let writer = ShuffleBlockWriter::try_new_with_offsets_as_lengths( + batch.schema_ref(), + self.codec.clone(), + self.offsets_as_lengths, + )?; + self.block_writer = Some((batch.schema(), writer)); + } + let (_, writer) = self + .block_writer + .as_ref() + .expect("block writer is set above"); + + let start = self.buffer.len(); + let mut cursor = Cursor::new(&mut self.buffer); + cursor.set_position(start as u64); + let written = + writer.write_batch(&batch, &mut cursor, &mut self.codec_context, &self.ipc_time)?; + if let Some(checksum) = checksum { + checksum.update(&self.buffer[start..]); + } + Ok(written) + } + + /// Builds a batch of `arrays`, reusing the last batch's schema while the columns keep their + /// types, so that a new schema means the column types changed. + fn make_batch( + &mut self, + arrays: Vec, + num_rows: usize, + ) -> Result { + let reusable = self.batch_schema.as_ref().is_some_and(|schema| { + schema + .fields() + .iter() + .zip(&arrays) + .all(|(field, array)| field.data_type() == array.data_type()) + }); + if !reusable { + let fields = arrays + .iter() + .enumerate() + .map(|(i, array)| Field::new(format!("c{i}"), array.data_type().clone(), true)) + .collect::>(); + self.batch_schema = Some(Arc::new(Schema::new(fields))); + } + let schema = Arc::clone( + self.batch_schema + .as_ref() + .expect("batch schema is set above"), + ); + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + RecordBatch::try_new_with_options(schema, arrays, &options) + } + + /// Writes the buffered blocks to `output` once they reach `OUTPUT_BUFFER_BYTES`. + fn write_if_full(&mut self, output: &mut impl Write) -> Result<(), CometError> { + if self.buffer.len() >= OUTPUT_BUFFER_BYTES { + self.write_buffered(output)?; + } + Ok(()) + } + + fn write_buffered(&mut self, output: &mut impl Write) -> Result<(), CometError> { + output.write_all(&self.buffer)?; + self.buffer.clear(); + Ok(()) + } + + fn encode_nanos(&self) -> i64 { + self.ipc_time.value() as i64 + } +} + /// Processes a sorted row partition and writes the result to the given output path. #[allow(clippy::too_many_arguments)] pub fn process_sorted_row_partition( @@ -1388,14 +1568,13 @@ pub fn process_sorted_row_partition( None }; - // Create builders once and reuse them across batches. - // After finish() is called, builders are reset and can be reused. - let mut data_builders: Vec> = Vec::with_capacity(schema.len()); - schema.iter().try_for_each(|dt| { - make_builders(dt, batch_size, prefer_dictionary_ratio) - .map(|builder| data_builders.push(builder))?; - Ok::<(), CometError>(()) - })?; + let mut encoder = SortedRowEncoder::try_new( + schema, + batch_size, + prefer_dictionary_ratio, + codec, + offsets_as_lengths, + )?; // Open the output file once and reuse it across batches let mut output_data = OpenOptions::new() @@ -1403,76 +1582,188 @@ pub fn process_sorted_row_partition( .append(true) .open(&output_path)?; - // Reusable buffer for serialized batch data - let mut frozen: Vec = Vec::new(); - - // Single ipc_time accumulates encode + compression time across all batches. - let ipc_time = Time::default(); - // One context for every batch this call encodes; the JVM calls in once per sorted - // partition, so there is no wider native scope to hoist it to. - let mut codec_context = ShuffleCodecContext::default(); - // The block writer pre-encodes the batch schema, so it is kept until the schema changes. - let mut block_writer: Option<(SchemaRef, ShuffleBlockWriter)> = None; - while current_row < row_num { let n = std::cmp::min(batch_size, row_num - current_row); + written += encoder.encode_batch( + row_addresses_ptr, + row_sizes_ptr, + current_row, + current_row + n, + current_checksum.as_mut(), + )?; + encoder.write_if_full(&mut output_data)?; + current_row += n; + } + encoder.write_buffered(&mut output_data)?; - // Appends rows to the array builders. - // For each column, iterating over rows and appending values to corresponding array - // builder. - for (idx, builder) in data_builders.iter_mut().enumerate() { - append_columns( - row_addresses_ptr, - row_sizes_ptr, - current_row, - current_row + n, - schema, - idx, - builder, - prefer_dictionary_ratio, - )?; - } + Ok(( + written as i64, + current_checksum.map(|c| c.finalize()), + encoder.encode_nanos(), + )) +} - // Writes a record batch generated from the array builders to the output file. - // Note: builder_to_array calls finish() which resets the builder, making it reusable for the next batch. - let array_refs: Result, _> = data_builders - .iter_mut() - .zip(schema.iter()) - .map(|(builder, datatype)| builder_to_array(builder, datatype, prefer_dictionary_ratio)) - .collect(); - let batch = make_batch(array_refs?, n)?; +/// Processes one sorted spill of the sort-based JVM shuffle writer: converts the rows of every +/// partition in it and appends the partitions' blocks to `output_path`, in partition order. +/// +/// `record_pointers` are Spark `PackedRecordPointer`s sorted by partition id, as +/// `ShuffleInMemorySorter` holds them. The page number in a pointer indexes `page_addresses`, +/// the base address of each page the sorter holds, and its offset locates the record in that +/// page. A record is stored as its length (`uao_size` bytes, written by `UnsafeAlignedOffset`), +/// then the row's size (4 bytes, written by `UnsafeRowSerializer`), then the row. +/// +/// Sets `partition_lengths[p]` to the bytes written for partition `p`. With a `checksum_algo`, +/// it also continues `partition_checksums[p]` over those bytes, where `i64::MIN` is a checksum not +/// yet started. Partitions without rows keep their values. Returns the encode and compression +/// time in nanoseconds. +/// +/// # Safety +/// Every pointer must resolve to a record in one of the pages, holding a readable row. +#[allow(clippy::too_many_arguments)] +pub unsafe fn process_sorted_partitions( + record_pointers: &[i64], + page_addresses: &[i64], + uao_size: usize, + schema: &[DataType], + output_path: &str, + prefer_dictionary_ratio: f64, + batch_size: usize, + checksum_algo: Option, + partition_checksums: &mut [i64], + partition_lengths: &mut [i64], + codec: &CompressionCodec, + offsets_as_lengths: bool, +) -> Result { + if uao_size != 4 && uao_size != 8 { + return Err(CometError::Internal(format!( + "Unsupported record length size {uao_size}" + ))); + } + if batch_size == 0 { + return Err(CometError::Internal( + "Batch size must be positive".to_string(), + )); + } + let num_partitions = partition_lengths.len(); + if checksum_algo.is_some() && partition_checksums.len() != num_partitions { + return Err(CometError::Internal(format!( + "{} partition checksums for {num_partitions} partitions", + partition_checksums.len() + ))); + } - frozen.clear(); - let mut cursor = Cursor::new(&mut frozen); + let mut encoder = SortedRowEncoder::try_new( + schema, + batch_size, + prefer_dictionary_ratio, + codec, + offsets_as_lengths, + )?; + let mut output = OpenOptions::new() + .create(true) + .append(true) + .open(output_path)?; + + // The rows of one batch, resolved from their record pointers. + let capacity = batch_size.min(record_pointers.len()); + let mut row_addresses: Vec = Vec::with_capacity(capacity); + let mut row_sizes: Vec = Vec::with_capacity(capacity); + + let mut start = 0; + let mut previous_partition = None; + while start < record_pointers.len() { + let partition = packed_partition_id(record_pointers[start]); + if partition >= num_partitions { + return Err(CometError::Internal(format!( + "Shuffle record for partition {partition} of {num_partitions}" + ))); + } + if previous_partition.is_some_and(|previous| partition <= previous) { + return Err(CometError::Internal( + "Shuffle records are not sorted by partition".to_string(), + )); + } + previous_partition = Some(partition); + let end = record_pointers[start..] + .iter() + .position(|&pointer| packed_partition_id(pointer) != partition) + .map_or(record_pointers.len(), |len| start + len); - // The schema changes between batches when a dictionary-encoded column falls back to a - // plain array (see `builder_to_array`). - if block_writer - .as_ref() - .is_none_or(|(cached, _)| cached != batch.schema_ref()) - { - let writer = ShuffleBlockWriter::try_new_with_offsets_as_lengths( - batch.schema_ref(), - codec.clone(), - offsets_as_lengths, + let mut checksum = checksum_algo + .map(|algo| { + let current = partition_checksums[partition]; + Checksum::try_new(algo, (current != i64::MIN).then_some(current as u32)) + }) + .transpose()?; + let mut written = 0; + for pointers in record_pointers[start..end].chunks(batch_size) { + row_addresses.clear(); + row_sizes.clear(); + for &pointer in pointers { + let (address, size) = resolve_record(pointer, page_addresses, uao_size)?; + row_addresses.push(address); + row_sizes.push(size); + } + written += encoder.encode_batch( + row_addresses.as_mut_ptr(), + row_sizes.as_mut_ptr(), + 0, + pointers.len(), + checksum.as_mut(), )?; - block_writer = Some((batch.schema(), writer)); + encoder.write_if_full(&mut output)?; } - let (_, writer) = block_writer.as_ref().expect("block writer is set above"); - written += writer.write_batch(&batch, &mut cursor, &mut codec_context, &ipc_time)?; - if let Some(checksum) = &mut current_checksum { - checksum.update(&mut cursor)?; + partition_lengths[partition] = written as i64; + if let Some(checksum) = checksum { + partition_checksums[partition] = checksum.finalize() as i64; } - - output_data.write_all(&frozen)?; - current_row += n; + start = end; } + encoder.write_buffered(&mut output)?; + + Ok(encoder.encode_nanos()) +} +fn packed_partition_id(pointer: i64) -> usize { + (pointer as u64 >> PACKED_PARTITION_ID_SHIFT) as usize +} + +/// Returns the address and size of the row the packed record pointer `pointer` points to. +/// +/// # Safety +/// If the pointer's page is in `page_addresses`, its record must be readable. +unsafe fn resolve_record( + pointer: i64, + page_addresses: &[i64], + uao_size: usize, +) -> Result<(jlong, jint), CometError> { + let pointer = pointer as u64; + let page_number = ((pointer >> PACKED_PAGE_NUMBER_SHIFT) & PACKED_PAGE_NUMBER_MASK) as usize; + let page_address = page_addresses + .get(page_number) + .copied() + .filter(|&address| address != 0) + .ok_or_else(|| { + CometError::Internal(format!( + "Shuffle record in page {page_number}, which the sorter does not hold" + )) + })?; + let record_address = page_address + (pointer & PACKED_OFFSET_MASK) as i64; + let record_length = if uao_size == 4 { + (record_address as *const i32).read_unaligned() + } else { + (record_address as *const i64).read_unaligned() as i32 + }; + let row_size = record_length - SERIALIZED_ROW_SIZE_BYTES as i32; + if row_size < 0 { + return Err(CometError::Internal(format!( + "Shuffle record of {record_length} bytes is too short to hold a row" + ))); + } Ok(( - written as i64, - current_checksum.map(|c| c.finalize()), - ipc_time.value() as i64, + record_address + uao_size as i64 + SERIALIZED_ROW_SIZE_BYTES, + row_size, )) } @@ -1526,17 +1817,6 @@ fn builder_to_array( } } -fn make_batch(arrays: Vec, row_count: usize) -> Result { - let fields = arrays - .iter() - .enumerate() - .map(|(i, array)| Field::new(format!("c{i}"), array.data_type().clone(), true)) - .collect::>(); - let schema = Arc::new(Schema::new(fields)); - let options = RecordBatchOptions::new().with_row_count(Option::from(row_count)); - RecordBatch::try_new_with_options(schema, arrays, &options) -} - #[cfg(test)] mod test { use arrow::array::AsArray; @@ -1686,4 +1966,45 @@ mod test { lists.iter().map(|l| l.map(<[_]>::to_vec)).collect(); assert_eq!(appended, expected); } + + #[test] + fn resolve_record_finds_the_row_a_packed_pointer_points_to() { + // A page holding two records the way Spark's sorter stores them: the record's length in + // `uao_size` bytes, then the row's 4-byte size, then a 16-byte row. + for uao_size in [4_usize, 8] { + let record_len = uao_size + 4 + 16; + let mut page = vec![0_u64; (2 * record_len).div_ceil(8)]; + let base = page.as_mut_ptr() as *mut u8; + for record in 0..2 { + let offset = record * record_len; + // SAFETY: both records lie inside `page`. + unsafe { + let length = base.add(offset); + if uao_size == 4 { + (length as *mut i32).write_unaligned(4 + 16); + } else { + (length as *mut i64).write_unaligned(4 + 16); + } + (base.add(offset + uao_size) as *mut i32).write_unaligned(16_i32.to_be()); + } + } + // The page is page number 1; page 0 is not held. + let page_addresses = [0, base as i64]; + + let pointer = (5_i64 << 40) | (1 << 27) | record_len as i64; + assert_eq!(packed_partition_id(pointer), 5); + let (address, size) = unsafe { resolve_record(pointer, &page_addresses, uao_size) } + .expect("record in a held page"); + assert_eq!(address, base as i64 + (record_len + uao_size + 4) as i64); + assert_eq!(size, 16); + + for missing_page in [0_i64, 2] { + let pointer = (5_i64 << 40) | (missing_page << 27); + let error = unsafe { resolve_record(pointer, &page_addresses, uao_size) } + .unwrap_err() + .to_string(); + assert!(error.contains(&format!("page {missing_page}")), "{error}"); + } + } + } } diff --git a/native/shuffle/src/writers/checksum.rs b/native/shuffle/src/writers/checksum.rs index 1dfd15eb373..dacb6def159 100644 --- a/native/shuffle/src/writers/checksum.rs +++ b/native/shuffle/src/writers/checksum.rs @@ -16,12 +16,10 @@ // under the License. use crate::spark_crc32c_hasher::SparkCrc32cHasher; -use bytes::Buf; use datafusion_comet_jni_bridge::errors::{CometError, CometResult}; use simd_adler32::Adler32; use std::default::Default; use std::hash::Hasher; -use std::io::{Cursor, SeekFrom}; /// Checksum algorithms for writing IPC bytes. #[derive(Clone)] @@ -69,23 +67,11 @@ impl Checksum { } } - pub(crate) fn update(&mut self, cursor: &mut Cursor<&mut Vec>) -> CometResult<()> { + pub(crate) fn update(&mut self, bytes: &[u8]) { match self { - Checksum::CRC32(hasher) => { - std::io::Seek::seek(cursor, SeekFrom::Start(0))?; - hasher.update(cursor.chunk()); - Ok(()) - } - Checksum::Adler32(hasher) => { - std::io::Seek::seek(cursor, SeekFrom::Start(0))?; - hasher.write(cursor.chunk()); - Ok(()) - } - Checksum::CRC32C(hasher) => { - std::io::Seek::seek(cursor, SeekFrom::Start(0))?; - hasher.write(cursor.chunk()); - Ok(()) - } + Checksum::CRC32(hasher) => hasher.update(bytes), + Checksum::Adler32(hasher) => hasher.write(bytes), + Checksum::CRC32C(hasher) => hasher.write(bytes), } } @@ -101,17 +87,13 @@ impl Checksum { #[cfg(test)] mod tests { use super::*; - use std::io::Cursor; #[test] fn test_crc32() { let mut checksum = Checksum::try_new(0, None).unwrap(); let message = b"123456789"; - let mut vector: Vec = message.to_vec(); - let mut buff = Cursor::new(&mut vector); - - checksum.update(&mut buff).unwrap(); + checksum.update(message); let result = checksum.finalize(); let expected_crc = 0xcbf43926u32; @@ -123,10 +105,7 @@ mod tests { let mut checksum = Checksum::try_new(1, None).unwrap(); let message = b"123456789"; - let mut vector: Vec = message.to_vec(); - let mut buff = Cursor::new(&mut vector); - - checksum.update(&mut buff).unwrap(); + checksum.update(message); let result = checksum.finalize(); let expected_crc = 0x091e01deu32; @@ -138,10 +117,7 @@ mod tests { let mut checksum = Checksum::try_new(2, None).unwrap(); let message = b"123456789"; - let mut vector: Vec = message.to_vec(); - let mut buff = Cursor::new(&mut vector); - - checksum.update(&mut buff).unwrap(); + checksum.update(message); let result = checksum.finalize(); let expected_crc = 0xe3069283u32; diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java index 4ce62af6e8c..e437acd3035 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java @@ -36,9 +36,10 @@ * serialized rows when Spark runs in on-heap mode. * *

Spark's own allocator cannot be used here. `TaskMemoryManager.allocatePage` hands out pages - * from `tungstenMemoryAllocator`, which is the on-heap allocator in this mode, and the row - * addresses derived from these pages are passed to `writeSortedFileNative` for Rust to dereference. - * The pages therefore have to be `Unsafe`-allocated regardless of Spark's memory mode. + * from `tungstenMemoryAllocator`, which is the on-heap allocator in this mode, and Rust reads the + * rows straight from these pages, through the row addresses passed to `writeSortedFileNative` or + * the page addresses passed to `writeSortedPartitionsNative`. The pages therefore have to be + * `Unsafe`-allocated regardless of Spark's memory mode. * *

Nothing bounds these allocations. On-heap mode exists so that the Spark SQL tests can run * against Comet without changing Spark's memory configuration; it is not a production diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java index 23368115ae9..2965815b360 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java +++ b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java @@ -33,7 +33,9 @@ import org.apache.spark.unsafe.Platform; import org.apache.spark.unsafe.UnsafeAlignedOffset; import org.apache.spark.unsafe.array.LongArray; +import org.apache.spark.unsafe.memory.MemoryBlock; +import org.apache.comet.CometConf; import org.apache.comet.Native; /** @@ -58,7 +60,8 @@ public interface SpillCallback { private final String compressionCodec; private final int compressionLevel; private final boolean offsetsAsLengths; - private final String checksumAlgorithm; + // The checksum algorithm's id for native code, or -1 when shuffle checksums are disabled. + private final int checksumAlgoId; // Shared state (mutable, passed by reference from parent) private final long[] partitionChecksums; @@ -113,7 +116,8 @@ public SpillSorter( this.compressionCodec = compressionCodec; this.compressionLevel = compressionLevel; this.offsetsAsLengths = offsetsAsLengths; - this.checksumAlgorithm = checksumAlgorithm; + this.checksumAlgoId = + partitionChecksums.length > 0 ? checksumAlgorithmId(checksumAlgorithm) : -1; this.partitionChecksums = partitionChecksums; this.writeMetrics = writeMetrics; this.taskContext = taskContext; @@ -235,19 +239,22 @@ public int numRecords() { return this.inMemSorter.numRecords(); } + /** + * Sorts the buffered records by partition and writes them to the spill file set by {@link + * #setSpillInfo}, every partition in one native call, recording each partition's length and + * checksum. + */ public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) throws IOException { - // This call performs the actual sort. - long arrayAddr = this.sorterArray.getBaseOffset(); - int pos = inMemSorter.numRecords(); - nativeLib.sortRowPartitionsNative(arrayAddr, pos, tracingEnabled); - ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords = - new ShuffleInMemorySorter.ShuffleSorterIterator(pos, this.sorterArray, 0); - - // If there are no sorted records, so we don't need to create an empty spill file. - if (!sortedRecords.hasNext()) { + final int numRecords = inMemSorter.numRecords(); + // If there are no records, we don't need to create an empty spill file. + if (numRecords == 0) { return; } + // This call performs the actual sort. + long arrayAddr = this.sorterArray.getBaseOffset(); + nativeLib.sortRowPartitionsNative(arrayAddr, numRecords, tracingEnabled); + final ShuffleWriteMetricsReporter writeMetricsToUse; if (isLastFile) { @@ -260,81 +267,42 @@ public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) th writeMetricsToUse = new ShuffleWriteMetrics(); } - int currentPartition = -1; - - final RowPartition rowPartition = new RowPartition(initialSize); - - while (sortedRecords.hasNext()) { - sortedRecords.loadNext(); - final int partition = sortedRecords.packedRecordPointer.getPartitionId(); - assert (partition >= currentPartition); - if (partition != currentPartition) { - // Switch to the new partition - if (currentPartition != -1) { - - if (partitionChecksums.length > 0) { - // If checksum is enabled, we need to update the checksum for the current partition. - setChecksum(partitionChecksums[currentPartition]); - setChecksumAlgo(checksumAlgorithm); - } - - long written = - doSpilling( - dataTypes, - spillInfo.file, - rowPartition, - writeMetricsToUse, - preferDictionaryRatio, - compressionCodec, - compressionLevel, - offsetsAsLengths, - tracingEnabled); - spillInfo.partitionLengths[currentPartition] = written; - - // Store the checksum for the current partition. - partitionChecksums[currentPartition] = getChecksum(); - } - currentPartition = partition; - } - - final long recordPointer = sortedRecords.packedRecordPointer.getRecordPointer(); - final long recordOffsetInPage = allocator.getOffsetInPage(recordPointer); - // Note that we need to skip over record key (partition id) - // Note that we already use off-heap memory for serialized rows, so recordPage is always - // null. - int recordSizeInBytes = UnsafeAlignedOffset.getSize(null, recordOffsetInPage) - 4; - long recordReadPosition = recordOffsetInPage + uaoSize + 4; // skip over record length too - rowPartition.addRow(recordReadPosition, recordSizeInBytes); + // Native code resolves the sorted record pointers through the page table, so the rows are + // read where they lie, and continues each partition's checksum over the bytes it writes. + long start = System.nanoTime(); + long encodeNanos = + nativeLib.writeSortedPartitionsNative( + arrayAddr, + numRecords, + pageAddresses(), + uaoSize, + dataTypes, + spillInfo.file.getAbsolutePath(), + preferDictionaryRatio, + CometConf.jvmShuffleBatchSize(), + checksumAlgoId, + partitionChecksums, + spillInfo.partitionLengths, + compressionCodec, + compressionLevel, + offsetsAsLengths, + tracingEnabled); + addEncodeNanos(encodeNanos); + + long written = 0; + for (long length : spillInfo.partitionLengths) { + written += length; } - if (currentPartition != -1) { - if (partitionChecksums.length > 0) { - // If checksum is enabled, we need to update the checksum for the last partition. - setChecksum(partitionChecksums[currentPartition]); - setChecksumAlgo(checksumAlgorithm); - } - - long written = - doSpilling( - dataTypes, - spillInfo.file, - rowPartition, - writeMetricsToUse, - preferDictionaryRatio, - compressionCodec, - compressionLevel, - offsetsAsLengths, - tracingEnabled); - spillInfo.partitionLengths[currentPartition] = written; - - // Store the checksum for the last partition. - if (partitionChecksums.length > 0) { - partitionChecksums[currentPartition] = getChecksum(); - } + // Other threads may be updating the metrics at the same time, so we need to synchronize it. + synchronized (writeMetricsToUse) { + writeMetricsToUse.incWriteTime(System.nanoTime() - start); + writeMetricsToUse.incRecordsWritten(numRecords); + writeMetricsToUse.incBytesWritten(written); + } - synchronized (spills) { - spills.add(spillInfo); - } + synchronized (spills) { + spills.add(spillInfo); } if (!isLastFile) { // i.e. this is a spill file @@ -367,6 +335,24 @@ public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) th } } + /** + * The base address of each page holding records, indexed by page number, for native code to + * resolve the record pointers through. Pages this sorter does not hold are 0. + */ + private long[] pageAddresses() { + synchronized (this) { + int maxPageNumber = -1; + for (MemoryBlock page : allocatedPages) { + maxPageNumber = Math.max(maxPageNumber, page.pageNumber); + } + long[] addresses = new long[maxPageNumber + 1]; + for (MemoryBlock page : allocatedPages) { + addresses[page.pageNumber] = page.getBaseOffset(); + } + return addresses; + } + } + public boolean hasSpaceForAnotherRecord() { return inMemSorter.hasSpaceForAnotherRecord(); } diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java index dc67838c813..48cc2d41f9b 100644 --- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java +++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java @@ -75,8 +75,8 @@ public abstract class SpillWriter { /** * Optional shared accumulator for encode + compression time in nanoseconds. When set, every - * doSpilling() call adds the native encode time (results[2]) to this accumulator. Multiple - * SpillWriter instances can share the same AtomicLong so their times aggregate correctly. + * native write adds the encode time it reports to this accumulator. Multiple SpillWriter + * instances can share the same AtomicLong so their times aggregate correctly. */ @Nullable private AtomicLong encodeNanosAcc = null; @@ -84,6 +84,13 @@ public void setEncodeNanosAccumulator(AtomicLong acc) { this.encodeNanosAcc = acc; } + /** Adds native encode + compression time to the shared accumulator, if one is set. */ + protected void addEncodeNanos(long nanos) { + if (encodeNanosAcc != null) { + encodeNanosAcc.addAndGet(nanos); + } + } + /** Serialize row schema to byte array. */ protected byte[][] serializeSchema(StructType schema) { byte[][] dataTypes = new byte[schema.length()][]; @@ -104,14 +111,19 @@ protected byte[][] serializeSchema(StructType schema) { } protected void setChecksumAlgo(String checksumAlgo) { + this.checksumAlgo = checksumAlgorithmId(checksumAlgo); + } + + /** The id native code knows the Spark shuffle checksum algorithm `checksumAlgo` by. */ + protected static int checksumAlgorithmId(String checksumAlgo) { String algo = checksumAlgo.toLowerCase(Locale.ROOT); if (algo.equals("crc32")) { - this.checksumAlgo = 0; + return 0; } else if (algo.equals("adler32")) { - this.checksumAlgo = 1; + return 1; } else if (algo.equals("crc32c")) { - this.checksumAlgo = 2; + return 2; } else { throw new UnsupportedOperationException( "Unsupported shuffle checksum algorithm: " + checksumAlgo); @@ -217,9 +229,7 @@ protected long doSpilling( long written = results[0]; checksum = results[1]; // results[2] = encode + compression time in nanoseconds (measured by native code) - if (encodeNanosAcc != null) { - encodeNanosAcc.addAndGet(results[2]); - } + addEncodeNanos(results[2]); rowPartition.reset(); diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index c12a85dee93..efeabf079a7 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -139,7 +139,7 @@ class Native extends NativeBase { @native def releasePlan(plan: Long): Unit /** - * Used by Comet shuffle external sorter to write sorted records to disk. + * Used by the bypass merge sort shuffle writer to write one partition's records to disk. * * @param addresses * the array of addresses of Spark unsafe rows. @@ -189,6 +189,62 @@ class Native extends NativeBase { tracingEnabled: Boolean): Array[Long] // scalastyle:on + /** + * Used by Comet shuffle external sorter to write one sorted spill to disk: the records of every + * partition, in partition order, in one call. + * + * @param recordPointersAddr + * the address of the sorter's packed record pointers, sorted by partition id. + * @param numRecords + * the number of record pointers. + * @param pageAddresses + * the base address of each page holding records, indexed by the page number in the record + * pointers. Pages the sorter does not hold are 0. + * @param uaoSize + * the size of the record length stored before each record, `UnsafeAlignedOffset.getUaoSize`. + * @param datatypes + * the datatypes of fields in Spark unsafe rows. + * @param file + * the file path to append to. + * @param preferDictionaryRatio + * as in `writeSortedFileNative`. + * @param batchSize + * the maximum number of rows in each batch written. + * @param checksumAlgo + * the checksum algorithm to use: 0 for CRC32, 1 for Adler32, 2 for CRC32C, or -1 for none. + * @param partitionChecksums + * each partition's checksum, updated for every partition written. `Long.MIN_VALUE` is a + * checksum not computed yet. Empty without a checksum. + * @param partitionLengths + * set to the number of bytes written for each partition. + * @param compressionCodec + * the compression codec + * @param compressionLevel + * the compression level + * @param offsetsAsLengths + * whether to store offsets as lengths where the codec takes it + * @return + * the time spent encoding and compressing, in nanoseconds + */ + // scalastyle:off + @native def writeSortedPartitionsNative( + recordPointersAddr: Long, + numRecords: Int, + pageAddresses: Array[Long], + uaoSize: Int, + datatypes: Array[Array[Byte]], + file: String, + preferDictionaryRatio: Double, + batchSize: Int, + checksumAlgo: Int, + partitionChecksums: Array[Long], + partitionLengths: Array[Long], + compressionCodec: String, + compressionLevel: Int, + offsetsAsLengths: Boolean, + tracingEnabled: Boolean): Long + // scalastyle:on + /** * Sorts partition ids of Spark unsafe rows in place. Used by Comet shuffle external sorter. * diff --git a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala index d2d67941262..19d3d149999 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala @@ -19,7 +19,14 @@ package org.apache.spark.shuffle.sort +import java.io.File +import java.nio.ByteBuffer +import java.nio.file.Files +import java.util.UUID import java.util.concurrent.atomic.AtomicInteger +import java.util.zip.Adler32 + +import scala.collection.mutable import org.scalatest.BeforeAndAfterEach import org.scalatest.funsuite.AnyFunSuite @@ -28,7 +35,10 @@ import org.apache.spark.{SparkConf, TaskContext} import org.apache.spark.executor.ShuffleWriteMetrics import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.shuffle.comet.{CometShuffleMemoryAllocator, CometShuffleMemoryAllocatorTrait} +import org.apache.spark.sql.catalyst.expressions.UnsafeRow +import org.apache.spark.sql.comet.execution.shuffle.{SpillInfo => CometSpillInfo} import org.apache.spark.sql.types._ +import org.apache.spark.storage.TempShuffleBlockId import org.apache.spark.unsafe.Platform import org.apache.spark.unsafe.UnsafeAlignedOffset @@ -267,6 +277,127 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { } } + test("a sorted spill writes every partition in one native call") { + // Rows of `createTestSchema` in partitions 0, 2, 3 and 5 of 7, over 1 KiB pages so that they + // span several pages. + val numPartitions = 7 + def partitionOf(i: Int): Int = Array(0, 2, 5, 3, 0, 2)(i % 6) + val records = (0 until 300).map { i => + val row = new UnsafeRow(1) + row.pointTo(new Array[Byte](16), 16) + row.setInt(0, i) + // A record as `UnsafeRowSerializer` writes it: the row's size, then the row. + (ByteBuffer.allocate(4 + 16).putInt(16).put(row.getBytes).array(), partitionOf(i)) + } + val usedPartitions = records.map(_._2).toSet + + def check(allocator: CometShuffleMemoryAllocatorTrait, checksumAlgorithm: String): Unit = { + val partitionChecksums = + if (checksumAlgorithm == null) Array.emptyLongArray + else Array.fill(numPartitions)(Long.MinValue) + // Renamed on import: Spark 3.x has a `SpillInfo` of its own in this package. + val spills = new java.util.LinkedList[CometSpillInfo]() + val writeMetrics = new ShuffleWriteMetrics() + val taskContext = TaskContext.empty() + val sorter = new SpillSorter( + allocator, + INITIAL_SIZE, + createTestSchema(), + UAO_SIZE, + 10.0, // preferDictionaryRatio + "lz4", // compressionCodec + 1, // compressionLevel + true, // offsetsAsLengths + checksumAlgorithm, + partitionChecksums, + writeMetrics, + taskContext, + spills, + () => {}) + + // A spill, then the final file, with the same records. + val files = mutable.ArrayBuffer[File]() + try { + Seq(false, true).foreach { isLastFile => + records.foreach { case (record, partition) => + assert(sorter.acquireNewPageIfNecessary(record.length + UAO_SIZE)) + sorter.insertRecord( + record, + Platform.BYTE_ARRAY_OFFSET.toLong, + record.length, + partition) + } + val file = File.createTempFile("comet-spill-sorter", ".data") + files += file + sorter.setSpillInfo( + new CometSpillInfo(numPartitions, file, TempShuffleBlockId(UUID.randomUUID()))) + sorter.writeSortedFileNative(isLastFile, false) + sorter.freeMemory() + if (!isLastFile) { + sorter.reset() + } + } + + assert(spills.size === 2) + val written = Seq(spills.get(0), spills.get(1)) + written.foreach { spill => + assert(spill.partitionLengths.sum === spill.file.length()) + (0 until numPartitions).foreach { p => + assert((spill.partitionLengths(p) > 0) === usedPartitions(p), s"partition $p") + } + } + + if (checksumAlgorithm != null) { + // A partition's checksum covers its bytes in every file, as Spark checksums the + // merged partition. Partitions without rows keep the unset value. + (0 until numPartitions).foreach { p => + if (usedPartitions(p)) { + val adler32 = new Adler32() + written.foreach { spill => + val bytes = Files.readAllBytes(spill.file.toPath) + val start = spill.partitionLengths.take(p).sum.toInt + adler32.update(bytes, start, spill.partitionLengths(p).toInt) + } + assert(partitionChecksums(p) === adler32.getValue, s"partition $p") + } else { + assert(partitionChecksums(p) === Long.MinValue, s"partition $p") + } + } + } + + // Both writes count their records, but only the final file counts as shuffle bytes. + assert(writeMetrics.recordsWritten === 2 * records.size) + assert(writeMetrics.bytesWritten === written(1).file.length()) + assert(taskContext.taskMetrics.diskBytesSpilled === written(0).file.length()) + } finally { + sorter.freeMemory() + sorter.freeArray() + files.foreach(_.delete()) + } + } + + val pageSize = 1024L + // On-heap mode: the allocator numbers pages in its own table. + check(CometShuffleMemoryAllocator.getInstance(taskMemoryManager, pageSize), "adler32") + // Shuffle checksums disabled: no checksum array and no algorithm. + check(CometShuffleMemoryAllocator.getInstance(taskMemoryManager, pageSize), null) + + // Off-heap mode: pages are numbered in the TaskMemoryManager's table. + val offHeapConf = new SparkConf(false) + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "64m") + val offHeapMemoryManager = new TestMemoryManager(offHeapConf) + offHeapMemoryManager.limit(64L * 1024 * 1024) + val offHeapTaskMemoryManager = new TaskMemoryManager(offHeapMemoryManager, 0) + try { + check( + CometShuffleMemoryAllocator.getInstance(offHeapTaskMemoryManager, pageSize), + "adler32") + } finally { + assert(offHeapTaskMemoryManager.cleanUpAllAllocatedMemory() === 0L) + } + } + test("pointer array growth is sized from the array, not from the data pages") { // Use the unified (off-heap) allocator so that every allocation made by this sorter is // visible through `getUsed` on an allocator that is private to this test. diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala new file mode 100644 index 00000000000..bff4405688f --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala @@ -0,0 +1,119 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * Measures Comet's JVM columnar shuffle at partition counts above Spark's + * `spark.shuffle.sort.bypassMergeThreshold` (200), where the sort-based writer is used. That + * writer sorts its buffered rows by partition and writes every partition holding a row, once per + * spill and once more at the end, so its cost per partition is paid `partitions x (spills + 1)` + * times. Each case hash-partitions the same rows and discards the result. + * + * The partition count is set through `spark.sql.shuffle.partitions` rather than `repartition(n, + * ...)`, so that AQE coalesces the read side into a few tasks and the measurement is not + * dominated by scheduling one reduce task per partition. + * + * `spark.comet.shuffle.jvm.spillThreshold` makes Comet spill a fixed number of rows at a time. It + * does not apply to Spark's shuffle, which does not spill in these cases. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometColumnarShuffleWriteBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometColumnarShuffleWriteBenchmark-**results.txt". + */ +object CometColumnarShuffleWriteBenchmark extends CometBenchmarkBase { + + private val numRows = 1024 * 1024 + private val spillThreshold = 100000 + + private def cometJvmShuffle(extraConfs: (String, String)*): Seq[(String, String)] = + Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") ++ extraConfs + + private def query(): DataFrame = spark.table("parquetV1Table").repartition(col("c1")) + + private def shuffleWrite(partitions: Int): Unit = { + val benchmark = + new Benchmark(s"Shuffle write, $partitions partitions", numRows, output = output) + + val shufflePartitions = SQLConf.SHUFFLE_PARTITIONS.key -> partitions.toString + val cases = Seq( + "Spark" -> Seq(CometConf.COMET_ENABLED.key -> "false"), + "Comet JVM shuffle" -> cometJvmShuffle(), + s"Comet JVM shuffle, spill every $spillThreshold rows" -> cometJvmShuffle( + CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> spillThreshold.toString)) + + cases.foreach { case (name, caseConfs) => + val confs = caseConfs :+ shufflePartitions + if (name.startsWith("Comet")) { + // Check the case measures the JVM columnar shuffle before timing it. + val plans = capturePlans { + withSQLConf(confs: _*) { + query().noop() + } + } + val shuffles = collectAcross(plans) { + case s: CometShuffleExchangeExec if s.shuffleType == CometColumnarShuffle => s + } + assert(shuffles.nonEmpty, s"$name did not use the JVM columnar shuffle: $plans") + } + + benchmark.addCase(name) { _ => + withSQLConf(confs: _*) { + query().noop() + } + } + } + + benchmark.run() + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + runBenchmarkWithTable("Sort-based JVM columnar shuffle write", numRows) { _ => + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.sql(s"""SELECT + | value AS c1, + | CAST(value % 1000 AS INT) AS c2, + | CAST(value AS DOUBLE) / 3 AS c3, + | CONCAT('status-', CAST(ABS(value % 100) AS STRING)) AS c4, + | DATE_ADD(DATE '2020-01-01', CAST(ABS(value % 1000) AS INT)) AS c5 + |FROM $tbl""".stripMargin)) + + Seq(1000, 10000).foreach(shuffleWrite) + } + } + } + } +} From 20a9dc3e5eb22c6a045414f1f98c81b2fe5ca28d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 15:59:42 -0600 Subject: [PATCH 2/5] test: pin dictionary changes and decoded rows in the sorted spill tests 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. --- docs/source/contributor-guide/jvm_shuffle.md | 14 +- native/core/src/execution/jni_api.rs | 128 +++++++++++++++++- native/shuffle/benches/row_columnar.rs | 3 +- .../spark/shuffle/sort/SpillSorterSuite.scala | 37 ++++- .../CometColumnarShuffleWriteBenchmark.scala | 2 +- 5 files changed, 168 insertions(+), 16 deletions(-) diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index e64e714d5ce..94125e7ea88 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -175,15 +175,15 @@ Selection logic in `CometShuffleManager.shouldBypassMergeSort()`: 1. `ShuffleWriteProcessor` calls `CometShuffleManager.getWriter()` 2. Writer receives `Iterator[Product2[K, V]]` where V is `UnsafeRow` 3. Rows are serialized and buffered in off-heap memory pages -4. When memory threshold or batch size is reached, the buffered rows are written: `CometDiskBlockWriter` - calls `SpillWriter.doSpilling()` for its partition, and `CometShuffleExternalSorter` calls - `SpillSorter.writeSortedFileNative()` for every partition +4. When memory threshold or batch size is reached, the buffered rows are written. `CometDiskBlockWriter` + calls `SpillWriter.doSpilling()` for its partition. `CometShuffleExternalSorter` calls + `SpillSorter.writeSortedFileNative()` once for each spill and once more for the final file, and each + call writes every partition 5. Native code converts the rows to Arrow arrays and writes them as Arrow IPC blocks. `doSpilling()` passes one partition's row addresses and sizes to `Native.writeSortedFileNative()`. `SpillSorter` - sorts its record pointers by partition with `Native.sortRowPartitionsNative()` and passes them, with - the address of each data page, to `Native.writeSortedPartitionsNative()`, which reads the rows in - place, writes every partition of the spill in one call, and returns each partition's length and - checksum + sorts its record pointers by partition with `Native.sortRowPartitionsNative()`, then passes them, + with the address of each data page, to `Native.writeSortedPartitionsNative()`. That call reads the + rows in place, writes every partition of the spill, and returns each partition's length and checksum 6. For bypass writer: partition files are concatenated into final output 7. For sort writer: spill files are merged diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 4def25abc6e..939993623c9 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -3610,9 +3610,9 @@ mod tests { /// lookup that ignored the page number would read the wrong page. const SORTER_PAGE_NUMBERS: [usize; 2] = [3, 1]; - /// 16-byte rows stored the way Spark's sorter stores them: in pages, each record its length - /// in 4 bytes, then the row's 4-byte big-endian size from `UnsafeRowSerializer`, then the - /// row. Rows alternate between two pages, and `pointers` holds the packed record pointer of + /// Rows stored the way Spark's sorter stores them: in pages, each record its length in 4 + /// bytes, then the row's 4-byte big-endian size from `UnsafeRowSerializer`, then the row. + /// Rows alternate between two pages, and `pointers` holds the packed record pointer of /// each row, sorted by partition as `sort_row_partitions` sorts them. struct SorterRecords { _pages: Vec>, @@ -3621,10 +3621,13 @@ mod tests { } impl SorterRecords { - fn new(rows: &[[u8; 16]], partitions: &[usize]) -> Self { + /// Each row must be whole 8-byte words, as an `UnsafeRow` is, so that every row stays + /// aligned in its page. + fn new>(rows: &[R], partitions: &[usize]) -> Self { let mut page_bytes: Vec> = vec![Vec::new(); SORTER_PAGE_NUMBERS.len()]; let mut pointers = Vec::new(); for (i, (row, &partition)) in rows.iter().zip(partitions).enumerate() { + let row = row.as_ref(); let page = i % SORTER_PAGE_NUMBERS.len(); let offset = page_bytes[page].len(); page_bytes[page].extend_from_slice(&(4 + row.len() as i32).to_ne_bytes()); @@ -3768,6 +3771,121 @@ mod tests { } } + /// Spark `UnsafeRow`s with one non-null string field each: an 8-byte null bitset, the + /// field's offset and size, then the string padded to a word. + fn string_rows(values: &[&str]) -> Vec> { + values + .iter() + .map(|value| { + let mut row = vec![0_u8; 16 + value.len().div_ceil(8) * 8]; + row[8..16].copy_from_slice(&(16_i64 << 32 | value.len() as i64).to_le_bytes()); + row[16..16 + value.len()].copy_from_slice(value.as_bytes()); + row + }) + .collect() + } + + #[test] + fn write_sorted_partitions_matches_writing_each_partition_as_dictionaries_change() { + // With a dictionary ratio of 1.5, a 2-row batch of two equal strings stays + // dictionary-encoded and one of two different strings is written plain + // (`builder_to_array`). So the block schema changes within partitions 0 and 3 and from + // partition 0 to 2, and stays the same from partition 2 to 3, where the encoder keeps + // its block writer. + let inserted = [ + ("a", 0), + ("e", 3), + ("a", 0), + ("d", 2), + ("e", 3), + ("b", 0), + ("d", 2), + ("f", 3), + ("c", 0), + ("g", 3), + ]; + let (values, partitions): (Vec<&str>, Vec) = inserted.into_iter().unzip(); + let rows = string_rows(&values); + let records = SorterRecords::new(&rows, &partitions); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("spill"); + let num_partitions = 4; + let mut checksums = vec![i64::MIN; num_partitions]; + let mut lengths = vec![0_i64; num_partitions]; + unsafe { + write_sorted_partitions( + &records.pointers, + &records.page_addresses, + 4, + &[DataType::Utf8], + path.to_str().unwrap(), + 1.5, + 2, + 1, + &mut checksums, + &mut lengths, + "lz4", + 1, + false, + ) + } + .unwrap(); + + // Writing each partition's rows on its own, with an encoder of its own. + let mut expected = Vec::new(); + for partition in 0..num_partitions { + let (addresses, sizes): (Vec, Vec) = rows + .iter() + .zip(&partitions) + .filter(|(_, &p)| p == partition) + .map(|(row, _)| (row.as_ptr() as i64, row.len() as i32)) + .unzip(); + if addresses.is_empty() { + assert_eq!(lengths[partition], 0, "partition {partition}"); + continue; + } + let partition_path = dir.path().join(format!("partition-{partition}")); + let [written, checksum, _] = unsafe { + write_sorted_file( + &addresses, + &sizes, + &[DataType::Utf8], + partition_path.to_str().unwrap().to_string(), + 1.5, + 2, + true, + 1, + i64::MIN, + "lz4", + 1, + false, + ) + } + .unwrap(); + assert_eq!(lengths[partition], written, "partition {partition}"); + assert_eq!(checksums[partition], checksum, "partition {partition}"); + expected.extend(std::fs::read(&partition_path).unwrap()); + } + let data = std::fs::read(&path).unwrap(); + assert_eq!(data, expected); + + let mut dictionary_blocks = Vec::new(); + let mut decoded = Vec::new(); + for block in shuffle_blocks(&data) { + let batch = read_ipc_compressed(block).unwrap(); + let column = batch.column(0); + dictionary_blocks.push(matches!(column.data_type(), DataType::Dictionary(_, _))); + let strings = arrow::compute::cast(column, &DataType::Utf8).unwrap(); + let strings = arrow::array::AsArray::as_string::(&strings); + decoded.extend(strings.iter().map(|s| s.unwrap().to_string())); + } + assert_eq!(dictionary_blocks, [true, false, true, true, false]); + let mut expected_values = inserted.to_vec(); + expected_values.sort_by_key(|&(_, partition)| partition); + let expected_values: Vec<&str> = expected_values.iter().map(|&(v, _)| v).collect(); + assert_eq!(decoded, expected_values); + } + #[test] fn write_sorted_partitions_continues_partition_checksums() { let rows = long_rows(&(0..23).collect::>()); @@ -3808,7 +3926,7 @@ mod tests { #[test] fn write_sorted_partitions_without_records() { - let records = SorterRecords::new(&[], &[]); + let records = SorterRecords::new::<[u8; 16]>(&[], &[]); let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("empty"); let mut checksums = vec![i64::MIN; 3]; diff --git a/native/shuffle/benches/row_columnar.rs b/native/shuffle/benches/row_columnar.rs index 5687281bc96..fef3b4e8fa5 100644 --- a/native/shuffle/benches/row_columnar.rs +++ b/native/shuffle/benches/row_columnar.rs @@ -19,7 +19,8 @@ //! //! Measures `process_sorted_row_partition()` performance for converting Spark //! UnsafeRow data to Arrow arrays, covering primitive, struct (flat/nested), -//! list, and map types. +//! list, and map types, and compares it with `process_sorted_partitions()` on +//! one sorted spill of the sort-based writer. use arrow::datatypes::{DataType as ArrowDataType, Field, Fields}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; diff --git a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala index 19d3d149999..7adacf7cef3 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.shuffle.sort -import java.io.File +import java.io.{ByteArrayInputStream, File} import java.nio.ByteBuffer import java.nio.file.Files import java.util.UUID @@ -36,12 +36,16 @@ import org.apache.spark.executor.ShuffleWriteMetrics import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.shuffle.comet.{CometShuffleMemoryAllocator, CometShuffleMemoryAllocatorTrait} import org.apache.spark.sql.catalyst.expressions.UnsafeRow -import org.apache.spark.sql.comet.execution.shuffle.{SpillInfo => CometSpillInfo} +import org.apache.spark.sql.comet.execution.shuffle.{NativeBatchDecoderIterator, SpillInfo => CometSpillInfo} +import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types._ import org.apache.spark.storage.TempShuffleBlockId import org.apache.spark.unsafe.Platform import org.apache.spark.unsafe.UnsafeAlignedOffset +import org.apache.comet.Native +import org.apache.comet.vector.NativeUtil + /** * Unit tests for [[SpillSorter]]. * @@ -85,6 +89,23 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { new StructType().add("id", IntegerType) } + /** Decodes the shuffle blocks in `bytes`, rows of `createTestSchema`, into their ids. */ + private def decodeIds(bytes: Array[Byte]): Seq[Int] = { + val nativeUtil = new NativeUtil + val decoder = NativeBatchDecoderIterator( + new ByteArrayInputStream(bytes), + new SQLMetric("decodeTime"), + new Native, + nativeUtil, + tracingEnabled = false) + try { + decoder.flatMap(batch => (0 until batch.numRows()).map(batch.column(0).getInt)).toList + } finally { + decoder.close() + nativeUtil.close() + } + } + private def createSpillSorter( spillCallback: SpillSorter.SpillCallback = () => {}, spills: java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo] = @@ -347,6 +368,18 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { } } + // Native code found each record through the page table: in both files, every partition + // decodes to its rows, in the order they were inserted. + written.foreach { spill => + val bytes = Files.readAllBytes(spill.file.toPath) + (0 until numPartitions).foreach { p => + val start = spill.partitionLengths.take(p).sum.toInt + val end = start + spill.partitionLengths(p).toInt + val ids = records.indices.filter(i => records(i)._2 == p) + assert(decodeIds(bytes.slice(start, end)) === ids, s"partition $p") + } + } + if (checksumAlgorithm != null) { // A partition's checksum covers its bytes in every file, as Spark checksums the // merged partition. Partitions without rows keep the unset value. diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala index bff4405688f..53138220045 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala @@ -63,7 +63,7 @@ object CometColumnarShuffleWriteBenchmark extends CometBenchmarkBase { private def shuffleWrite(partitions: Int): Unit = { val benchmark = - new Benchmark(s"Shuffle write, $partitions partitions", numRows, output = output) + new Benchmark(s"Shuffle write, $partitions partitions", numRows.toLong, output = output) val shufflePartitions = SQLConf.SHUFFLE_PARTITIONS.key -> partitions.toString val cases = Seq( From a0f90de2520969b2cce2bf2e10e13f1d90b4fb6c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 16:06:50 -0600 Subject: [PATCH 3/5] test: run the dictionary change test with zstd offsets as lengths 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. --- native/core/src/execution/jni_api.rs | 137 ++++++++++++++------------- 1 file changed, 71 insertions(+), 66 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 939993623c9..50a6157a2ab 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -3807,83 +3807,88 @@ mod tests { let (values, partitions): (Vec<&str>, Vec) = inserted.into_iter().unzip(); let rows = string_rows(&values); let records = SorterRecords::new(&rows, &partitions); - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("spill"); let num_partitions = 4; - let mut checksums = vec![i64::MIN; num_partitions]; - let mut lengths = vec![0_i64; num_partitions]; - unsafe { - write_sorted_partitions( - &records.pointers, - &records.page_addresses, - 4, - &[DataType::Utf8], - path.to_str().unwrap(), - 1.5, - 2, - 1, - &mut checksums, - &mut lengths, - "lz4", - 1, - false, - ) - } - .unwrap(); + let mut expected_values = inserted.to_vec(); + expected_values.sort_by_key(|&(_, partition)| partition); + let expected_values: Vec<&str> = expected_values.iter().map(|&(v, _)| v).collect(); + let dir = tempfile::tempdir().unwrap(); - // Writing each partition's rows on its own, with an encoder of its own. - let mut expected = Vec::new(); - for partition in 0..num_partitions { - let (addresses, sizes): (Vec, Vec) = rows - .iter() - .zip(&partitions) - .filter(|(_, &p)| p == partition) - .map(|(row, _)| (row.as_ptr() as i64, row.len() as i32)) - .unzip(); - if addresses.is_empty() { - assert_eq!(lengths[partition], 0, "partition {partition}"); - continue; - } - let partition_path = dir.path().join(format!("partition-{partition}")); - let [written, checksum, _] = unsafe { - write_sorted_file( - &addresses, - &sizes, + // zstd stores the offsets of a plain string column as lengths, and of a dictionary one + // as offsets, so a block writer kept for the wrong schema writes different bytes. + for (codec, offsets_as_lengths) in [("lz4", false), ("zstd", true)] { + let path = dir.path().join(codec); + let mut checksums = vec![i64::MIN; num_partitions]; + let mut lengths = vec![0_i64; num_partitions]; + unsafe { + write_sorted_partitions( + &records.pointers, + &records.page_addresses, + 4, &[DataType::Utf8], - partition_path.to_str().unwrap().to_string(), + path.to_str().unwrap(), 1.5, 2, - true, 1, - i64::MIN, - "lz4", + &mut checksums, + &mut lengths, + codec, 1, - false, + offsets_as_lengths, ) } .unwrap(); - assert_eq!(lengths[partition], written, "partition {partition}"); - assert_eq!(checksums[partition], checksum, "partition {partition}"); - expected.extend(std::fs::read(&partition_path).unwrap()); - } - let data = std::fs::read(&path).unwrap(); - assert_eq!(data, expected); - - let mut dictionary_blocks = Vec::new(); - let mut decoded = Vec::new(); - for block in shuffle_blocks(&data) { - let batch = read_ipc_compressed(block).unwrap(); - let column = batch.column(0); - dictionary_blocks.push(matches!(column.data_type(), DataType::Dictionary(_, _))); - let strings = arrow::compute::cast(column, &DataType::Utf8).unwrap(); - let strings = arrow::array::AsArray::as_string::(&strings); - decoded.extend(strings.iter().map(|s| s.unwrap().to_string())); + + // Writing each partition's rows on its own, with an encoder of its own. + let mut expected = Vec::new(); + for partition in 0..num_partitions { + let (addresses, sizes): (Vec, Vec) = rows + .iter() + .zip(&partitions) + .filter(|(_, &p)| p == partition) + .map(|(row, _)| (row.as_ptr() as i64, row.len() as i32)) + .unzip(); + if addresses.is_empty() { + assert_eq!(lengths[partition], 0, "{codec}: {partition}"); + continue; + } + let partition_path = dir.path().join(format!("{codec}-{partition}")); + let [written, checksum, _] = unsafe { + write_sorted_file( + &addresses, + &sizes, + &[DataType::Utf8], + partition_path.to_str().unwrap().to_string(), + 1.5, + 2, + true, + 1, + i64::MIN, + codec, + 1, + offsets_as_lengths, + ) + } + .unwrap(); + assert_eq!(lengths[partition], written, "{codec}: {partition}"); + assert_eq!(checksums[partition], checksum, "{codec}: {partition}"); + expected.extend(std::fs::read(&partition_path).unwrap()); + } + let data = std::fs::read(&path).unwrap(); + assert_eq!(data, expected, "{codec}"); + + let mut dictionaries = Vec::new(); + let mut decoded = Vec::new(); + for block in shuffle_blocks(&data) { + let batch = read_ipc_compressed(block).unwrap(); + let column = batch.column(0); + dictionaries.push(matches!(column.data_type(), DataType::Dictionary(_, _))); + let strings = arrow::compute::cast(column, &DataType::Utf8).unwrap(); + let strings = arrow::array::AsArray::as_string::(&strings); + decoded.extend(strings.iter().map(|s| s.unwrap().to_string())); + } + assert_eq!(dictionaries, [true, false, true, true, false], "{codec}"); + assert_eq!(decoded, expected_values, "{codec}"); } - assert_eq!(dictionary_blocks, [true, false, true, true, false]); - let mut expected_values = inserted.to_vec(); - expected_values.sort_by_key(|&(_, partition)| partition); - let expected_values: Vec<&str> = expected_values.iter().map(|&(v, _)| v).collect(); - assert_eq!(decoded, expected_values); } #[test] From c4ae0a86c1fd1285f127bc62c289412749c28c9b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 16:37:05 -0600 Subject: [PATCH 4/5] refactor: simplify the sorted row encoder and the spill metrics 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. --- native/core/src/execution/jni_api.rs | 6 +- native/shuffle/Cargo.toml | 2 +- native/shuffle/benches/row_columnar.rs | 22 +-- native/shuffle/src/spark_unsafe/row.rs | 165 +++++++----------- .../spark/shuffle/sort/SpillSorter.java | 57 ++---- .../spark/shuffle/sort/SpillSorterSuite.scala | 64 +++---- 6 files changed, 115 insertions(+), 201 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 50a6157a2ab..31ea7d84078 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -3759,11 +3759,7 @@ mod tests { // The blocks hold every row, partition by partition. let mut decoded = Vec::new(); for block in shuffle_blocks(&data) { - let batch = read_ipc_compressed(block).unwrap(); - let column = arrow::array::AsArray::as_primitive::( - batch.column(0), - ); - decoded.extend(column.values().iter().copied()); + decoded.extend(long_values(&read_ipc_compressed(block).unwrap())); } let mut expected_values = values.clone(); expected_values.sort_by_key(|&v| partition_of(v as usize)); diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml index 8106ef1fec2..03ad0a91ff2 100644 --- a/native/shuffle/Cargo.toml +++ b/native/shuffle/Cargo.toml @@ -33,7 +33,6 @@ arrow = { workspace = true } arrow-data = { workspace = true } arrow-select = { workspace = true } async-trait = { workspace = true } -bytes = { workspace = true } clap = { version = "4", features = ["derive"], optional = true } crc32c = "0.6.8" crc32fast = "1.3.2" @@ -54,6 +53,7 @@ tokio = { version = "1", features = ["rt-multi-thread"] } zstd = { version = "0.13.3", features = ["experimental"] } [dev-dependencies] +bytes = { workspace = true } criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] } datafusion = { workspace = true, features = ["parquet_encryption", "sql"] } itertools = "0.15.0" diff --git a/native/shuffle/benches/row_columnar.rs b/native/shuffle/benches/row_columnar.rs index fef3b4e8fa5..1a2ebba9fcc 100644 --- a/native/shuffle/benches/row_columnar.rs +++ b/native/shuffle/benches/row_columnar.rs @@ -463,17 +463,17 @@ fn build_sorted_spill(num_partitions: usize) -> SortedSpill { } let page_addr = page.as_ptr() as i64; - let mut partitions: Vec<(Vec, Vec)> = Vec::new(); - let mut current = usize::MAX; - for &(partition, record_offset, row_size) in &records { - if partition != current { - partitions.push((Vec::new(), Vec::new())); - current = partition; - } - let (addrs, sizes) = partitions.last_mut().unwrap(); - addrs.push(page_addr + (record_offset + UAO_SIZE + SERIALIZED_ROW_SIZE_BYTES) as i64); - sizes.push(row_size as i32); - } + let partitions = records + .chunk_by(|a, b| a.0 == b.0) + .map(|run| { + run.iter() + .map(|&(_, record_offset, row_size)| { + let row_offset = record_offset + UAO_SIZE + SERIALIZED_ROW_SIZE_BYTES; + (page_addr + row_offset as i64, row_size as i32) + }) + .unzip() + }) + .collect(); // Every record is in page 0, so a pointer is the partition id and the offset in the page. let record_pointers = records .iter() diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index 8bdbb66d87a..4c9726df5c8 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -38,12 +38,11 @@ use arrow::array::{ }; use arrow::compute::cast; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; -use arrow::error::ArrowError; use datafusion::physical_plan::metrics::Time; use datafusion_comet_jni_bridge::errors::CometError; use jni::sys::{jint, jlong}; use std::{ - fs::OpenOptions, + fs::{File, OpenOptions}, io::{Cursor, Write}, sync::Arc, }; @@ -1371,10 +1370,10 @@ const SERIALIZED_ROW_SIZE_BYTES: i64 = 4; /// partitions is not written one partition at a time. const OUTPUT_BUFFER_BYTES: usize = 1 << 20; -/// Converts batches of Spark `UnsafeRow`s to Arrow and encodes each batch as a shuffle block. -/// It holds the state that does not depend on the rows, so one encoder serves every partition of -/// a spill: the array builders, the block writer with its pre-encoded schema, and the -/// compression context. +/// Converts batches of Spark `UnsafeRow`s to Arrow and appends each batch to a file as a shuffle +/// block. It holds the state that does not depend on the rows, so one encoder serves every +/// partition of a spill: the array builders, the block writer with its pre-encoded schema, the +/// compression context and the open file. struct SortedRowEncoder<'a> { schema: &'a [DataType], prefer_dictionary_ratio: f64, @@ -1382,46 +1381,53 @@ struct SortedRowEncoder<'a> { offsets_as_lengths: bool, /// `builder_to_array` finishes every builder, which resets it for the next batch. data_builders: Vec>, - /// The schema of the last batch, reused while the columns keep their types. - batch_schema: Option, - /// The block writer pre-encodes the batch schema, so it is kept until the schema changes. + /// The schema of the last batch and the block writer that pre-encodes it, kept while the + /// columns keep their types. block_writer: Option<(SchemaRef, ShuffleBlockWriter)>, codec_context: ShuffleCodecContext, - /// Encoded blocks not yet written to the output. + output: File, + /// Encoded blocks not yet written to `output`. buffer: Vec, /// Encode and compression time across all batches. ipc_time: Time, } impl<'a> SortedRowEncoder<'a> { + /// Creates an encoder that appends to the file at `output_path`, with builders sized for + /// `capacity` rows. fn try_new( schema: &'a [DataType], - batch_size: usize, + capacity: usize, prefer_dictionary_ratio: f64, codec: &'a CompressionCodec, offsets_as_lengths: bool, + output_path: &str, ) -> Result { let data_builders = schema .iter() - .map(|dt| make_builders(dt, batch_size, prefer_dictionary_ratio)) + .map(|dt| make_builders(dt, capacity, prefer_dictionary_ratio)) .collect::, _>>()?; + let output = OpenOptions::new() + .create(true) + .append(true) + .open(output_path)?; Ok(Self { schema, prefer_dictionary_ratio, codec, offsets_as_lengths, data_builders, - batch_schema: None, block_writer: None, codec_context: ShuffleCodecContext::default(), + output, buffer: Vec::new(), ipc_time: Time::default(), }) } /// Converts rows `row_start..row_end` of the rows at `row_addresses_ptr` and `row_sizes_ptr` - /// to one batch and buffers its shuffle block, adding the block to `checksum`. Returns the - /// block's size. + /// to one batch and buffers its shuffle block, adding the block to `checksum`. Writes the + /// buffered blocks once they reach `OUTPUT_BUFFER_BYTES`. Returns the block's size. /// /// # Safety /// @@ -1455,26 +1461,36 @@ impl<'a> SortedRowEncoder<'a> { builder_to_array(builder, datatype, self.prefer_dictionary_ratio) }) .collect::, _>>()?; - let batch = self.make_batch(arrays, row_end - row_start)?; // The schema changes between batches when a dictionary-encoded column falls back to a // plain array (see `builder_to_array`). - if self - .block_writer - .as_ref() - .is_none_or(|(cached, _)| !Arc::ptr_eq(cached, batch.schema_ref())) - { + let reusable = self.block_writer.as_ref().is_some_and(|(schema, _)| { + schema + .fields() + .iter() + .zip(&arrays) + .all(|(field, array)| field.data_type() == array.data_type()) + }); + if !reusable { + let fields = arrays + .iter() + .enumerate() + .map(|(i, array)| Field::new(format!("c{i}"), array.data_type().clone(), true)) + .collect::>(); + let schema = Arc::new(Schema::new(fields)); let writer = ShuffleBlockWriter::try_new_with_offsets_as_lengths( - batch.schema_ref(), + &schema, self.codec.clone(), self.offsets_as_lengths, )?; - self.block_writer = Some((batch.schema(), writer)); + self.block_writer = Some((schema, writer)); } - let (_, writer) = self + let (schema, writer) = self .block_writer .as_ref() .expect("block writer is set above"); + let options = RecordBatchOptions::new().with_row_count(Some(row_end - row_start)); + let batch = RecordBatch::try_new_with_options(Arc::clone(schema), arrays, &options)?; let start = self.buffer.len(); let mut cursor = Cursor::new(&mut self.buffer); @@ -1484,57 +1500,24 @@ impl<'a> SortedRowEncoder<'a> { if let Some(checksum) = checksum { checksum.update(&self.buffer[start..]); } - Ok(written) - } - - /// Builds a batch of `arrays`, reusing the last batch's schema while the columns keep their - /// types, so that a new schema means the column types changed. - fn make_batch( - &mut self, - arrays: Vec, - num_rows: usize, - ) -> Result { - let reusable = self.batch_schema.as_ref().is_some_and(|schema| { - schema - .fields() - .iter() - .zip(&arrays) - .all(|(field, array)| field.data_type() == array.data_type()) - }); - if !reusable { - let fields = arrays - .iter() - .enumerate() - .map(|(i, array)| Field::new(format!("c{i}"), array.data_type().clone(), true)) - .collect::>(); - self.batch_schema = Some(Arc::new(Schema::new(fields))); + if self.buffer.len() >= OUTPUT_BUFFER_BYTES { + self.write_buffered()?; } - let schema = Arc::clone( - self.batch_schema - .as_ref() - .expect("batch schema is set above"), - ); - let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); - RecordBatch::try_new_with_options(schema, arrays, &options) + Ok(written) } - /// Writes the buffered blocks to `output` once they reach `OUTPUT_BUFFER_BYTES`. - fn write_if_full(&mut self, output: &mut impl Write) -> Result<(), CometError> { - if self.buffer.len() >= OUTPUT_BUFFER_BYTES { - self.write_buffered(output)?; - } - Ok(()) + /// Writes the blocks still buffered and returns the encode and compression time across all + /// batches, in nanoseconds. + fn finish(mut self) -> Result { + self.write_buffered()?; + Ok(self.ipc_time.value() as i64) } - fn write_buffered(&mut self, output: &mut impl Write) -> Result<(), CometError> { - output.write_all(&self.buffer)?; + fn write_buffered(&mut self) -> Result<(), CometError> { + self.output.write_all(&self.buffer)?; self.buffer.clear(); Ok(()) } - - fn encode_nanos(&self) -> i64 { - self.ipc_time.value() as i64 - } } /// Processes a sorted row partition and writes the result to the given output path. @@ -1570,18 +1553,13 @@ pub fn process_sorted_row_partition( let mut encoder = SortedRowEncoder::try_new( schema, - batch_size, + batch_size.min(row_num), prefer_dictionary_ratio, codec, offsets_as_lengths, + &output_path, )?; - // Open the output file once and reuse it across batches - let mut output_data = OpenOptions::new() - .create(true) - .append(true) - .open(&output_path)?; - while current_row < row_num { let n = std::cmp::min(batch_size, row_num - current_row); written += encoder.encode_batch( @@ -1591,15 +1569,14 @@ pub fn process_sorted_row_partition( current_row + n, current_checksum.as_mut(), )?; - encoder.write_if_full(&mut output_data)?; current_row += n; } - encoder.write_buffered(&mut output_data)?; + let encode_nanos = encoder.finish()?; Ok(( written as i64, current_checksum.map(|c| c.finalize()), - encoder.encode_nanos(), + encode_nanos, )) } @@ -1652,42 +1629,36 @@ pub unsafe fn process_sorted_partitions( ))); } + // No batch holds more rows than this. + let capacity = batch_size.min(record_pointers.len()); let mut encoder = SortedRowEncoder::try_new( schema, - batch_size, + capacity, prefer_dictionary_ratio, codec, offsets_as_lengths, + output_path, )?; - let mut output = OpenOptions::new() - .create(true) - .append(true) - .open(output_path)?; // The rows of one batch, resolved from their record pointers. - let capacity = batch_size.min(record_pointers.len()); let mut row_addresses: Vec = Vec::with_capacity(capacity); let mut row_sizes: Vec = Vec::with_capacity(capacity); - let mut start = 0; - let mut previous_partition = None; - while start < record_pointers.len() { - let partition = packed_partition_id(record_pointers[start]); + // Each run of pointers with one partition id holds that partition's records. + let mut next_partition = 0; + for run in record_pointers.chunk_by(|&a, &b| packed_partition_id(a) == packed_partition_id(b)) { + let partition = packed_partition_id(run[0]); if partition >= num_partitions { return Err(CometError::Internal(format!( "Shuffle record for partition {partition} of {num_partitions}" ))); } - if previous_partition.is_some_and(|previous| partition <= previous) { + if partition < next_partition { return Err(CometError::Internal( "Shuffle records are not sorted by partition".to_string(), )); } - previous_partition = Some(partition); - let end = record_pointers[start..] - .iter() - .position(|&pointer| packed_partition_id(pointer) != partition) - .map_or(record_pointers.len(), |len| start + len); + next_partition = partition + 1; let mut checksum = checksum_algo .map(|algo| { @@ -1696,7 +1667,7 @@ pub unsafe fn process_sorted_partitions( }) .transpose()?; let mut written = 0; - for pointers in record_pointers[start..end].chunks(batch_size) { + for pointers in run.chunks(batch_size) { row_addresses.clear(); row_sizes.clear(); for &pointer in pointers { @@ -1711,18 +1682,14 @@ pub unsafe fn process_sorted_partitions( pointers.len(), checksum.as_mut(), )?; - encoder.write_if_full(&mut output)?; } partition_lengths[partition] = written as i64; if let Some(checksum) = checksum { partition_checksums[partition] = checksum.finalize() as i64; } - start = end; } - encoder.write_buffered(&mut output)?; - - Ok(encoder.encode_nanos()) + encoder.finish() } fn packed_partition_id(pointer: i64) -> usize { diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java index 2965815b360..49bcff480cb 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java +++ b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java @@ -24,7 +24,6 @@ import javax.annotation.Nullable; import org.apache.spark.TaskContext; -import org.apache.spark.executor.ShuffleWriteMetrics; import org.apache.spark.shuffle.ShuffleWriteMetricsReporter; import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocatorTrait; import org.apache.spark.sql.comet.execution.shuffle.SpillInfo; @@ -255,18 +254,6 @@ public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) th long arrayAddr = this.sorterArray.getBaseOffset(); nativeLib.sortRowPartitionsNative(arrayAddr, numRecords, tracingEnabled); - final ShuffleWriteMetricsReporter writeMetricsToUse; - - if (isLastFile) { - // We're writing the final non-spill file, so we _do_ want to count this as shuffle bytes. - writeMetricsToUse = writeMetrics; - } else { - // We're spilling, so bytes written should be counted towards spill rather than write. - // Create a dummy WriteMetrics object to absorb these metrics, since we don't want to count - // them towards shuffle bytes written. - writeMetricsToUse = new ShuffleWriteMetrics(); - } - // Native code resolves the sorted record pointers through the page table, so the rows are // read where they lie, and continues each partition's checksum over the bytes it writes. long start = System.nanoTime(); @@ -295,44 +282,22 @@ public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) th } // Other threads may be updating the metrics at the same time, so we need to synchronize it. - synchronized (writeMetricsToUse) { - writeMetricsToUse.incWriteTime(System.nanoTime() - start); - writeMetricsToUse.incRecordsWritten(numRecords); - writeMetricsToUse.incBytesWritten(written); + // As in Spark's ShuffleExternalSorter, a spill's records count as written, but its bytes count + // as spilled rather than as shuffle bytes, and its time does not count as shuffle write time + // (SPARK-3577 tracks spill time separately). + synchronized (writeMetrics) { + writeMetrics.incRecordsWritten(numRecords); + if (isLastFile) { + writeMetrics.incWriteTime(System.nanoTime() - start); + writeMetrics.incBytesWritten(written); + } else { + taskContext.taskMetrics().incDiskBytesSpilled(written); + } } synchronized (spills) { spills.add(spillInfo); } - - if (!isLastFile) { // i.e. this is a spill file - // The current semantics of `shuffleRecordsWritten` seem to be that it's updated when - // records - // are written to disk, not when they enter the shuffle sorting code. DiskBlockObjectWriter - // relies on its `recordWritten()` method being called in order to trigger periodic updates - // to - // `shuffleBytesWritten`. If we were to remove the `recordWritten()` call and increment that - // counter at a higher-level, then the in-progress metrics for records written and bytes - // written would get out of sync. - // - // When writing the last file, we pass `writeMetrics` directly to the DiskBlockObjectWriter; - // in all other cases, we pass in a dummy write metrics to capture metrics, then copy those - // metrics to the true write metrics here. The reason for performing this copying is so that - // we can avoid reporting spilled bytes as shuffle write bytes. - // - // Note that we intentionally ignore the value of `writeMetricsToUse.shuffleWriteTime()`. - // Consistent with ExternalSorter, we do not count this IO towards shuffle write time. - // SPARK-3577 tracks the spill time separately. - - // This is guaranteed to be a ShuffleWriteMetrics based on the if check in the beginning - // of this method. - synchronized (writeMetrics) { - writeMetrics.incRecordsWritten(((ShuffleWriteMetrics) writeMetricsToUse).recordsWritten()); - taskContext - .taskMetrics() - .incDiskBytesSpilled(((ShuffleWriteMetrics) writeMetricsToUse).bytesWritten()); - } - } } /** diff --git a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala index 7adacf7cef3..f23efa88f37 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala @@ -108,23 +108,23 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { private def createSpillSorter( spillCallback: SpillSorter.SpillCallback = () => {}, - spills: java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo] = - new java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo](), - partitionChecksums: Array[Long] = new Array[Long](10)): SpillSorter = { - val schema = createTestSchema() - val writeMetrics = new ShuffleWriteMetrics() - val taskContext = TaskContext.empty() - + // Renamed on import: Spark 3.x has a `SpillInfo` of its own in this package. + spills: java.util.LinkedList[CometSpillInfo] = new java.util.LinkedList[CometSpillInfo](), + partitionChecksums: Array[Long] = new Array[Long](10), + allocator: CometShuffleMemoryAllocatorTrait = this.allocator, + checksumAlgorithm: String = "adler32", + writeMetrics: ShuffleWriteMetrics = new ShuffleWriteMetrics(), + taskContext: TaskContext = TaskContext.empty()): SpillSorter = { new SpillSorter( allocator, INITIAL_SIZE, - schema, + createTestSchema(), UAO_SIZE, 0.5, // preferDictionaryRatio "zstd", // compressionCodec 1, // compressionLevel true, // offsetsAsLengths - "adler32", // checksumAlgorithm + checksumAlgorithm, partitionChecksums, writeMetrics, taskContext, @@ -316,25 +316,16 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { val partitionChecksums = if (checksumAlgorithm == null) Array.emptyLongArray else Array.fill(numPartitions)(Long.MinValue) - // Renamed on import: Spark 3.x has a `SpillInfo` of its own in this package. val spills = new java.util.LinkedList[CometSpillInfo]() val writeMetrics = new ShuffleWriteMetrics() val taskContext = TaskContext.empty() - val sorter = new SpillSorter( - allocator, - INITIAL_SIZE, - createTestSchema(), - UAO_SIZE, - 10.0, // preferDictionaryRatio - "lz4", // compressionCodec - 1, // compressionLevel - true, // offsetsAsLengths - checksumAlgorithm, - partitionChecksums, - writeMetrics, - taskContext, - spills, - () => {}) + val sorter = createSpillSorter( + spills = spills, + partitionChecksums = partitionChecksums, + allocator = allocator, + checksumAlgorithm = checksumAlgorithm, + writeMetrics = writeMetrics, + taskContext = taskContext) // A spill, then the final file, with the same records. val files = mutable.ArrayBuffer[File]() @@ -361,22 +352,21 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { assert(spills.size === 2) val written = Seq(spills.get(0), spills.get(1)) - written.foreach { spill => + // The bytes of each partition in each file. + val partitionBytes = written.map { spill => assert(spill.partitionLengths.sum === spill.file.length()) - (0 until numPartitions).foreach { p => - assert((spill.partitionLengths(p) > 0) === usedPartitions(p), s"partition $p") - } + val bytes = Files.readAllBytes(spill.file.toPath) + val starts = spill.partitionLengths.scanLeft(0L)(_ + _).map(_.toInt) + (0 until numPartitions).map(p => bytes.slice(starts(p), starts(p + 1))) } // Native code found each record through the page table: in both files, every partition // decodes to its rows, in the order they were inserted. - written.foreach { spill => - val bytes = Files.readAllBytes(spill.file.toPath) + partitionBytes.foreach { partitions => (0 until numPartitions).foreach { p => - val start = spill.partitionLengths.take(p).sum.toInt - val end = start + spill.partitionLengths(p).toInt val ids = records.indices.filter(i => records(i)._2 == p) - assert(decodeIds(bytes.slice(start, end)) === ids, s"partition $p") + assert(partitions(p).nonEmpty === ids.nonEmpty, s"partition $p") + assert(decodeIds(partitions(p)) === ids, s"partition $p") } } @@ -386,11 +376,7 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { (0 until numPartitions).foreach { p => if (usedPartitions(p)) { val adler32 = new Adler32() - written.foreach { spill => - val bytes = Files.readAllBytes(spill.file.toPath) - val start = spill.partitionLengths.take(p).sum.toInt - adler32.update(bytes, start, spill.partitionLengths(p).toInt) - } + partitionBytes.foreach(partitions => adler32.update(partitions(p))) assert(partitionChecksums(p) === adler32.getValue, s"partition $p") } else { assert(partitionChecksums(p) === Long.MinValue, s"partition $p") From fc7595251e0e5909ebe322ac96ed36caf169ea89 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 18 Oct 2026 18:42:34 +0100 Subject: [PATCH 5/5] refactor: rename SpillSorter.writeSortedFileNative and document the output buffer Rename the spill-wide write to writeSortedPartitions so that writeSortedFileNative only names the per-partition native call, document that the encoder's output buffer is not reserved from a memory pool, and explain in the benchmark's scaladoc why it is separate from CometShuffleBenchmark. --- docs/source/contributor-guide/jvm_shuffle.md | 2 +- native/shuffle/src/spark_unsafe/row.rs | 4 +++- .../spark/shuffle/sort/CometShuffleExternalSorter.java | 4 ++-- .../java/org/apache/spark/shuffle/sort/SpillSorter.java | 2 +- .../org/apache/spark/shuffle/sort/SpillSorterSuite.scala | 2 +- .../sql/benchmark/CometColumnarShuffleWriteBenchmark.scala | 7 +++++++ 6 files changed, 15 insertions(+), 6 deletions(-) diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index 94125e7ea88..badbaba19a9 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -177,7 +177,7 @@ Selection logic in `CometShuffleManager.shouldBypassMergeSort()`: 3. Rows are serialized and buffered in off-heap memory pages 4. When memory threshold or batch size is reached, the buffered rows are written. `CometDiskBlockWriter` calls `SpillWriter.doSpilling()` for its partition. `CometShuffleExternalSorter` calls - `SpillSorter.writeSortedFileNative()` once for each spill and once more for the final file, and each + `SpillSorter.writeSortedPartitions()` once for each spill and once more for the final file, and each call writes every partition 5. Native code converts the rows to Arrow arrays and writes them as Arrow IPC blocks. `doSpilling()` passes one partition's row addresses and sizes to `Native.writeSortedFileNative()`. `SpillSorter` diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index 4c9726df5c8..694acbe6557 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -1386,7 +1386,9 @@ struct SortedRowEncoder<'a> { block_writer: Option<(SchemaRef, ShuffleBlockWriter)>, codec_context: ShuffleCodecContext, output: File, - /// Encoded blocks not yet written to `output`. + /// Encoded blocks not yet written to `output`. Native memory that no memory pool reserves, like + /// the builders. Its contents are written out as soon as they reach `OUTPUT_BUFFER_BYTES`, so + /// they never exceed that size plus one block. The `Vec` keeps the capacity it grew to. buffer: Vec, /// Encode and compression time across all batches. ipc_time: Time, diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java index 2d7b20859cd..288045f215a 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java +++ b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java @@ -199,7 +199,7 @@ public void spill() throws IOException { activeSpillSorter.setSpillInfo(spillInfo); - activeSpillSorter.writeSortedFileNative(false, tracingEnabled); + activeSpillSorter.writeSortedPartitions(false, tracingEnabled); updatePeakMemoryUsed(); final long spillSize = activeSpillSorter.freeMemory(); activeSpillSorter.reset(); @@ -338,7 +338,7 @@ public SpillInfo[] closeAndGetSpills() throws IOException { final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId); activeSpillSorter.setSpillInfo(spillInfo); - activeSpillSorter.writeSortedFileNative(true, tracingEnabled); + activeSpillSorter.writeSortedPartitions(true, tracingEnabled); freeMemory(); } diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java index 49bcff480cb..d68ef454369 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java +++ b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java @@ -243,7 +243,7 @@ public int numRecords() { * #setSpillInfo}, every partition in one native call, recording each partition's length and * checksum. */ - public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) throws IOException { + public void writeSortedPartitions(boolean isLastFile, boolean tracingEnabled) throws IOException { final int numRecords = inMemSorter.numRecords(); // If there are no records, we don't need to create an empty spill file. if (numRecords == 0) { diff --git a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala index f23efa88f37..bf691530ccd 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala @@ -343,7 +343,7 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { files += file sorter.setSpillInfo( new CometSpillInfo(numPartitions, file, TempShuffleBlockId(UUID.randomUUID()))) - sorter.writeSortedFileNative(isLastFile, false) + sorter.writeSortedPartitions(isLastFile, false) sorter.freeMemory() if (!isLastFile) { sorter.reset() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala index 53138220045..0e159984968 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala @@ -41,6 +41,13 @@ import org.apache.comet.CometConf * `spark.comet.shuffle.jvm.spillThreshold` makes Comet spill a fixed number of rows at a time. It * does not apply to Spark's shuffle, which does not spill in these cases. * + * This is separate from `CometShuffleBenchmark`, which also covers the sort-based writer, because + * it needs a session of its own. The master is fixed when the session starts, and this benchmark + * uses `local[1]` so that the map tasks run one at a time and the time follows their cost, where + * `CometShuffleBenchmark` uses `local[5]`. That benchmark also sets a default spill threshold, so + * its no-spill cases would spill, and this one leaves the default so that its no-spill case does + * not. + * * To run this benchmark: * {{{ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometColumnarShuffleWriteBenchmark