diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index f41877b08dc..badbaba19a9 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. `CometShuffleExternalSorter` calls + `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` + 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 cabbb58d2bc..31ea7d84078 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,391 @@ 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]; + + /// 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 { + /// 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()); + 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) { + 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)); + assert_eq!(decoded, expected_values, "{name}"); + } + } + + /// 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 num_partitions = 4; + 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(); + + // 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], + path.to_str().unwrap(), + 1.5, + 2, + 1, + &mut checksums, + &mut lengths, + codec, + 1, + offsets_as_lengths, + ) + } + .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, "{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}"); + } + } + + #[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::<[u8; 16]>(&[], &[]); + 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/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 e6a04b822cf..1a2ebba9fcc 100644 --- a/native/shuffle/benches/row_columnar.rs +++ b/native/shuffle/benches/row_columnar.rs @@ -19,11 +19,14 @@ //! //! 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}; -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 +381,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 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() + .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 +569,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..694acbe6557 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, }; @@ -1357,6 +1356,172 @@ 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 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, + 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 and the block writer that pre-encodes it, kept while the + /// columns keep their types. + block_writer: Option<(SchemaRef, ShuffleBlockWriter)>, + codec_context: ShuffleCodecContext, + output: File, + /// 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, +} + +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], + 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, 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, + 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`. Writes the + /// buffered blocks once they reach `OUTPUT_BUFFER_BYTES`. 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::, _>>()?; + + // The schema changes between batches when a dictionary-encoded column falls back to a + // plain array (see `builder_to_array`). + 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( + &schema, + self.codec.clone(), + self.offsets_as_lengths, + )?; + self.block_writer = Some((schema, writer)); + } + 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); + 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..]); + } + if self.buffer.len() >= OUTPUT_BUFFER_BYTES { + self.write_buffered()?; + } + Ok(written) + } + + /// 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) -> Result<(), CometError> { + self.output.write_all(&self.buffer)?; + self.buffer.clear(); + Ok(()) + } +} + /// 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,91 +1553,186 @@ 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>(()) - })?; - - // Open the output file once and reuse it across batches - let mut output_data = OpenOptions::new() - .create(true) - .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; + let mut encoder = SortedRowEncoder::try_new( + schema, + batch_size.min(row_num), + prefer_dictionary_ratio, + codec, + offsets_as_lengths, + &output_path, + )?; 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(), + )?; + current_row += n; + } + let encode_nanos = encoder.finish()?; - // 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, - )?; - } - - // 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)?; + Ok(( + written as i64, + current_checksum.map(|c| c.finalize()), + encode_nanos, + )) +} - frozen.clear(); - let mut cursor = Cursor::new(&mut frozen); +/// 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() + ))); + } - // 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, + // No batch holds more rows than this. + let capacity = batch_size.min(record_pointers.len()); + let mut encoder = SortedRowEncoder::try_new( + schema, + capacity, + prefer_dictionary_ratio, + codec, + offsets_as_lengths, + output_path, + )?; + + // The rows of one batch, resolved from their record pointers. + let mut row_addresses: Vec = Vec::with_capacity(capacity); + let mut row_sizes: Vec = Vec::with_capacity(capacity); + + // 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 partition < next_partition { + return Err(CometError::Internal( + "Shuffle records are not sorted by partition".to_string(), + )); + } + next_partition = partition + 1; + + 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 run.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)); } - 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; } + encoder.finish() +} +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 +1786,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 +1935,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/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 23368115ae9..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 @@ -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; @@ -33,7 +32,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 +59,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 +115,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,135 +238,83 @@ public int numRecords() { return this.inMemSorter.numRecords(); } - 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()) { + /** + * 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 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) { return; } - 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(); + // This call performs the actual sort. + long arrayAddr = this.sorterArray.getBaseOffset(); + nativeLib.sortRowPartitionsNative(arrayAddr, numRecords, tracingEnabled); + + // 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; } - 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; + // Other threads may be updating the metrics at the same time, so we need to synchronize it. + // 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); } - - 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); } - 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(); - } - - synchronized (spills) { - spills.add(spillInfo); - } + 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()); + /** + * 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; } } 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..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 @@ -19,7 +19,14 @@ package org.apache.spark.shuffle.sort +import java.io.{ByteArrayInputStream, 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,10 +35,17 @@ 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.{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]]. * @@ -75,25 +89,42 @@ 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] = - 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, @@ -267,6 +298,125 @@ 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) + val spills = new java.util.LinkedList[CometSpillInfo]() + val writeMetrics = new ShuffleWriteMetrics() + val taskContext = TaskContext.empty() + 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]() + 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.writeSortedPartitions(isLastFile, false) + sorter.freeMemory() + if (!isLastFile) { + sorter.reset() + } + } + + assert(spills.size === 2) + val written = Seq(spills.get(0), spills.get(1)) + // The bytes of each partition in each file. + val partitionBytes = written.map { spill => + assert(spill.partitionLengths.sum === spill.file.length()) + 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. + partitionBytes.foreach { partitions => + (0 until numPartitions).foreach { p => + val ids = records.indices.filter(i => records(i)._2 == p) + assert(partitions(p).nonEmpty === ids.nonEmpty, s"partition $p") + assert(decodeIds(partitions(p)) === 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. + (0 until numPartitions).foreach { p => + if (usedPartitions(p)) { + val adler32 = new Adler32() + partitionBytes.foreach(partitions => adler32.update(partitions(p))) + 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..0e159984968 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala @@ -0,0 +1,126 @@ +/* + * 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. + * + * 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 + * }}} + * 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.toLong, 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) + } + } + } + } +}