|
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}; |
@@ -94,6 +94,27 @@ fn spark_read_side_padding2( |
94 | 94 | array_pad_string, |
95 | 95 | ) |
96 | 96 | } |
| 97 | + // Dictionary support required for SPARK-48498 |
| 98 | + (DataType::Dictionary(_, value_type), DataType::Utf8) => { |
| 99 | + let dict = as_dictionary_array::<Int32Type>(array); |
| 100 | + let values = if value_type.as_ref() == &DataType::Utf8 { |
| 101 | + spark_read_side_padding_internal::<i32, i32, i32>( |
| 102 | + dict.values(), |
| 103 | + truncate, |
| 104 | + array_int, |
| 105 | + array_pad_string, |
| 106 | + )? |
| 107 | + } else { |
| 108 | + spark_read_side_padding_internal::<i64, i32, i64>( |
| 109 | + dict.values(), |
| 110 | + truncate, |
| 111 | + array_int, |
| 112 | + array_pad_string, |
| 113 | + )? |
| 114 | + }; |
| 115 | + let result = DictionaryArray::try_new(dict.keys().clone(), values)?; |
| 116 | + Ok(make_array(result.into())) |
| 117 | + } |
97 | 118 | other => Err(DataFusionError::Internal(format!( |
98 | 119 | "Unsupported data type {other:?} for function rpad/read_side_padding", |
99 | 120 | ))), |
|
0 commit comments