Optimize Spark hex null handling

Reuse the input null buffer and split nullable and no-null encoding paths to avoid rebuilding validity and per-value Option checks. Add pointer-reuse and no-null regressions plus dedicated Criterion cases.

Author: floze-the-genius

datafusion/spark/benches/hex.rs

@@ -135,11 +135,21 @@ fn criterion_benchmark(c: &mut Criterion) {
         run_benchmark(c, "hex_utf8", size, Arc::new(data));
     }
 
+    for &size in &sizes {
+        let data = generate_utf8_data(size, 0.0);
+        run_benchmark(c, "hex_utf8_no_nulls", size, Arc::new(data));
+    }
+
     for &size in &sizes {
         let data = generate_binary_data(size, null_density);
         run_benchmark(c, "hex_binary", size, Arc::new(data));
     }
 
+    for &size in &sizes {
+        let data = generate_binary_data(size, 0.0);
+        run_benchmark(c, "hex_binary_no_nulls", size, Arc::new(data));
+    }
+
     for &size in &sizes {
         let data = generate_int64_dict_data(size, null_density);
         run_benchmark(c, "hex_int64_dict", size, Arc::new(data));

datafusion/spark/src/function/math/hex.rs

@@ -18,7 +18,7 @@
 use std::str::from_utf8_unchecked;
 use std::sync::Arc;
 
-use arrow::array::{Array, ArrayRef, NullBufferBuilder, StringArray, StringBuilder};
+use arrow::array::{Array, ArrayAccessor, ArrayRef, StringArray, StringBuilder};
 use arrow::buffer::{Buffer, OffsetBuffer};
 use arrow::datatypes::DataType;
 use arrow::{
@@ -111,21 +111,40 @@ impl ScalarUDFImpl for SparkHex {
     }
 }
 
+#[inline]
+fn append_hex_bytes(
+    values: &mut Vec<u8>,
+    bytes: &[u8],
+    case: HexCase,
+) -> Result<i32, DataFusionError> {
+    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}")
+    })?;
+    encode_bytes_into(bytes, case, values);
+    i32::try_from(values.len())
+        .map_err(|_| exec_datafusion_err!("hex output exceeds i32 offset range"))
+}
+
 /// Generic hex encoding for byte array types
-fn hex_encode_bytes<'a, I, T>(
-    iter: I,
+fn hex_encode_bytes<'a, A, T>(
+    array: &A,
     lowercase: bool,
-    len: usize,
 ) -> Result<ArrayRef, DataFusionError>
 where
-    I: Iterator<Item = Option<T>>,
-    T: AsRef<[u8]> + 'a,
+    A: ArrayAccessor<Item = &'a T>,
+    T: AsRef<[u8]> + ?Sized + 'a,
 {
     let case = if lowercase {
         HexCase::Lower
     } else {
         HexCase::Upper
     };
+    let len = array.len();
+    let nulls = array.nulls().cloned();
 
     // Write hex digits directly into one growing value buffer, tracking offsets
     // ourselves. Each input byte becomes exactly two output bytes, so there is
@@ -134,30 +153,25 @@ 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 {
-            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}"
-                )
-            })?;
-            encode_bytes_into(bytes, case, &mut values);
-            nulls.append_non_null();
-        } else {
-            nulls.append_null();
+    if let Some(ref nulls) = nulls {
+        for i in 0..len {
+            if nulls.is_valid(i) {
+                // SAFETY: `i` is in bounds and the validity buffer marks it valid.
+                let bytes = unsafe { array.value_unchecked(i) }.as_ref();
+                offsets.push(append_hex_bytes(&mut values, bytes, case)?);
+            } else {
+                offsets.push(i32::try_from(values.len()).map_err(|_| {
+                    exec_datafusion_err!("hex output exceeds i32 offset range")
+                })?);
+            }
+        }
+    } else {
+        for i in 0..len {
+            // SAFETY: `i` is in bounds and no null buffer means every value is valid.
+            let bytes = unsafe { array.value_unchecked(i) }.as_ref();
+            offsets.push(append_hex_bytes(&mut values, bytes, case)?);
         }
-        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
@@ -168,7 +182,7 @@ where
         StringArray::new_unchecked(
             OffsetBuffer::new(offsets.into()),
             Buffer::from_vec(values),
-            nulls.finish(),
+            nulls,
         )
     };
     Ok(Arc::new(array))
@@ -227,51 +241,27 @@ pub fn compute_hex(
             }
             DataType::Utf8 => {
                 let array = as_string_array(array);
-                Ok(ColumnarValue::Array(hex_encode_bytes(
-                    array.iter(),
-                    lowercase,
-                    array.len(),
-                )?))
+                Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?))
             }
             DataType::Utf8View => {
                 let array = as_string_view_array(array)?;
-                Ok(ColumnarValue::Array(hex_encode_bytes(
-                    array.iter(),
-                    lowercase,
-                    array.len(),
-                )?))
+                Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?))
             }
             DataType::LargeUtf8 => {
                 let array = as_largestring_array(array);
-                Ok(ColumnarValue::Array(hex_encode_bytes(
-                    array.iter(),
-                    lowercase,
-                    array.len(),
-                )?))
+                Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?))
             }
             DataType::Binary => {
                 let array = as_binary_array(array)?;
-                Ok(ColumnarValue::Array(hex_encode_bytes(
-                    array.iter(),
-                    lowercase,
-                    array.len(),
-                )?))
+                Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?))
             }
             DataType::LargeBinary => {
                 let array = as_large_binary_array(array)?;
-                Ok(ColumnarValue::Array(hex_encode_bytes(
-                    array.iter(),
-                    lowercase,
-                    array.len(),
-                )?))
+                Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?))
             }
             DataType::FixedSizeBinary(_) => {
                 let array = as_fixed_size_binary_array(array)?;
-                Ok(ColumnarValue::Array(hex_encode_bytes(
-                    array.iter(),
-                    lowercase,
-                    array.len(),
-                )?))
+                Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?))
             }
             DataType::Dictionary(key_type, _) => {
                 if **key_type != DataType::Int32 {
@@ -291,27 +281,27 @@ 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, lowercase)?
                     }
                     DataType::LargeUtf8 => {
                         let arr = as_largestring_array(dict_values);
-                        hex_encode_bytes(arr.iter(), lowercase, arr.len())?
+                        hex_encode_bytes(&arr, lowercase)?
                     }
                     DataType::Utf8View => {
                         let arr = as_string_view_array(dict_values)?;
-                        hex_encode_bytes(arr.iter(), lowercase, arr.len())?
+                        hex_encode_bytes(&arr, lowercase)?
                     }
                     DataType::Binary => {
                         let arr = as_binary_array(dict_values)?;
-                        hex_encode_bytes(arr.iter(), lowercase, arr.len())?
+                        hex_encode_bytes(&arr, lowercase)?
                     }
                     DataType::LargeBinary => {
                         let arr = as_large_binary_array(dict_values)?;
-                        hex_encode_bytes(arr.iter(), lowercase, arr.len())?
+                        hex_encode_bytes(&arr, lowercase)?
                     }
                     DataType::FixedSizeBinary(_) => {
                         let arr = as_fixed_size_binary_array(dict_values)?;
-                        hex_encode_bytes(arr.iter(), lowercase, arr.len())?
+                        hex_encode_bytes(&arr, lowercase)?
                     }
                     _ => {
                         return exec_err!(
@@ -465,7 +455,8 @@ mod test {
         // is reachable only via `spark_sha2_hex`, which has no in-workspace
         // caller, so it otherwise has no coverage. Drive it directly here.
         let input = StringArray::from(vec![Some("hi"), Some("bye"), None, Some("rust")]);
-        let result = super::hex_encode_bytes(input.iter(), true, input.len()).unwrap();
+        let input_ref = &input;
+        let result = super::hex_encode_bytes(&input_ref, true).unwrap();
         let result = as_string_array(&result);
 
         let expected =
@@ -495,6 +486,56 @@ mod test {
         assert_eq!(strings.value(0), expected);
     }
 
+    #[test]
+    fn test_spark_hex_binary_no_nulls() {
+        let input = BinaryArray::from(vec![
+            b"".as_slice(),
+            b"\x00\x7f\x80\xff".as_slice(),
+            b"DataFusion".as_slice(),
+        ]);
+
+        let result = super::spark_hex(&[ColumnarValue::Array(Arc::new(input))]).unwrap();
+        let array = match result {
+            ColumnarValue::Array(array) => array,
+            _ => panic!("Expected array"),
+        };
+        let strings = as_string_array(&array);
+
+        assert_eq!(strings.nulls(), None);
+        assert_eq!(
+            strings,
+            &StringArray::from(vec!["", "007F80FF", "44617461467573696F6E"])
+        );
+    }
+
+    #[test]
+    fn test_spark_hex_binary_reuses_input_nulls() {
+        let input = BinaryArray::from(vec![
+            Some(b"skip".as_slice()),
+            None,
+            Some(b"\x00\xff".as_slice()),
+            Some(b"hex".as_slice()),
+            None,
+        ])
+        .slice(1, 4);
+        let input_nulls = input.nulls().unwrap().clone();
+
+        let result = super::spark_hex(&[ColumnarValue::Array(Arc::new(input))]).unwrap();
+        let array = match result {
+            ColumnarValue::Array(array) => array,
+            _ => panic!("Expected array"),
+        };
+        let strings = as_string_array(&array);
+        let output_nulls = strings.nulls().unwrap();
+
+        assert_eq!(output_nulls, &input_nulls);
+        assert!(output_nulls.inner().ptr_eq(input_nulls.inner()));
+        assert_eq!(
+            strings,
+            &StringArray::from(vec![None, Some("00FF"), Some("686578"), None])
+        );
+    }
+
     #[test]
     fn test_spark_hex_int64() {
         let int_array = Int64Array::from(vec![Some(1), Some(2), None, Some(3)]);
@@ -540,4 +581,25 @@ mod test {
 
         assert_eq!(&expected, result);
     }
+
+    #[test]
+    fn test_dict_binary_values_null() {
+        let keys = Int32Array::from(vec![Some(0), None, Some(1)]);
+        let vals = BinaryArray::from(vec![Some(b"hi".as_slice()), None]);
+        // [b"hi", null, null]
+        let dict = DictionaryArray::new(keys, Arc::new(vals));
+
+        let result = super::spark_hex(&[ColumnarValue::Array(Arc::new(dict))]).unwrap();
+        let result = match result {
+            ColumnarValue::Array(array) => array,
+            _ => panic!("Expected array"),
+        };
+        let result = as_dictionary_array(&result).unwrap();
+
+        let keys = Int32Array::from(vec![Some(0), None, Some(1)]);
+        let vals = StringArray::from(vec![Some("6869"), None]);
+        let expected = DictionaryArray::new(keys, Arc::new(vals));
+
+        assert_eq!(&expected, result);
+    }
 }