Skip to content

Commit 174068e

Browse files
committed
Use error macros and import types instead of inline resolving
1 parent 3a58334 commit 174068e

1 file changed

Lines changed: 24 additions & 23 deletions

File tree

datafusion/functions/src/core/getfield.rs

Lines changed: 24 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,20 @@
1818
use std::sync::Arc;
1919

2020
use 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
};
2424
use 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+
};
2629
use arrow_buffer::NullBuffer;
2730

2831
use datafusion_common::cast::{as_map_array, as_struct_array};
2932
use 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
};
3236
use datafusion_expr::expr::ScalarFunction;
3337
use 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

Comments
 (0)