1818use std:: sync:: Arc ;
1919
2020use arrow:: array:: {
21- Array , BooleanArray , Capacities , MutableArrayData , Scalar , make_array ,
22- make_comparator,
21+ Array , BooleanArray , Capacities , DictionaryArray , MutableArrayData , Scalar ,
22+ make_array , make_comparator,
2323} ;
2424use arrow:: compute:: SortOptions ;
25- use arrow:: datatypes:: { DataType , Field , FieldRef } ;
25+ use arrow:: datatypes:: {
26+ DataType , Field , FieldRef , Int8Type , Int16Type , Int32Type , Int64Type , UInt8Type ,
27+ UInt16Type , UInt32Type , UInt64Type ,
28+ } ;
2629use arrow_buffer:: NullBuffer ;
2730
2831use datafusion_common:: cast:: { as_map_array, as_struct_array} ;
2932use datafusion_common:: {
30- Result , ScalarValue , exec_err, internal_err, plan_datafusion_err,
33+ Result , ScalarValue , exec_datafusion_err, exec_err, internal_datafusion_err,
34+ internal_err, plan_datafusion_err,
3135} ;
3236use datafusion_expr:: expr:: ScalarFunction ;
3337use datafusion_expr:: simplify:: ExprSimplifyResult ;
@@ -210,22 +214,21 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result<Column
210214 ( $key_ty: ty) => { {
211215 let dict = array
212216 . as_any( )
213- . downcast_ref:: <arrow :: array :: DictionaryArray <$key_ty>>( )
217+ . downcast_ref:: <DictionaryArray <$key_ty>>( )
214218 . ok_or_else( || {
215- datafusion_common:: DataFusionError :: Internal ( format!(
216- "Failed to downcast dictionary with key type {}" ,
217- key_type
218- ) )
219+ internal_datafusion_err!(
220+ "Failed to downcast dictionary with key type {key_type}"
221+ )
219222 } ) ?;
220223 let values_struct = as_struct_array( dict. values( ) ) ?;
221224 let field_col =
222225 values_struct. column_by_name( & field_name) . ok_or_else( || {
223- datafusion_common :: DataFusionError :: Execution ( format !(
226+ exec_datafusion_err !(
224227 "Field {field_name} not found in dictionary struct"
225- ) )
228+ )
226229 } ) ?;
227230 // Rebuild dictionary: same keys, extracted field as values.
228- let new_dict = arrow :: array :: DictionaryArray :: <$key_ty>:: try_new(
231+ let new_dict = DictionaryArray :: <$key_ty>:: try_new(
229232 dict. keys( ) . clone( ) ,
230233 Arc :: clone( field_col) ,
231234 ) ?;
@@ -234,14 +237,14 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result<Column
234237 }
235238
236239 match key_type. as_ref ( ) {
237- DataType :: Int8 => extract_dict_field ! ( arrow :: datatypes :: Int8Type ) ,
238- DataType :: Int16 => extract_dict_field ! ( arrow :: datatypes :: Int16Type ) ,
239- DataType :: Int32 => extract_dict_field ! ( arrow :: datatypes :: Int32Type ) ,
240- DataType :: Int64 => extract_dict_field ! ( arrow :: datatypes :: Int64Type ) ,
241- DataType :: UInt8 => extract_dict_field ! ( arrow :: datatypes :: UInt8Type ) ,
242- DataType :: UInt16 => extract_dict_field ! ( arrow :: datatypes :: UInt16Type ) ,
243- DataType :: UInt32 => extract_dict_field ! ( arrow :: datatypes :: UInt32Type ) ,
244- DataType :: UInt64 => extract_dict_field ! ( arrow :: datatypes :: UInt64Type ) ,
240+ DataType :: Int8 => extract_dict_field ! ( Int8Type ) ,
241+ DataType :: Int16 => extract_dict_field ! ( Int16Type ) ,
242+ DataType :: Int32 => extract_dict_field ! ( Int32Type ) ,
243+ DataType :: Int64 => extract_dict_field ! ( Int64Type ) ,
244+ DataType :: UInt8 => extract_dict_field ! ( UInt8Type ) ,
245+ DataType :: UInt16 => extract_dict_field ! ( UInt16Type ) ,
246+ DataType :: UInt32 => extract_dict_field ! ( UInt32Type ) ,
247+ DataType :: UInt64 => extract_dict_field ! ( UInt64Type ) ,
245248 other => exec_err ! ( "Unsupported dictionary key type: {other}" ) ,
246249 }
247250 }
@@ -395,9 +398,7 @@ impl ScalarUDFImpl for GetFieldFunc {
395398 sv. try_as_str ( ) . flatten ( ) . filter ( |s| !s. is_empty ( ) )
396399 } )
397400 . ok_or_else ( || {
398- datafusion_common:: DataFusionError :: Execution (
399- "Field name must be a non-empty string" . to_string ( ) ,
400- )
401+ exec_datafusion_err ! ( "Field name must be a non-empty string" )
401402 } ) ?;
402403
403404 let child_field = fields
0 commit comments