Skip to content

Commit 2d86c7f

Browse files
Weijun-Happletreeisyellow
authored andcommitted
Refactor array_union function to use a generic (apache#8381)
union_arrays function
1 parent b63f403 commit 2d86c7f

File tree

1 file changed

+21
-20
lines changed

1 file changed

+21
-20
lines changed

datafusion/physical-expr/src/array_expressions.rs

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1525,32 +1525,33 @@ pub fn array_union(args: &[ArrayRef]) -> Result<ArrayRef> {
15251525
}
15261526
let array1 = &args[0];
15271527
let array2 = &args[1];
1528+
1529+
fn union_arrays<O: OffsetSizeTrait>(
1530+
array1: &ArrayRef,
1531+
array2: &ArrayRef,
1532+
l_field_ref: &Arc<Field>,
1533+
r_field_ref: &Arc<Field>,
1534+
) -> Result<ArrayRef> {
1535+
match (l_field_ref.data_type(), r_field_ref.data_type()) {
1536+
(DataType::Null, _) => Ok(array2.clone()),
1537+
(_, DataType::Null) => Ok(array1.clone()),
1538+
(_, _) => {
1539+
let list1 = array1.as_list::<O>();
1540+
let list2 = array2.as_list::<O>();
1541+
let result = union_generic_lists::<O>(list1, list2, l_field_ref)?;
1542+
Ok(Arc::new(result))
1543+
}
1544+
}
1545+
}
1546+
15281547
match (array1.data_type(), array2.data_type()) {
15291548
(DataType::Null, _) => Ok(array2.clone()),
15301549
(_, DataType::Null) => Ok(array1.clone()),
15311550
(DataType::List(l_field_ref), DataType::List(r_field_ref)) => {
1532-
match (l_field_ref.data_type(), r_field_ref.data_type()) {
1533-
(DataType::Null, _) => Ok(array2.clone()),
1534-
(_, DataType::Null) => Ok(array1.clone()),
1535-
(_, _) => {
1536-
let list1 = array1.as_list::<i32>();
1537-
let list2 = array2.as_list::<i32>();
1538-
let result = union_generic_lists::<i32>(list1, list2, l_field_ref)?;
1539-
Ok(Arc::new(result))
1540-
}
1541-
}
1551+
union_arrays::<i32>(array1, array2, l_field_ref, r_field_ref)
15421552
}
15431553
(DataType::LargeList(l_field_ref), DataType::LargeList(r_field_ref)) => {
1544-
match (l_field_ref.data_type(), r_field_ref.data_type()) {
1545-
(DataType::Null, _) => Ok(array2.clone()),
1546-
(_, DataType::Null) => Ok(array1.clone()),
1547-
(_, _) => {
1548-
let list1 = array1.as_list::<i64>();
1549-
let list2 = array2.as_list::<i64>();
1550-
let result = union_generic_lists::<i64>(list1, list2, l_field_ref)?;
1551-
Ok(Arc::new(result))
1552-
}
1553-
}
1554+
union_arrays::<i64>(array1, array2, l_field_ref, r_field_ref)
15541555
}
15551556
_ => {
15561557
internal_err!(

0 commit comments

Comments
 (0)