Skip to content
Open
Changes from all commits
Commits
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
44 changes: 32 additions & 12 deletions datafusion/spark/src/function/math/hex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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
Expand Down
Loading