diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 42a7bcb97b..4addc026d8 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -124,6 +124,10 @@ harness = false name = "to_time" harness = false +[[bench]] +name = "unhex" +harness = false + [[bench]] name = "array_size" harness = false diff --git a/native/spark-expr/benches/unhex.rs b/native/spark-expr/benches/unhex.rs new file mode 100644 index 0000000000..9aca65a8de --- /dev/null +++ b/native/spark-expr/benches/unhex.rs @@ -0,0 +1,109 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::builder::StringBuilder; +use arrow::array::ArrayRef; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::spark_unhex; +use std::hint::black_box; +use std::sync::Arc; + +/// Build a Utf8 array of `size` rows by cycling through `samples`, inserting a +/// null every `null_every` rows (0 = no nulls). +fn build(samples: &[&str], size: usize, null_every: usize) -> ArrayRef { + let mut builder = StringBuilder::new(); + for i in 0..size { + if null_every != 0 && i % null_every == 0 { + builder.append_null(); + } else { + builder.append_value(samples[i % samples.len()]); + } + } + Arc::new(builder.finish()) +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 8192; + + // Dense column of valid, even-length hex strings, no nulls. + let all_valid = build( + &[ + "537061726B2053514C", + "737472696E67", + "1C", + "deadbeef", + "0011ccddeeff", + ], + size, + 0, + ); + // Same shape but with ~6% nulls interleaved. + let with_nulls = build( + &[ + "537061726B2053514C", + "737472696E67", + "1C", + "deadbeef", + "0011ccddeeff", + ], + size, + 17, + ); + // Long valid hex strings, exercising the decode loop and allocation. + let long_valid = build( + &[ + "0011223344556677889900aabbccddeeff0011223344556677889900aabbccddeeff", + "537061726b2053514c537061726b2053514c537061726b2053514c537061726b2053514c", + ], + size, + 0, + ); + // Mostly invalid inputs (odd length or non-hex), which decode to null. + let invalid = build(&["ZZ", "A1B", "xyz", "12G4", "gg"], size, 0); + // A diverse mix of the above shapes plus empty strings. + let mixed = build( + &[ + "537061726B2053514C", + "1C", + "A1B", + "737472696E67", + "0011223344556677889900aabbccddeeff", + "deadBEEF", + "ZZ", + "", + ], + size, + 17, + ); + + let mut bench = |name: &str, arr: &ArrayRef| { + let args = vec![ColumnarValue::Array(Arc::clone(arr))]; + c.bench_function(name, |b| { + b.iter(|| black_box(spark_unhex(black_box(&args)).unwrap())) + }); + }; + + bench("spark_unhex: all valid", &all_valid); + bench("spark_unhex: with nulls", &with_nulls); + bench("spark_unhex: long strings", &long_valid); + bench("spark_unhex: invalid inputs", &invalid); + bench("spark_unhex: mixed hex column", &mixed); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/math_funcs/unhex.rs b/native/spark-expr/src/math_funcs/unhex.rs index 46f899fcc2..b8292bd8b3 100644 --- a/native/spark-expr/src/math_funcs/unhex.rs +++ b/native/spark-expr/src/math_funcs/unhex.rs @@ -17,40 +17,66 @@ use std::sync::Arc; -use arrow::array::OffsetSizeTrait; +use arrow::array::{Array, BinaryBuilder, OffsetSizeTrait}; use arrow::datatypes::DataType; use datafusion::common::{cast::as_generic_string_array, exec_err, DataFusionError, ScalarValue}; use datafusion::logical_expr::ColumnarValue; -/// Helper function to convert a hex digit to a binary value. -fn unhex_digit(c: u8) -> Result { - match c { - b'0'..=b'9' => Ok(c - b'0'), - b'A'..=b'F' => Ok(10 + c - b'A'), - b'a'..=b'f' => Ok(10 + c - b'a'), - _ => Err(DataFusionError::Execution( - "Input to unhex_digit is not a valid hex digit".to_string(), - )), +/// Sentinel stored in `HEX_LUT` for bytes that are not valid hex digits. Its value `0xFF` has all +/// bits set, so OR-ing the two nibbles of a byte and comparing the result against `INVALID_HEX` +/// rejects the byte whenever either nibble is invalid: valid nibbles are `<= 0x0F`, so a valid +/// pair can never OR up to `0xFF`. This lets the decode loop test both nibbles with one comparison. +const INVALID_HEX: u8 = 0xFF; + +/// Lookup table mapping each possible input byte to its hex value, or `INVALID_HEX` for bytes +/// that are not `0-9`, `A-F`, or `a-f`. Built at compile time so decoding is a single indexed +/// load per digit rather than a chain of range comparisons. +const HEX_LUT: [u8; 256] = build_hex_lut(); + +const fn build_hex_lut() -> [u8; 256] { + let mut lut = [INVALID_HEX; 256]; + let mut i = 0; + while i < 256 { + lut[i] = match i as u8 { + c @ b'0'..=b'9' => c - b'0', + c @ b'A'..=b'F' => c - b'A' + 10, + c @ b'a'..=b'f' => c - b'a' + 10, + _ => INVALID_HEX, + }; + i += 1; } + lut +} + +#[inline] +fn invalid_hex_digit() -> DataFusionError { + DataFusionError::Execution("Input to unhex is not a valid hex digit".to_string()) } /// Convert a hex string to binary and store the result in `result`. Returns an error if the input /// is not a valid hex string. fn unhex(hex_str: &str, result: &mut Vec) -> Result<(), DataFusionError> { let bytes = hex_str.as_bytes(); + // Each output byte consumes two hex digits (the optional leading nibble rounds up). + result.reserve(bytes.len().div_ceil(2)); let mut i = 0; if (bytes.len() & 0x01) != 0 { - let v = unhex_digit(bytes[0])?; - + let v = HEX_LUT[bytes[0] as usize]; + if v == INVALID_HEX { + return Err(invalid_hex_digit()); + } result.push(v); i += 1; } while i < bytes.len() { - let first = unhex_digit(bytes[i])?; - let second = unhex_digit(bytes[i + 1])?; + let first = HEX_LUT[bytes[i] as usize]; + let second = HEX_LUT[bytes[i + 1] as usize]; + if (first | second) == INVALID_HEX { + return Err(invalid_hex_digit()); + } result.push((first << 4) | second); i += 2; @@ -68,7 +94,13 @@ fn spark_unhex_inner( let string_array = as_generic_string_array::(array)?; let mut encoded = Vec::new(); - let mut builder = arrow::array::BinaryBuilder::new(); + // Every two input hex digits produce one output byte, so the decoded data is at most + // half the total input length. Preallocating both the offset and value buffers avoids + // the repeated doublings a fresh `BinaryBuilder` would incur across the whole column. + let mut builder = BinaryBuilder::with_capacity( + string_array.len(), + string_array.value_data().len() / 2, + ); for item in string_array.iter() { if let Some(s) = item {