You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/19 14:07:19 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-7065-80b518173d0631d31028b732b385c68f69ace326 created (now 9d7344ec26)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-7065-80b518173d0631d31028b732b385c68f69ace326
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 9d7344ec26 perf: write every partition of a JVM shuffle spill in one native call (#7065)
This branch includes the following new commits:
new 9d7344ec26 perf: write every partition of a JVM shuffle spill in one native call (#7065)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: perf: write every partition of a JVM shuffle spill in one native call (#7065)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-7065-80b518173d0631d31028b732b385c68f69ace326
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 9d7344ec26fd81e8019c756ba58da19e4019649f
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon Oct 19 14:02:51 2026 +0000
perf: write every partition of a JVM shuffle spill in one native call (#7065)
* 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).
* 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.
* 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.
* 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.
* 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 | 26 +-
native/core/src/execution/jni_api.rs | 538 ++++++++++++++++++++-
native/shuffle/Cargo.toml | 2 +-
native/shuffle/benches/row_columnar.rs | 188 ++++++-
native/shuffle/src/spark_unsafe/row.rs | 460 ++++++++++++++----
native/shuffle/src/writers/checksum.rs | 38 +-
.../CometUnboundedShuffleMemoryAllocator.java | 7 +-
.../shuffle/sort/CometShuffleExternalSorter.java | 4 +-
.../org/apache/spark/shuffle/sort/SpillSorter.java | 193 +++-----
.../sql/comet/execution/shuffle/SpillWriter.java | 26 +-
spark/src/main/scala/org/apache/comet/Native.scala | 58 ++-
.../spark/shuffle/sort/SpillSorterSuite.scala | 168 ++++++-
.../CometColumnarShuffleWriteBenchmark.scala | 126 +++++
13 files changed, 1552 insertions(+), 282 deletions(-)
diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md
index f41877b08d..badbaba19a 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 cabbb58d2b..31ea7d8407 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::<i64>(),
+ 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<i64> {
+ 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<Vec<u64>>,
+ page_addresses: Vec<i64>,
+ pointers: Vec<i64>,
+ }
+
+ 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<R: AsRef<[u8]>>(rows: &[R], partitions: &[usize]) -> Self {
+ let mut page_bytes: Vec<Vec<u8>> = 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<Vec<u64>> = 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<i64> {
+ 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<i64> = (0..23).collect();
+ let rows = long_rows(&values);
+ let partitions: Vec<usize> = (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<Vec<u8>> {
+ 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<usize>) = 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<i64>, Vec<i32>) = 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::<i32>(&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::<Vec<_>>());
+ let partitions: Vec<usize> = (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::<i64>() 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 8106ef1fec..03ad0a91ff 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 e6a04b822c..1a2ebba9fc 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<ArrowDataType> {
+ 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<u64>,
+ /// The packed record pointers, sorted by partition, as the sorter holds them.
+ record_pointers: Vec<i64>,
+ /// 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<i64>, Vec<i32>)>,
+}
+
+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<u8> = 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 219f09131e..a7b3a2fab7 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,6 +1370,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<Box<dyn ArrayBuilder>>,
+ /// 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<u8>,
+ /// 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<Self, CometError> {
+ let data_builders = schema
+ .iter()
+ .map(|dt| make_builders(dt, capacity, prefer_dictionary_ratio))
+ .collect::<Result<Vec<_>, _>>()?;
+ 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<usize, CometError> {
+ 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::<Result<Vec<_>, _>>()?;
+
+ // 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::<Vec<_>>();
+ 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<i64, CometError> {
+ 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(
@@ -1402,91 +1567,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<Box<dyn ArrayBuilder>> = 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<u8> = 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<Vec<ArrayRef>, _> = 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<i32>,
+ partition_checksums: &mut [i64],
+ partition_lengths: &mut [i64],
+ codec: &CompressionCodec,
+ offsets_as_lengths: bool,
+) -> Result<i64, CometError> {
+ 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<jlong> = Vec::with_capacity(capacity);
+ let mut row_sizes: Vec<jint> = 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,
))
}
@@ -1540,17 +1800,6 @@ fn builder_to_array(
}
}
-fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> Result<RecordBatch, ArrowError> {
- let fields = arrays
- .iter()
- .enumerate()
- .map(|(i, array)| Field::new(format!("c{i}"), array.data_type().clone(), true))
- .collect::<Vec<_>>();
- 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;
@@ -1701,6 +1950,47 @@ mod test {
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}");
+ }
+ }
+ }
+
/// Lays out an 8-byte-aligned `UnsafeRow` with one struct field: the null bitset, the field's
/// offset and size, then the nested `UnsafeRow` holding `fields`. A one-word field is stored
/// in its slot, a longer one in the variable-length region, and `None` is written as null, as
diff --git a/native/shuffle/src/writers/checksum.rs b/native/shuffle/src/writers/checksum.rs
index 1dfd15eb37..dacb6def15 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<u8>>) -> 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<u8> = 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<u8> = 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<u8> = 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 4ce62af6e8..e437acd303 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 @@ import org.apache.spark.unsafe.memory.UnsafeMemoryAllocator;
* serialized rows when Spark runs in on-heap mode.
*
* <p>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.
*
* <p>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 2d7b20859c..288045f215 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 final class CometShuffleExternalSorter implements CometShuffleChecksumSup
activeSpillSorter.setSpillInfo(spillInfo);
- activeSpillSorter.writeSortedFileNative(false, tracingEnabled);
+ activeSpillSorter.writeSortedPartitions(false, tracingEnabled);
updatePeakMemoryUsed();
final long spillSize = activeSpillSorter.freeMemory();
activeSpillSorter.reset();
@@ -338,7 +338,7 @@ public final class CometShuffleExternalSorter implements CometShuffleChecksumSup
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 23368115ae..d68ef45436 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 java.util.LinkedList;
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.sql.types.StructType;
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 class SpillSorter extends SpillWriter {
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 class SpillSorter extends SpillWriter {
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 class SpillSorter extends SpillWriter {
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 dc67838c81..48cc2d41f9 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 abstract class SpillWriter {
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 @@ public abstract class SpillWriter {
}
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 @@ public abstract class SpillWriter {
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 c12a85dee9..efeabf079a 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 d2d6794126..bf691530cc 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 0000000000..0e15998496
--- /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)
+ }
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org