perf: avoid per-row copy in Spark hex byte encoding

Author: andygrove

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

@@ -18,7 +18,8 @@
 use std::str::from_utf8_unchecked;
 use std::sync::Arc;
 
-use arrow::array::{Array, ArrayRef, StringBuilder};
+use arrow::array::{Array, ArrayRef, NullBufferBuilder, StringArray, StringBuilder};
+use arrow::buffer::{Buffer, OffsetBuffer};
 use arrow::datatypes::DataType;
 use arrow::{
     array::{as_dictionary_array, as_largestring_array, as_string_array},
@@ -166,40 +167,59 @@ where
     I: Iterator<Item = Option<T>>,
     T: AsRef<[u8]> + 'a,
 {
-    let mut builder = StringBuilder::with_capacity(len, len * 64);
-    let mut buffer = Vec::with_capacity(64);
     let lookup = if lowercase {
         &HEX_LOOKUP_LOWER
     } else {
         &HEX_LOOKUP_UPPER
     };
 
+    // Write hex digits directly into one growing value buffer, tracking offsets
+    // ourselves. Each input byte becomes exactly two output bytes, so there is
+    // no per-row `String`/`StringBuilder` copy — the hex digits are written once
+    // into the final buffer.
+    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();
-            buffer.clear();
             let additional = bytes
                 .len()
                 .checked_mul(2)
                 .ok_or_else(|| exec_datafusion_err!("hex output size overflow"))?;
-            buffer.try_reserve(additional).map_err(|e| {
+            values.try_reserve(additional).map_err(|e| {
                 exec_datafusion_err!(
                     "failed to reserve {additional} bytes for hex output: {e}"
                 )
             })?;
             for &byte in bytes {
-                buffer.extend_from_slice(&lookup[byte as usize]);
-            }
-            // SAFETY: buffer contains only ASCII hex digits, which are valid UTF-8.
-            unsafe {
-                builder.append_value(from_utf8_unchecked(&buffer));
+                values.extend_from_slice(&lookup[byte as usize]);
             }
+            nulls.append_non_null();
         } else {
-            builder.append_null();
+            nulls.append_null();
         }
+        offsets.push(
+            i32::try_from(values.len()).map_err(|_| {
+                exec_datafusion_err!("hex output exceeds i32 offset range")
+            })?,
+        );
     }
 
-    Ok(Arc::new(builder.finish()))
+    // 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.
+    let array = unsafe {
+        StringArray::new_unchecked(
+            OffsetBuffer::new(offsets.into()),
+            Buffer::from_vec(values),
+            nulls.finish(),
+        )
+    };
+    Ok(Arc::new(array))
 }
 
 /// Generic hex encoding for int64 type