Skip to content

perf: reuse input null buffer and add no-nulls fast path in Spark hex #23721

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jackylee-ch wants to merge 1 commit into from
+101 −19
Closed

perf: reuse input null buffer and add no-nulls fast path in Spark hex #23721

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 101 additions & 19 deletions datafusion/spark/src/function/math/hex.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@
use std::str::from_utf8_unchecked;
use std::sync::Arc;

use arrow::array::{Array, ArrayRef, NullBufferBuilder, StringArray, StringBuilder};
use arrow::buffer::{Buffer, OffsetBuffer};
use arrow::array::{Array, ArrayRef, StringArray, StringBuilder};
use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer};
use arrow::datatypes::DataType;
use arrow::{
array::{as_dictionary_array, as_largestring_array, as_string_array},
Expand Down Expand Up @@ -159,8 +159,13 @@ fn hex_int64(num: i64, buffer: &mut [u8; 16]) -> &[u8] {
}

/// Generic hex encoding for byte array types
///
/// `nulls` is the input array's null buffer. Hex encoding maps rows 1:1 (a
/// null in maps to a null out), so it is reused for the output array as-is
/// instead of being rebuilt row by row.
fn hex_encode_bytes<'a, I, T>(
iter: I,
nulls: Option<NullBuffer>,
lowercase: bool,
len: usize,
) -> Result<ArrayRef, DataFusionError>
Expand All @@ -181,10 +186,11 @@ where
let mut values: Vec<u8> = Vec::with_capacity(len * 64);
let mut offsets: Vec<i32> = Vec::with_capacity(len + 1);
offsets.push(0);
let mut nulls = NullBufferBuilder::new(len);

for v in iter {
if let Some(b) = v {
if nulls.is_none() {
// Fast path: the input has no nulls, so every item is `Some` and the
// per-row `Option` check and null bookkeeping can be skipped.
for b in iter.flatten() {
let bytes = b.as_ref();
let additional = bytes
.len()
Expand All @@ -198,26 +204,46 @@ where
for &byte in bytes {
values.extend_from_slice(&lookup[byte as usize]);
}
nulls.append_non_null();
} else {
nulls.append_null();
offsets.push(i32::try_from(values.len()).map_err(|_| {
exec_datafusion_err!("hex output exceeds i32 offset range")
})?);
}
offsets.push(
i32::try_from(values.len()).map_err(|_| {
} else {
for v in iter {
if let Some(b) = v {
let bytes = b.as_ref();
let additional = bytes
.len()
.checked_mul(2)
.ok_or_else(|| exec_datafusion_err!("hex output size overflow"))?;
values.try_reserve(additional).map_err(|e| {
exec_datafusion_err!(
"failed to reserve {additional} bytes for hex output: {e}"
)
})?;
for &byte in bytes {
values.extend_from_slice(&lookup[byte as usize]);
}
}
// For a null slot no hex digits are written, so the row becomes an
// empty value; the reused input null buffer marks it as null.
offsets.push(i32::try_from(values.len()).map_err(|_| {
exec_datafusion_err!("hex output exceeds i32 offset range")
})?,
);
})?);
}
}

// SAFETY: the value buffer contains only ASCII hex digits (valid UTF-8) and
// the offsets are monotonically increasing and end at `values.len()`, so the
// array invariants hold. This mirrors the previous `from_utf8_unchecked`
// path and avoids a redundant UTF-8 validation pass over the whole buffer.
// The reused null buffer has the input array's length, matching the number
// of encoded rows.
let array = unsafe {
StringArray::new_unchecked(
OffsetBuffer::new(offsets.into()),
Buffer::from_vec(values),
nulls.finish(),
nulls,
)
};
Ok(Arc::new(array))
Expand Down Expand Up @@ -278,6 +304,7 @@ pub fn compute_hex(
let array = as_string_array(array);
Ok(ColumnarValue::Array(hex_encode_bytes(
array.iter(),
array.nulls().cloned(),
lowercase,
array.len(),
)?))
Expand All @@ -286,6 +313,7 @@ pub fn compute_hex(
let array = as_string_view_array(array)?;
Ok(ColumnarValue::Array(hex_encode_bytes(
array.iter(),
array.nulls().cloned(),
lowercase,
array.len(),
)?))
Expand All @@ -294,6 +322,7 @@ pub fn compute_hex(
let array = as_largestring_array(array);
Ok(ColumnarValue::Array(hex_encode_bytes(
array.iter(),
array.nulls().cloned(),
lowercase,
array.len(),
)?))
Expand All @@ -302,6 +331,7 @@ pub fn compute_hex(
let array = as_binary_array(array)?;
Ok(ColumnarValue::Array(hex_encode_bytes(
array.iter(),
array.nulls().cloned(),
lowercase,
array.len(),
)?))
Expand All @@ -310,6 +340,7 @@ pub fn compute_hex(
let array = as_large_binary_array(array)?;
Ok(ColumnarValue::Array(hex_encode_bytes(
array.iter(),
array.nulls().cloned(),
lowercase,
array.len(),
)?))
Expand All @@ -318,6 +349,7 @@ pub fn compute_hex(
let array = as_fixed_size_binary_array(array)?;
Ok(ColumnarValue::Array(hex_encode_bytes(
array.iter(),
array.nulls().cloned(),
lowercase,
array.len(),
)?))
Expand All @@ -340,27 +372,57 @@ pub fn compute_hex(
}
DataType::Utf8 => {
let arr = as_string_array(dict_values);
hex_encode_bytes(arr.iter(), lowercase, arr.len())?
hex_encode_bytes(
arr.iter(),
arr.nulls().cloned(),
lowercase,
arr.len(),
)?
}
DataType::LargeUtf8 => {
let arr = as_largestring_array(dict_values);
hex_encode_bytes(arr.iter(), lowercase, arr.len())?
hex_encode_bytes(
arr.iter(),
arr.nulls().cloned(),
lowercase,
arr.len(),
)?
}
DataType::Utf8View => {
let arr = as_string_view_array(dict_values)?;
hex_encode_bytes(arr.iter(), lowercase, arr.len())?
hex_encode_bytes(
arr.iter(),
arr.nulls().cloned(),
lowercase,
arr.len(),
)?
}
DataType::Binary => {
let arr = as_binary_array(dict_values)?;
hex_encode_bytes(arr.iter(), lowercase, arr.len())?
hex_encode_bytes(
arr.iter(),
arr.nulls().cloned(),
lowercase,
arr.len(),
)?
}
DataType::LargeBinary => {
let arr = as_large_binary_array(dict_values)?;
hex_encode_bytes(arr.iter(), lowercase, arr.len())?
hex_encode_bytes(
arr.iter(),
arr.nulls().cloned(),
lowercase,
arr.len(),
)?
}
DataType::FixedSizeBinary(_) => {
let arr = as_fixed_size_binary_array(dict_values)?;
hex_encode_bytes(arr.iter(), lowercase, arr.len())?
hex_encode_bytes(
arr.iter(),
arr.nulls().cloned(),
lowercase,
arr.len(),
)?
}
_ => {
return exec_err!(
Expand Down Expand Up @@ -554,6 +616,26 @@ mod test {
assert_eq!(strings.value(0), expected);
}

#[test]
fn test_spark_hex_utf8_with_nulls() {
// Exercises the has-nulls path: the output must reuse the input's
// null buffer and keep values aligned around the null slots.
let input =
StringArray::from(vec![Some("hi"), None, Some(""), None, Some("rust")]);

let result = super::spark_hex(&[ColumnarValue::Array(Arc::new(input))]).unwrap();
let result = match result {
ColumnarValue::Array(array) => array,
_ => panic!("Expected array"),
};

let string_array = as_string_array(&result);
let expected =
StringArray::from(vec![Some("6869"), None, Some(""), None, Some("72757374")]);

assert_eq!(string_array, &expected);
}

#[test]
fn test_spark_hex_int64() {
let int_array = Int64Array::from(vec![Some(1), Some(2), None, Some(3)]);
Expand Down
Loading