diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 6ecd41b0b9a5c..d6c5b054cd858 100644 --- a/datafusion/functions/src/string/common.rs +++ b/datafusion/functions/src/string/common.rs @@ -24,7 +24,7 @@ use crate::strings::{ StringViewArrayBuilder, append_view, }; use arrow::array::{ - Array, ArrayRef, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, AsArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, StringViewArray, new_null_array, }; use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; @@ -387,6 +387,21 @@ fn case_conversion( Ok(ColumnarValue::Array(Arc::new(builder.finish(nulls)?))) } + DataType::Dictionary(_, _) => { + let dict = array.as_any_dictionary(); + let values = dict.values(); + let converted = case_conversion( + &[ColumnarValue::Array(Arc::clone(values))], + lower, + name, + )?; + match converted { + ColumnarValue::Array(new_values) => { + Ok(ColumnarValue::Array(dict.with_values(new_values))) + } + _ => unreachable!("Array input should produce Array output"), + } + } other => exec_err!("Unsupported data type {other:?} for function {name}"), }, ColumnarValue::Scalar(scalar) => match scalar { diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 57cbe1d8779f0..f4c25bc35da12 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -91,7 +91,9 @@ impl ScalarUDFImpl for LowerFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, ArrayRef, StringArray, StringViewArray}; + use arrow::array::{ + Array, ArrayRef, DictionaryArray, StringArray, StringViewArray, UInt8Array, + }; use arrow::datatypes::Field; use datafusion_common::config::ConfigOptions; use std::sync::Arc; @@ -370,4 +372,22 @@ mod tests { assert_eq!(result_sa.value_data().len(), 10); Ok(()) } + + #[test] + fn test_dict() -> Result<()> { + let input = Arc::new(DictionaryArray::new( + UInt8Array::from_iter_values([0, 1, 1, 0, 1, 0]), + Arc::new(StringArray::from_iter_values(["A", "b"])), + )); + + let result = invoke_lower(input)?; + + let expected = DictionaryArray::new( + UInt8Array::from_iter_values([0, 1, 1, 0, 1, 0]), + Arc::new(StringArray::from_iter_values(["a", "b"])), + ); + + assert_eq!(result.as_ref(), &expected); + Ok(()) + } } diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index c0ac90b1bc598..a7ee12b4c6e49 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -90,7 +90,9 @@ impl ScalarUDFImpl for UpperFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, ArrayRef, StringArray, StringViewArray}; + use arrow::array::{ + Array, ArrayRef, DictionaryArray, StringArray, StringViewArray, UInt8Array, + }; use arrow::datatypes::Field; use datafusion_common::config::ConfigOptions; use std::sync::Arc; @@ -369,4 +371,22 @@ mod tests { assert_eq!(result_sa.value_data().len(), 10); Ok(()) } + + #[test] + fn test_dict() -> Result<()> { + let input = Arc::new(DictionaryArray::new( + UInt8Array::from_iter_values([0, 1, 1, 0, 1, 0]), + Arc::new(StringArray::from_iter_values(["A", "b"])), + )); + + let result = invoke_upper(input)?; + + let expected = DictionaryArray::new( + UInt8Array::from_iter_values([0, 1, 1, 0, 1, 0]), + Arc::new(StringArray::from_iter_values(["A", "B"])), + ); + + assert_eq!(result.as_ref(), &expected); + Ok(()) + } }