fix: handle large strings

apache · May 29, 2024 · 129f00a · 129f00a
1 parent e7062c6
commit 129f00a
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/core/src/execution/datafusion/expressions/scalar_funcs/hex.rs b/core/src/execution/datafusion/expressions/scalar_funcs/hex.rs
@@ -18,7 +18,7 @@
 use std::sync::Arc;
 
 use arrow::{
-    array::{as_dictionary_array, as_string_array},
+    array::{as_dictionary_array, as_largestring_array, as_string_array},
     datatypes::Int32Type,
 };
 use arrow_array::StringArray;
@@ -60,7 +60,7 @@ pub(super) fn spark_hex(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFus
 
                 Ok(ColumnarValue::Array(Arc::new(hexed_array)))
             }
-            DataType::Utf8 | DataType::LargeUtf8 => {
+            DataType::Utf8 => {
                 let array = as_string_array(array);
 
                 let hexed: StringArray = array
@@ -70,6 +70,16 @@ pub(super) fn spark_hex(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFus
 
                 Ok(ColumnarValue::Array(Arc::new(hexed)))
             }
+            DataType::LargeUtf8 => {
+                let array = as_largestring_array(array);
+
+                let hexed: StringArray = array
+                    .iter()
+                    .map(|v| v.map(hex_bytes).transpose())
+                    .collect::<Result<_, _>>()?;
+
+                Ok(ColumnarValue::Array(Arc::new(hexed)))
+            }
             DataType::Binary => {
                 let array = as_binary_array(array)?;