|
18 | 18 | use crate::utils::make_scalar_function; |
19 | 19 | use arrow::array::builder::GenericStringBuilder; |
20 | 20 | use arrow::array::types::Int32Type; |
21 | | -use arrow::array::{Array, AsArray}; |
| 21 | +use arrow::array::{as_dictionary_array, make_array, Array, AsArray, DictionaryArray}; |
22 | 22 | use arrow::array::{ArrayRef, OffsetSizeTrait}; |
23 | 23 | use arrow::datatypes::DataType; |
24 | 24 | use datafusion::common::{cast::as_generic_string_array, DataFusionError}; |
@@ -56,6 +56,25 @@ fn spark_read_side_padding2( |
56 | 56 | DataType::LargeUtf8 => { |
57 | 57 | spark_read_side_padding_space_internal::<i64>(array, truncate, array_int) |
58 | 58 | } |
| 59 | + // Dictionary support required for SPARK-48498 |
| 60 | + DataType::Dictionary(_, value_type) => { |
| 61 | + let dict = as_dictionary_array::<Int32Type>(array); |
| 62 | + let values = if value_type.as_ref() == &DataType::Utf8 { |
| 63 | + spark_read_side_padding_space_internal::<i32>( |
| 64 | + dict.values(), |
| 65 | + truncate, |
| 66 | + array_int, |
| 67 | + )? |
| 68 | + } else { |
| 69 | + spark_read_side_padding_space_internal::<i64>( |
| 70 | + dict.values(), |
| 71 | + truncate, |
| 72 | + array_int, |
| 73 | + )? |
| 74 | + }; |
| 75 | + let result = DictionaryArray::try_new(dict.keys().clone(), values)?; |
| 76 | + Ok(make_array(result.into())) |
| 77 | + } |
59 | 78 | other => Err(DataFusionError::Internal(format!( |
60 | 79 | "Unsupported data type {other:?} for function rpad/read_side_padding", |
61 | 80 | ))), |
|
0 commit comments