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/05 04:23:40 UTC
(datafusion-comet) 01/01: perf: speed up string, map and array conversion in JVM columnar shuffle (#6604)
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-6604-93ce5c3471dbd89aa5cd5f6f9e26df27adab6a8a
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 8368bb90c324374fa41e3627a1585709ee883162
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon Oct 5 04:23:16 2026 +0000
perf: speed up string, map and array conversion in JVM columnar shuffle (#6604)
* perf: speed up string, map and array conversion in JVM columnar shuffle
The JVM columnar shuffle converts Spark UnsafeRows to Arrow natively, and
three per-value costs dominated its string, map and array conversion:
- decode_utf8_spark_lossy validated every string with str::from_utf8, which
costs several times more than an is_ascii check on short strings. ASCII
strings now skip it, and the check is inlined into callers in other crates.
- append_map_elements built a MapFieldNames, allocating and freeing six
strings, for every map value just to read the key and value fields.
- Nullable element arrays were always appended one element at a time. An
array of eight or more elements with no nulls is now copied in one go.
The timestamp and date appenders, which duplicated the macro, now use it.
* perf: don't scan long strings twice in decode_utf8_spark_lossy
decode_utf8_spark_lossy checked the whole value with is_ascii() and, when that
failed, validated it again from the start with str::from_utf8. So a long value
whose first non-ASCII byte came late was read twice: 4096-byte values ending in
é converted 12-16% slower than on main on x86, and 35% slower on aarch64. On
aarch64, is_ascii checks one word at a time and is slower than str::from_utf8
on long values, so long ASCII values were slower than on main too.
Values of up to 32 bytes still get the inline is_ascii check. Longer values are
checked out of line 32 bytes at a time, ORing each block's words, and
str::from_utf8 validates them only from the first block that is not ASCII. The
block check is out of line because inlining it made the shuffle's get_string
too large to inline into its per-row loops, and the out-of-line function is
cold so that those loops keep the short ASCII path as the fall-through.
---
native/common/src/utils.rs | 106 +++++++++++++-
native/shuffle/src/spark_unsafe/list.rs | 236 +++++++++++++++++---------------
native/shuffle/src/spark_unsafe/map.rs | 52 ++++---
native/shuffle/src/spark_unsafe/row.rs | 4 +-
4 files changed, 256 insertions(+), 142 deletions(-)
diff --git a/native/common/src/utils.rs b/native/common/src/utils.rs
index 1838d87d40..ab2534c5e5 100644
--- a/native/common/src/utils.rs
+++ b/native/common/src/utils.rs
@@ -52,10 +52,59 @@ pub fn bytes_to_i128(slice: &[u8]) -> i128 {
/// per-class malformed lengths below (E0/ED overlong & surrogate handling, F0/F4 range checks)
/// match the observable replacement behavior of the JDK UTF-8 decoder; they were determined from
/// observed `new String(bytes, UTF_8)` output, not by reviewing the OpenJDK source.
+// Inlined so that callers in other crates, such as the JVM columnar shuffle converting one value at
+// a time, get the ASCII check without a call.
+#[inline]
pub fn decode_utf8_spark_lossy(bytes: &[u8]) -> Cow<'_, str> {
+ // ASCII is valid UTF-8, and on short strings checking for it is several times cheaper than
+ // `str::from_utf8`, whose cost dominated the JVM columnar shuffle's string conversion. Longer
+ // values are checked out of line: checking them here made the shuffle's `get_string` too large
+ // to inline into its per-row loops.
+ if bytes.len() <= ASCII_BLOCK && bytes.is_ascii() {
+ // SAFETY: every ASCII byte sequence is valid UTF-8.
+ return Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(bytes) });
+ }
+ decode_long_or_non_ascii_utf8_spark_lossy(bytes)
+}
+
+/// How many bytes [`decode_utf8_spark_lossy`] checks for ASCII at a time.
+const ASCII_BLOCK: usize = 32;
+
+/// Whether every byte of `block` is ASCII. ORing its words keeps every byte's high bit and compiles
+/// to a few vector instructions. On aarch64, `<[u8]>::is_ascii` branches on each word instead, and
+/// on long values it is slower than `str::from_utf8`.
+#[inline]
+fn is_ascii_block(block: &[u8; ASCII_BLOCK]) -> bool {
+ let (words, _) = block.as_chunks::<8>();
+ let bits = words
+ .iter()
+ .fold(0, |acc, word| acc | u64::from_ne_bytes(*word));
+ bits & 0x8080_8080_8080_8080 == 0
+}
+
+/// [`decode_utf8_spark_lossy`] for bytes that are longer than an ASCII block or not all ASCII.
+// Cold so that callers lay the short ASCII path out as the fall-through. Without it, the shuffle's
+// per-row string loops took an extra jump for every short ASCII value.
+#[cold]
+fn decode_long_or_non_ascii_utf8_spark_lossy(bytes: &[u8]) -> Cow<'_, str> {
+ // A long value is checked for ASCII a block at a time, and only the bytes from the first block
+ // that is not ASCII are validated, so a value whose first non-ASCII byte comes late is not
+ // scanned twice. An ASCII byte is a whole character, so those bytes validate on their own.
+ let mut ascii_len = 0;
+ if bytes.len() > ASCII_BLOCK {
+ let (blocks, tail) = bytes.as_chunks::<ASCII_BLOCK>();
+ let ascii_blocks = blocks.iter().take_while(|b| is_ascii_block(b)).count();
+ if ascii_blocks == blocks.len() && tail.is_ascii() {
+ // SAFETY: every ASCII byte sequence is valid UTF-8.
+ return Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(bytes) });
+ }
+ ascii_len = ascii_blocks * ASCII_BLOCK;
+ }
+
// Fast path: well-formed UTF-8 borrows with zero copy (the overwhelmingly common case).
- if let Ok(s) = std::str::from_utf8(bytes) {
- return Cow::Borrowed(s);
+ if std::str::from_utf8(&bytes[ascii_len..]).is_ok() {
+ // SAFETY: ASCII followed by valid UTF-8 is valid UTF-8.
+ return Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(bytes) });
}
const RC: char = '\u{FFFD}';
@@ -198,6 +247,59 @@ mod tests {
}
}
+ #[test]
+ fn decode_utf8_spark_lossy_ascii_is_borrowed_zero_copy() {
+ // Lengths below and above a machine word and an ASCII block, so both the check of short
+ // values and the block check of long ones run, the latter with bytes after the last block.
+ for s in [
+ "",
+ "a",
+ "value-1",
+ "value-12345",
+ "abcdefghijklmnopqrstuvwxyz0123456789",
+ &"0123456789".repeat(10),
+ ] {
+ match decode_utf8_spark_lossy(s.as_bytes()) {
+ Cow::Borrowed(b) => assert_eq!(b, s),
+ Cow::Owned(_) => panic!("ASCII must borrow, not allocate"),
+ }
+ }
+ // A single byte past ASCII anywhere still takes the validating paths.
+ let mut bytes = b"abcdefghijklmnop".to_vec();
+ bytes[11] = 0xFF;
+ assert_eq!(decode_utf8_spark_lossy(&bytes), "abcdefghijk\u{FFFD}mnop");
+ }
+
+ #[test]
+ fn decode_utf8_spark_lossy_non_ascii_at_every_offset() {
+ // A long value is validated from its first block that is not ASCII. Put a character or an
+ // ill-formed byte at every offset of values up to a few blocks long, so it lands in short
+ // values and in each block of long ones, across block boundaries and after the last block.
+ let ascii = |n: usize| "x".repeat(n);
+ for len in 1..=100 {
+ for at in 0..len {
+ for c in ["é", "日", "🦀"] {
+ let s = format!("{}{c}{}", ascii(at), ascii(len - at));
+ match decode_utf8_spark_lossy(s.as_bytes()) {
+ Cow::Borrowed(b) => assert_eq!(b, s),
+ Cow::Owned(_) => panic!("valid UTF-8 must borrow, not allocate"),
+ }
+ // An ill-formed byte after it still makes the value take the lossy decode.
+ let mut bytes = s.as_bytes().to_vec();
+ bytes.push(0xFF);
+ assert_eq!(decode_utf8_spark_lossy(&bytes), format!("{s}\u{FFFD}"));
+ }
+ // 0xFF is never valid, and 0x80 is a continuation byte with no lead byte.
+ for b in [0xFF, 0x80] {
+ let mut bytes = ascii(len).into_bytes();
+ bytes[at] = b;
+ let expected = format!("{}\u{FFFD}{}", ascii(at), ascii(len - at - 1));
+ assert_eq!(decode_utf8_spark_lossy(&bytes), expected);
+ }
+ }
+ }
+ }
+
#[test]
fn decode_utf8_spark_lossy_valid_multibyte_around_invalid_bytes_decodes() {
// 'a' | é (C3 A9) | stray 0xFF | 'b' | 🦀 (F0 9F A6 80) -> valid chars preserved, one FFFD.
diff --git a/native/shuffle/src/spark_unsafe/list.rs b/native/shuffle/src/spark_unsafe/list.rs
index 22dce2cf79..3d8e89ebc5 100644
--- a/native/shuffle/src/spark_unsafe/list.rs
+++ b/native/shuffle/src/spark_unsafe/list.rs
@@ -32,6 +32,13 @@ use arrow::array::{
use arrow::datatypes::{DataType, TimeUnit};
use datafusion_comet_jni_bridge::errors::CometError;
+/// Element count from which a nullable array with no null elements is appended with one copy
+/// instead of element by element. Below it the loop is as fast: checking the null bitset and
+/// calling `append_slice` cost about as much as a few appends, and arrays of up to five elements
+/// with some null elements got slower with the copy. Arrays of 16 elements take about 40% of the
+/// time with the copy, and arrays of 50 about a fifth.
+const MIN_BULK_APPEND_ELEMENTS: usize = 8;
+
/// Generates bulk append methods for primitive types in SparkUnsafeArray.
///
/// # Safety invariants for all generated methods:
@@ -46,7 +53,8 @@ macro_rules! impl_append_to_builder {
return;
}
- if NULLABLE {
+ // A nullable array without nulls takes the copy below too, once it is long enough.
+ if NULLABLE && (num_elements < MIN_BULK_APPEND_ELEMENTS || self.has_null()) {
let mut ptr = self.element_offset as *const $element_type;
let null_words = self.null_bitset_ptr();
debug_assert!(!null_words.is_null(), "null_bitset_ptr is null");
@@ -167,6 +175,16 @@ impl SparkUnsafeArray {
(self.row_addr + 8) as *const i64
}
+ /// Returns true if any element is null, checking the null bitset a word at a time.
+ #[inline]
+ fn has_null(&self) -> bool {
+ let null_words = self.null_bitset_ptr();
+ // SAFETY: the null bitset holds ceil(num_elements/64) words. Spark zeroes the bits past
+ // the last element, and a stray one would only send the array down the per-element path.
+ (0..self.num_elements.div_ceil(64))
+ .any(|word| unsafe { null_words.add(word).read_unaligned() } != 0)
+ }
+
/// Checks whether the null bit at `idx` is set in the given null bitset pointer.
///
/// # Safety
@@ -185,6 +203,12 @@ impl SparkUnsafeArray {
impl_append_to_builder!(append_bytes_to_builder, Int8Builder, i8);
impl_append_to_builder!(append_floats_to_builder, Float32Builder, f32);
impl_append_to_builder!(append_doubles_to_builder, Float64Builder, f64);
+ impl_append_to_builder!(
+ append_timestamps_to_builder,
+ TimestampMicrosecondBuilder,
+ i64
+ );
+ impl_append_to_builder!(append_dates_to_builder, Date32Builder, i32);
/// Bulk append boolean values to builder.
/// Booleans are stored as 1 byte each in SparkUnsafeArray, requiring special handling.
@@ -230,114 +254,6 @@ impl SparkUnsafeArray {
}
}
}
-
- /// Bulk append timestamp values to builder (stored as i64 microseconds).
- pub(crate) fn append_timestamps_to_builder<const NULLABLE: bool>(
- &self,
- builder: &mut TimestampMicrosecondBuilder,
- ) {
- let num_elements = self.num_elements;
- if num_elements == 0 {
- return;
- }
-
- if NULLABLE {
- let mut ptr = self.element_offset as *const i64;
- let null_words = self.null_bitset_ptr();
- debug_assert!(
- !null_words.is_null(),
- "append_timestamps: null_bitset_ptr is null"
- );
- debug_assert!(
- !ptr.is_null(),
- "append_timestamps: element_offset pointer is null"
- );
- for idx in 0..num_elements {
- // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements
- let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) };
-
- if is_null {
- builder.append_null();
- } else {
- // SAFETY: ptr is within element data bounds
- builder.append_value(unsafe { ptr.read_unaligned() });
- }
- // SAFETY: ptr stays within bounds, iterating num_elements times
- ptr = unsafe { ptr.add(1) };
- }
- } else {
- // SAFETY: element_offset points to contiguous i64 data of length num_elements
- debug_assert!(
- self.element_offset != 0,
- "append_timestamps: element_offset is null"
- );
- let ptr = self.element_offset as *const i64;
- if (ptr as usize).is_multiple_of(std::mem::align_of::<i64>()) {
- let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) };
- builder.append_slice(slice);
- } else {
- let mut ptr = ptr;
- for _ in 0..num_elements {
- builder.append_value(unsafe { ptr.read_unaligned() });
- ptr = unsafe { ptr.add(1) };
- }
- }
- }
- }
-
- /// Bulk append date values to builder (stored as i32 days since epoch).
- pub(crate) fn append_dates_to_builder<const NULLABLE: bool>(
- &self,
- builder: &mut Date32Builder,
- ) {
- let num_elements = self.num_elements;
- if num_elements == 0 {
- return;
- }
-
- if NULLABLE {
- let mut ptr = self.element_offset as *const i32;
- let null_words = self.null_bitset_ptr();
- debug_assert!(
- !null_words.is_null(),
- "append_dates: null_bitset_ptr is null"
- );
- debug_assert!(
- !ptr.is_null(),
- "append_dates: element_offset pointer is null"
- );
- for idx in 0..num_elements {
- // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements
- let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) };
-
- if is_null {
- builder.append_null();
- } else {
- // SAFETY: ptr is within element data bounds
- builder.append_value(unsafe { ptr.read_unaligned() });
- }
- // SAFETY: ptr stays within bounds, iterating num_elements times
- ptr = unsafe { ptr.add(1) };
- }
- } else {
- // SAFETY: element_offset points to contiguous i32 data of length num_elements
- debug_assert!(
- self.element_offset != 0,
- "append_dates: element_offset is null"
- );
- let ptr = self.element_offset as *const i32;
- if (ptr as usize).is_multiple_of(std::mem::align_of::<i32>()) {
- let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) };
- builder.append_slice(slice);
- } else {
- let mut ptr = ptr;
- for _ in 0..num_elements {
- builder.append_value(unsafe { ptr.read_unaligned() });
- ptr = unsafe { ptr.add(1) };
- }
- }
- }
- }
}
pub fn append_to_builder<const NULLABLE: bool>(
@@ -494,3 +410,105 @@ pub fn append_list_element(
Ok(())
}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+ use arrow::array::builder::PrimitiveBuilder;
+ use arrow::datatypes::{
+ ArrowPrimitiveType, Date32Type, Int16Type, Int32Type, Int64Type, Int8Type,
+ TimestampMicrosecondType,
+ };
+
+ /// Lays `values` out as a Spark `UnsafeArrayData` of `width`-byte elements: the element
+ /// count, the null bitset, then the elements.
+ fn unsafe_array_bytes(values: &[Option<i64>], width: usize) -> Vec<u8> {
+ let bitset_words = values.len().div_ceil(64);
+ let mut bitset = vec![0u64; bitset_words];
+ let mut data = vec![0u8; (values.len() * width).div_ceil(8) * 8];
+ for (i, value) in values.iter().enumerate() {
+ match value {
+ None => bitset[i / 64] |= 1 << (i % 64),
+ Some(v) => {
+ data[i * width..(i + 1) * width].copy_from_slice(&v.to_le_bytes()[..width])
+ }
+ }
+ }
+ let mut bytes = (values.len() as u64).to_le_bytes().to_vec();
+ bitset
+ .iter()
+ .for_each(|word| bytes.extend_from_slice(&word.to_le_bytes()));
+ bytes.extend_from_slice(&data);
+ bytes
+ }
+
+ /// Appends `values`, laid out `shift` bytes past an 8-byte boundary, through
+ /// `append_to_builder::<true>` and reads them back.
+ fn round_trip<T>(data_type: DataType, values: &[Option<i64>], shift: usize) -> Vec<Option<i64>>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: Into<i64>,
+ {
+ let bytes = unsafe_array_bytes(values, std::mem::size_of::<T::Native>());
+ let mut words = vec![0u64; (shift + bytes.len()).div_ceil(8)];
+ // SAFETY: `words` covers `shift + bytes.len()` bytes.
+ let base = unsafe { (words.as_mut_ptr() as *mut u8).add(shift) };
+ unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), base, bytes.len()) };
+ let array = SparkUnsafeArray::new(base as i64);
+ let mut builder = PrimitiveBuilder::<T>::new();
+ append_to_builder::<true>(&data_type, &mut builder, &array).unwrap();
+ builder.finish().iter().map(|v| v.map(Into::into)).collect()
+ }
+
+ #[test]
+ fn nullable_primitive_arrays_round_trip_on_every_append_path() {
+ let long: Vec<Option<i64>> = (0..70).map(Some).collect();
+ let mut null_in_second_bitset_word = long.clone();
+ null_in_second_bitset_word[66] = None;
+ let cases = [
+ vec![],
+ // Shorter than MIN_BULK_APPEND_ELEMENTS: element by element.
+ vec![Some(1), Some(-2), Some(3)],
+ // Holds a null: element by element.
+ vec![Some(1), None, Some(3), Some(-4), Some(5)],
+ // No nulls and long enough: one copy.
+ vec![Some(-1), Some(2), Some(-3), Some(4)],
+ long,
+ // The null is only in the second bitset word, which the copy check must read.
+ null_in_second_bitset_word,
+ ];
+ for values in &cases {
+ // An 8-byte boundary, as inside an UnsafeRow, and a misaligned one for every width.
+ for shift in [0, 1, 4] {
+ assert_eq!(
+ round_trip::<Int8Type>(DataType::Int8, values, shift),
+ *values
+ );
+ assert_eq!(
+ round_trip::<Int16Type>(DataType::Int16, values, shift),
+ *values
+ );
+ assert_eq!(
+ round_trip::<Int32Type>(DataType::Int32, values, shift),
+ *values
+ );
+ assert_eq!(
+ round_trip::<Int64Type>(DataType::Int64, values, shift),
+ *values
+ );
+ assert_eq!(
+ round_trip::<Date32Type>(DataType::Date32, values, shift),
+ *values
+ );
+ assert_eq!(
+ round_trip::<TimestampMicrosecondType>(
+ DataType::Timestamp(TimeUnit::Microsecond, None),
+ values,
+ shift
+ ),
+ *values
+ );
+ }
+ }
+ }
+}
diff --git a/native/shuffle/src/spark_unsafe/map.rs b/native/shuffle/src/spark_unsafe/map.rs
index 026e6f71d5..dc7e8b234c 100644
--- a/native/shuffle/src/spark_unsafe/map.rs
+++ b/native/shuffle/src/spark_unsafe/map.rs
@@ -72,7 +72,7 @@ pub fn append_map_elements(
map_builder: &mut MapBuilder<Box<dyn ArrayBuilder>, Box<dyn ArrayBuilder>>,
map: &SparkUnsafeMap,
) -> Result<(), CometError> {
- let (key_field, value_field, _) = get_map_key_value_fields(field)?;
+ let (key_field, value_field) = map_key_value_fields(field)?;
let keys = &map.keys;
let values = &map.values;
@@ -86,37 +86,31 @@ pub fn append_map_elements(
Ok(())
}
-#[allow(clippy::field_reassign_with_default)]
+/// Returns the key and value fields of a map's entries field. This runs once for every map value
+/// the shuffle converts, so unlike [`get_map_key_value_fields`] it builds no `MapFieldNames`,
+/// whose strings would be allocated and freed each time.
+pub fn map_key_value_fields(field: &FieldRef) -> Result<(&FieldRef, &FieldRef), CometError> {
+ match field.data_type() {
+ DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0], &fields[1])),
+ DataType::Struct(fields) => Err(CometError::Internal(format!(
+ "Map field should have 2 fields, but got {}",
+ fields.len()
+ ))),
+ data_type => Err(CometError::Internal(format!(
+ "Map field should be a struct, but got {data_type:?}"
+ ))),
+ }
+}
+
+/// Returns the key and value fields of a map's entries field, with the names a `MapBuilder` needs.
pub fn get_map_key_value_fields(
field: &FieldRef,
) -> Result<(&FieldRef, &FieldRef, MapFieldNames), CometError> {
- let mut map_fieldnames = MapFieldNames::default();
- map_fieldnames.entry = field.name().to_string();
-
- let (key_field, value_field) = match field.data_type() {
- DataType::Struct(fields) => {
- if fields.len() != 2 {
- return Err(CometError::Internal(format!(
- "Map field should have 2 fields, but got {}",
- fields.len()
- )));
- }
-
- let key = &fields[0];
- let value = &fields[1];
-
- map_fieldnames.key = key.name().to_string();
- map_fieldnames.value = value.name().to_string();
-
- (key, value)
- }
- _ => {
- return Err(CometError::Internal(format!(
- "Map field should be a struct, but got {:?}",
- field.data_type()
- )));
- }
+ let (key_field, value_field) = map_key_value_fields(field)?;
+ let map_fieldnames = MapFieldNames {
+ entry: field.name().to_string(),
+ key: key_field.name().to_string(),
+ value: value_field.name().to_string(),
};
-
Ok((key_field, value_field, map_fieldnames))
}
diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs
index 1918ce3b18..180673c37f 100644
--- a/native/shuffle/src/spark_unsafe/row.rs
+++ b/native/shuffle/src/spark_unsafe/row.rs
@@ -21,7 +21,7 @@ use crate::codec_context::ShuffleCodecContext;
use crate::spark_unsafe::unsafe_object::{impl_primitive_accessors, SparkUnsafeObject};
use crate::spark_unsafe::{
list::append_list_element,
- map::{append_map_elements, get_map_key_value_fields},
+ map::{append_map_elements, get_map_key_value_fields, map_key_value_fields},
};
use crate::writers::Checksum;
use crate::writers::ShuffleBlockWriter;
@@ -705,7 +705,7 @@ fn append_map_column_batch(
map_builder: &mut MapBuilder<Box<dyn ArrayBuilder>, Box<dyn ArrayBuilder>>,
) -> Result<(), CometError> {
let mut row = SparkUnsafeRow::new(schema);
- let (key_field, value_field, _) = get_map_key_value_fields(field)?;
+ let (key_field, value_field) = map_key_value_fields(field)?;
let key_type = key_field.data_type();
let value_type = value_field.data_type();
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org