datafusion_spark/function/math/
unhex.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow::array::{Array, ArrayRef, BinaryBuilder};
19use arrow::datatypes::DataType;
20use datafusion_common::cast::{
21    as_large_string_array, as_string_array, as_string_view_array,
22};
23use datafusion_common::types::logical_string;
24use datafusion_common::utils::take_function_args;
25use datafusion_common::{
26    DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err,
27};
28use datafusion_expr::{
29    Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature,
30    TypeSignatureClass, Volatility,
31};
32use std::sync::Arc;
33
34/// <https://spark.apache.org/docs/latest/api/sql/index.html#unhex>
35#[derive(Debug, PartialEq, Eq, Hash)]
36pub struct SparkUnhex {
37    signature: Signature,
38}
39
40impl Default for SparkUnhex {
41    fn default() -> Self {
42        Self::new()
43    }
44}
45
46impl SparkUnhex {
47    pub fn new() -> Self {
48        let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string()));
49
50        Self {
51            signature: Signature::coercible(vec![string], Volatility::Immutable),
52        }
53    }
54}
55
56impl ScalarUDFImpl for SparkUnhex {
57    fn name(&self) -> &str {
58        "unhex"
59    }
60
61    fn signature(&self) -> &Signature {
62        &self.signature
63    }
64
65    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
66        Ok(DataType::Binary)
67    }
68
69    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
70        spark_unhex(&args.args)
71    }
72}
73
74#[inline]
75fn hex_nibble(c: u8) -> Option<u8> {
76    match c {
77        b'0'..=b'9' => Some(c - b'0'),
78        b'a'..=b'f' => Some(c - b'a' + 10),
79        b'A'..=b'F' => Some(c - b'A' + 10),
80        _ => None,
81    }
82}
83
84/// Decodes a hex-encoded byte slice into binary data.
85/// Returns `true` if decoding succeeded, `false` if the input contains invalid hex characters.
86fn unhex_common(bytes: &[u8], out: &mut Vec<u8>) -> bool {
87    if bytes.is_empty() {
88        return true;
89    }
90
91    let mut i = 0usize;
92
93    // If the hex string length is odd, implicitly left-pad with '0'.
94    if (bytes.len() & 1) == 1 {
95        match hex_nibble(bytes[0]) {
96            // Equivalent to (0 << 4) | lo
97            Some(lo) => out.push(lo),
98            None => return false,
99        }
100        i = 1;
101    }
102
103    while i + 1 < bytes.len() {
104        match (hex_nibble(bytes[i]), hex_nibble(bytes[i + 1])) {
105            (Some(hi), Some(lo)) => out.push((hi << 4) | lo),
106            _ => return false,
107        }
108        i += 2;
109    }
110
111    true
112}
113
114/// Converts an iterator of hex strings to a binary array.
115fn unhex_array<I, T>(
116    iter: I,
117    len: usize,
118    capacity: usize,
119) -> Result<ArrayRef, DataFusionError>
120where
121    I: Iterator<Item = Option<T>>,
122    T: AsRef<str>,
123{
124    let mut builder = BinaryBuilder::with_capacity(len, capacity);
125    let mut buffer = Vec::new();
126
127    for v in iter {
128        if let Some(s) = v {
129            buffer.clear();
130            let additional = s.as_ref().len().div_ceil(2);
131            buffer.try_reserve(additional).map_err(|e| {
132                exec_datafusion_err!(
133                    "failed to reserve {additional} bytes for unhex output: {e}"
134                )
135            })?;
136            if unhex_common(s.as_ref().as_bytes(), &mut buffer) {
137                builder.append_value(&buffer);
138            } else {
139                builder.append_null();
140            }
141        } else {
142            builder.append_null();
143        }
144    }
145
146    Ok(Arc::new(builder.finish()))
147}
148
149/// Convert a single hex string to binary
150fn unhex_scalar(s: &str) -> Option<Vec<u8>> {
151    let mut buffer = Vec::with_capacity(s.len().div_ceil(2));
152    if unhex_common(s.as_bytes(), &mut buffer) {
153        Some(buffer)
154    } else {
155        None
156    }
157}
158
159fn spark_unhex(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
160    let [args] = take_function_args("unhex", args)?;
161
162    match args {
163        ColumnarValue::Array(array) => match array.data_type() {
164            DataType::Utf8 => {
165                let array = as_string_array(array)?;
166                let capacity = array.values().len().div_ceil(2);
167                Ok(ColumnarValue::Array(unhex_array(
168                    array.iter(),
169                    array.len(),
170                    capacity,
171                )?))
172            }
173            DataType::Utf8View => {
174                let array = as_string_view_array(array)?;
175                // Estimate capacity since StringViewArray data can be scattered or inlined.
176                let capacity = array.len() * 32;
177                Ok(ColumnarValue::Array(unhex_array(
178                    array.iter(),
179                    array.len(),
180                    capacity,
181                )?))
182            }
183            DataType::LargeUtf8 => {
184                let array = as_large_string_array(array)?;
185                let capacity = array.values().len().div_ceil(2);
186                Ok(ColumnarValue::Array(unhex_array(
187                    array.iter(),
188                    array.len(),
189                    capacity,
190                )?))
191            }
192            _ => exec_err!(
193                "unhex only supports string argument, but got: {}",
194                array.data_type()
195            ),
196        },
197        ColumnarValue::Scalar(sv) => match sv {
198            ScalarValue::Utf8(None)
199            | ScalarValue::Utf8View(None)
200            | ScalarValue::LargeUtf8(None) => {
201                Ok(ColumnarValue::Scalar(ScalarValue::Binary(None)))
202            }
203            ScalarValue::Utf8(Some(s))
204            | ScalarValue::Utf8View(Some(s))
205            | ScalarValue::LargeUtf8(Some(s)) => {
206                Ok(ColumnarValue::Scalar(ScalarValue::Binary(unhex_scalar(s))))
207            }
208            _ => {
209                exec_err!(
210                    "unhex only supports string argument, but got: {}",
211                    sv.data_type()
212                )
213            }
214        },
215    }
216}
datafusion_spark/function/math/unhex.rs

datafusion_spark/function/math/
unhex.rs
