Skip to content

Commit 1d61f31

Browse files
tlm365tlm365
authored
Add support UInt64 and other integer data types for to_hex (#16335)
* Support UInt64 data type for to_hex * Refactor unit tests * Add more test cases
1 parent e4166b3 commit 1d61f31

File tree

2 files changed

+117
-53
lines changed

2 files changed

+117
-53
lines changed

‎datafusion/functions/src/string/to_hex.rs‎

Lines changed: 112 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -19,25 +19,29 @@ use std::any::Any;
1919
use std::fmt::Write;
2020
use std::sync::Arc;
2121

22-
use arrow::array::{ArrayRef, GenericStringBuilder, OffsetSizeTrait};
22+
use crate::utils::make_scalar_function;
23+
use arrow::array::{ArrayRef, GenericStringBuilder};
24+
use arrow::datatypes::DataType::{
25+
Int16, Int32, Int64, Int8, UInt16, UInt32, UInt64, UInt8, Utf8,
26+
};
2327
use arrow::datatypes::{
24-
ArrowNativeType, ArrowPrimitiveType, DataType, Int32Type, Int64Type,
28+
ArrowNativeType, ArrowPrimitiveType, DataType, Int16Type, Int32Type, Int64Type,
29+
Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
2530
};
26-
27-
use crate::utils::make_scalar_function;
2831
use datafusion_common::cast::as_primitive_array;
2932
use datafusion_common::Result;
3033
use datafusion_common::{exec_err, plan_err};
3134

3235
use datafusion_expr::{ColumnarValue, Documentation};
3336
use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
37+
use datafusion_expr_common::signature::TypeSignature::Exact;
3438
use datafusion_macros::user_doc;
3539

3640
/// Converts the number to its equivalent hexadecimal representation.
3741
/// to_hex(2147483647) = '7fffffff'
3842
pub fn to_hex<T: ArrowPrimitiveType>(args: &[ArrayRef]) -> Result<ArrayRef>
3943
where
40-
T::Native: OffsetSizeTrait,
44+
T::Native: std::fmt::LowerHex,
4145
{
4246
let integer_array = as_primitive_array::<T>(&args[0])?;
4347

@@ -96,9 +100,20 @@ impl Default for ToHexFunc {
96100

97101
impl ToHexFunc {
98102
pub fn new() -> Self {
99-
use DataType::*;
100103
Self {
101-
signature: Signature::uniform(1, vec![Int64], Volatility::Immutable),
104+
signature: Signature::one_of(
105+
vec![
106+
Exact(vec![Int8]),
107+
Exact(vec![Int16]),
108+
Exact(vec![Int32]),
109+
Exact(vec![Int64]),
110+
Exact(vec![UInt8]),
111+
Exact(vec![UInt16]),
112+
Exact(vec![UInt32]),
113+
Exact(vec![UInt64]),
114+
],
115+
Volatility::Immutable,
116+
),
102117
}
103118
}
104119
}
@@ -117,10 +132,8 @@ impl ScalarUDFImpl for ToHexFunc {
117132
}
118133

119134
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
120-
use DataType::*;
121-
122135
Ok(match arg_types[0] {
123-
Int8 | Int16 | Int32 | Int64 => Utf8,
136+
Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 => Utf8,
124137
_ => {
125138
return plan_err!("The to_hex function can only accept integers.");
126139
}
@@ -129,12 +142,14 @@ impl ScalarUDFImpl for ToHexFunc {
129142

130143
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
131144
match args.args[0].data_type() {
132-
DataType::Int32 => {
133-
make_scalar_function(to_hex::<Int32Type>, vec![])(&args.args)
134-
}
135-
DataType::Int64 => {
136-
make_scalar_function(to_hex::<Int64Type>, vec![])(&args.args)
137-
}
145+
Int64 => make_scalar_function(to_hex::<Int64Type>, vec![])(&args.args),
146+
UInt64 => make_scalar_function(to_hex::<UInt64Type>, vec![])(&args.args),
147+
Int32 => make_scalar_function(to_hex::<Int32Type>, vec![])(&args.args),
148+
UInt32 => make_scalar_function(to_hex::<UInt32Type>, vec![])(&args.args),
149+
Int16 => make_scalar_function(to_hex::<Int16Type>, vec![])(&args.args),
150+
UInt16 => make_scalar_function(to_hex::<UInt16Type>, vec![])(&args.args),
151+
Int8 => make_scalar_function(to_hex::<Int8Type>, vec![])(&args.args),
152+
UInt8 => make_scalar_function(to_hex::<UInt8Type>, vec![])(&args.args),
138153
other => exec_err!("Unsupported data type {other:?} for function to_hex"),
139154
}
140155
}
@@ -146,48 +161,92 @@ impl ScalarUDFImpl for ToHexFunc {
146161

147162
#[cfg(test)]
148163
mod tests {
149-
use arrow::array::{Int32Array, StringArray};
150-
164+
use arrow::array::{
165+
Int16Array, Int32Array, Int64Array, Int8Array, StringArray, UInt16Array,
166+
UInt32Array, UInt64Array, UInt8Array,
167+
};
151168
use datafusion_common::cast::as_string_array;
152169

153170
use super::*;
154171

155-
#[test]
156-
// Test to_hex function for zero
157-
fn to_hex_zero() -> Result<()> {
158-
let array = vec![0].into_iter().collect::<Int32Array>();
159-
let array_ref = Arc::new(array);
160-
let hex_value_arc = to_hex::<Int32Type>(&[array_ref])?;
161-
let hex_value = as_string_array(&hex_value_arc)?;
162-
let expected = StringArray::from(vec![Some("0")]);
163-
assert_eq!(&expected, hex_value);
164-
165-
Ok(())
172+
macro_rules! test_to_hex_type {
173+
// Default test with standard input/output
174+
($name:ident, $arrow_type:ty, $array_type:ty) => {
175+
test_to_hex_type!(
176+
$name,
177+
$arrow_type,
178+
$array_type,
179+
vec![Some(100), Some(0), None],
180+
vec![Some("64"), Some("0"), None]
181+
);
182+
};
183+
184+
// Custom test with custom input/output (eg: positive number)
185+
($name:ident, $arrow_type:ty, $array_type:ty, $input:expr, $expected:expr) => {
186+
#[test]
187+
fn $name() -> Result<()> {
188+
let input = $input;
189+
let expected = $expected;
190+
191+
let array = <$array_type>::from(input);
192+
let array_ref = Arc::new(array);
193+
let hex_result = to_hex::<$arrow_type>(&[array_ref])?;
194+
let hex_array = as_string_array(&hex_result)?;
195+
let expected_array = StringArray::from(expected);
196+
197+
assert_eq!(&expected_array, hex_array);
198+
Ok(())
199+
}
200+
};
166201
}
167202

168-
#[test]
169-
// Test to_hex function for positive number
170-
fn to_hex_positive_number() -> Result<()> {
171-
let array = vec![100].into_iter().collect::<Int32Array>();
172-
let array_ref = Arc::new(array);
173-
let hex_value_arc = to_hex::<Int32Type>(&[array_ref])?;
174-
let hex_value = as_string_array(&hex_value_arc)?;
175-
let expected = StringArray::from(vec![Some("64")]);
176-
assert_eq!(&expected, hex_value);
177-
178-
Ok(())
179-
}
203+
test_to_hex_type!(
204+
to_hex_int8,
205+
Int8Type,
206+
Int8Array,
207+
vec![Some(100), Some(0), None, Some(-1)],
208+
vec![Some("64"), Some("0"), None, Some("ffffffffffffffff")]
209+
);
210+
test_to_hex_type!(
211+
to_hex_int16,
212+
Int16Type,
213+
Int16Array,
214+
vec![Some(100), Some(0), None, Some(-1)],
215+
vec![Some("64"), Some("0"), None, Some("ffffffffffffffff")]
216+
);
217+
test_to_hex_type!(
218+
to_hex_int32,
219+
Int32Type,
220+
Int32Array,
221+
vec![Some(100), Some(0), None, Some(-1)],
222+
vec![Some("64"), Some("0"), None, Some("ffffffffffffffff")]
223+
);
224+
test_to_hex_type!(
225+
to_hex_int64,
226+
Int64Type,
227+
Int64Array,
228+
vec![Some(100), Some(0), None, Some(-1)],
229+
vec![Some("64"), Some("0"), None, Some("ffffffffffffffff")]
230+
);
180231

181-
#[test]
182-
// Test to_hex function for negative number
183-
fn to_hex_negative_number() -> Result<()> {
184-
let array = vec![-1].into_iter().collect::<Int32Array>();
185-
let array_ref = Arc::new(array);
186-
let hex_value_arc = to_hex::<Int32Type>(&[array_ref])?;
187-
let hex_value = as_string_array(&hex_value_arc)?;
188-
let expected = StringArray::from(vec![Some("ffffffffffffffff")]);
189-
assert_eq!(&expected, hex_value);
190-
191-
Ok(())
192-
}
232+
test_to_hex_type!(to_hex_uint8, UInt8Type, UInt8Array);
233+
test_to_hex_type!(to_hex_uint16, UInt16Type, UInt16Array);
234+
test_to_hex_type!(to_hex_uint32, UInt32Type, UInt32Array);
235+
test_to_hex_type!(to_hex_uint64, UInt64Type, UInt64Array);
236+
237+
test_to_hex_type!(
238+
to_hex_large_signed,
239+
Int64Type,
240+
Int64Array,
241+
vec![Some(i64::MAX), Some(i64::MIN)],
242+
vec![Some("7fffffffffffffff"), Some("8000000000000000")]
243+
);
244+
245+
test_to_hex_type!(
246+
to_hex_large_unsigned,
247+
UInt64Type,
248+
UInt64Array,
249+
vec![Some(u64::MAX), Some(u64::MIN)],
250+
vec![Some("ffffffffffffffff"), Some("0")]
251+
);
193252
}

‎datafusion/sqllogictest/test_files/expr.slt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -698,6 +698,11 @@ SELECT to_hex(2147483647)
698698
----
699699
7fffffff
700700

701+
query T
702+
SELECT to_hex(CAST(2147483647 as BIGINT UNSIGNED))
703+
----
704+
7fffffff
705+
701706
query T
702707
SELECT to_hex(9223372036854775807)
703708
----

0 commit comments

Comments
 (0)