Skip to content

Commit 9c6b42e

Browse files
Add support to cast from UnionArray (#9544)
# Which issue does this PR close? - Closes #9225 # Rationale for this change This PR adds native support for casting Union arrays in the cast kernel. Previously, `can_cast_types` and `cast_with_options` had no handling for union types at all
1 parent ec771cc commit 9c6b42e

File tree

2 files changed

+499
-1
lines changed

2 files changed

+499
-1
lines changed

arrow-cast/src/cast/mod.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,13 +43,15 @@ mod list;
4343
mod map;
4444
mod run_array;
4545
mod string;
46+
mod union;
4647

4748
use crate::cast::decimal::*;
4849
use crate::cast::dictionary::*;
4950
use crate::cast::list::*;
5051
use crate::cast::map::*;
5152
use crate::cast::run_array::*;
5253
use crate::cast::string::*;
54+
pub use crate::cast::union::*;
5355

5456
use arrow_buffer::IntervalMonthDayNano;
5557
use arrow_data::ByteView;
@@ -108,6 +110,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
108110
can_cast_types(from_value_type, to_value_type)
109111
}
110112
(Dictionary(_, value_type), _) => can_cast_types(value_type, to_type),
113+
(Union(fields, _), _) => union::resolve_child_array(fields, to_type).is_some(),
114+
(_, Union(_, _)) => false,
111115
(RunEndEncoded(_, value_type), _) => can_cast_types(value_type.data_type(), to_type),
112116
(_, RunEndEncoded(_, value_type)) => can_cast_types(from_type, value_type.data_type()),
113117
(_, Dictionary(_, value_type)) => can_cast_types(from_type, value_type),
@@ -230,7 +234,6 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
230234
}
231235
(Struct(_), _) => false,
232236
(_, Struct(_)) => false,
233-
234237
(_, Boolean) => from_type.is_integer() || from_type.is_floating() || from_type.is_string(),
235238
(Boolean, _) => to_type.is_integer() || to_type.is_floating() || to_type.is_string(),
236239

@@ -781,6 +784,14 @@ pub fn cast_with_options(
781784
))),
782785
}
783786
}
787+
(Union(_, _), _) => union_extract_by_type(
788+
array.as_any().downcast_ref::<UnionArray>().unwrap(),
789+
to_type,
790+
cast_options,
791+
),
792+
(_, Union(_, _)) => Err(ArrowError::CastError(format!(
793+
"Casting from {from_type} to {to_type} not supported"
794+
))),
784795
(Dictionary(index_type, _), _) => match **index_type {
785796
Int8 => dictionary_cast::<Int8Type>(array, to_type, cast_options),
786797
Int16 => dictionary_cast::<Int16Type>(array, to_type, cast_options),

0 commit comments

Comments
 (0)