This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 79721ec408 Minor: support cast values to fixedsizelist (#5340)
79721ec408 is described below
commit 79721ec40884e66cbbbd2aa4f301b7c03a0aed5e
Author: Alex Huang <[email protected]>
AuthorDate: Sun Feb 4 22:11:30 2024 +0800
Minor: support cast values to fixedsizelist (#5340)
* support cast utf8 to fixedsizelist
* fix ci
* fix chore
* Apply suggestions from code review
Co-authored-by: Raphael Taylor-Davies
<[email protected]>
* fix compile
---------
Co-authored-by: Andrew Lamb <[email protected]>
Co-authored-by: Raphael Taylor-Davies
<[email protected]>
---
arrow-cast/src/cast.rs | 41 +++++++++++++++++++++++++++++++++++++++++
1 file changed, 41 insertions(+)
diff --git a/arrow-cast/src/cast.rs b/arrow-cast/src/cast.rs
index bd35096e06..e321ae1b1b 100644
--- a/arrow-cast/src/cast.rs
+++ b/arrow-cast/src/cast.rs
@@ -141,6 +141,8 @@ pub fn can_cast_types(from_type: &DataType, to_type:
&DataType) -> bool {
}
(_, List(list_to)) => can_cast_types(from_type, list_to.data_type()),
(_, LargeList(list_to)) => can_cast_types(from_type,
list_to.data_type()),
+ (_, FixedSizeList(list_to,size)) if *size == 1 => {
+ can_cast_types(from_type, list_to.data_type())},
// cast one decimal type to another decimal type
(Decimal128(_, _), Decimal128(_, _)) => true,
(Decimal256(_, _), Decimal256(_, _)) => true,
@@ -802,6 +804,9 @@ pub fn cast_with_options(
}
(_, List(ref to)) => cast_values_to_list::<i32>(array, to,
cast_options),
(_, LargeList(ref to)) => cast_values_to_list::<i64>(array, to,
cast_options),
+ (_, FixedSizeList(ref to, size)) if *size == 1 => {
+ cast_values_to_fixed_size_list(array, to, *size, cast_options)
+ }
(Decimal128(_, s1), Decimal128(p2, s2)) => {
cast_decimal_to_decimal_same_type::<Decimal128Type>(
array.as_primitive(),
@@ -3040,6 +3045,18 @@ fn cast_values_to_list<O: OffsetSizeTrait>(
Ok(Arc::new(list))
}
+/// Helper function that takes a primitive array and casts to a fixed size
list array.
+fn cast_values_to_fixed_size_list(
+ array: &dyn Array,
+ to: &FieldRef,
+ size: i32,
+ cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError> {
+ let values = cast_with_options(array, to.data_type(), cast_options)?;
+ let list = FixedSizeListArray::new(to.clone(), size, values, None);
+ Ok(Arc::new(list))
+}
+
/// A specified helper to cast from `GenericBinaryArray` to
`GenericStringArray` when they have same
/// offset size so re-encoding offset is unnecessary.
fn cast_binary_to_string<O: OffsetSizeTrait>(
@@ -7609,6 +7626,30 @@ mod tests {
assert_eq!(expected.values(), actual.values());
}
+ #[test]
+ fn test_cast_utf8_to_list() {
+ // DataType::List
+ let array = Arc::new(StringArray::from(vec!["5"])) as ArrayRef;
+ let field = Arc::new(Field::new("", DataType::Int32, false));
+ let list_array = cast(&array, &DataType::List(field.clone())).unwrap();
+ let actual = list_array.as_list_opt::<i32>().unwrap();
+ let expect = ListArray::from_iter_primitive::<Int32Type, _,
_>([Some([Some(5)])]);
+ assert_eq!(&expect.value(0), &actual.value(0));
+
+ // DataType::LargeList
+ let list_array = cast(&array,
&DataType::LargeList(field.clone())).unwrap();
+ let actual = list_array.as_list_opt::<i64>().unwrap();
+ let expect = LargeListArray::from_iter_primitive::<Int32Type, _,
_>([Some([Some(5)])]);
+ assert_eq!(&expect.value(0), &actual.value(0));
+
+ // DataType::FixedSizeList
+ let list_array = cast(&array, &DataType::FixedSizeList(field.clone(),
1)).unwrap();
+ let actual = list_array.as_fixed_size_list_opt().unwrap();
+ let expect =
+ FixedSizeListArray::from_iter_primitive::<Int32Type, _,
_>([Some([Some(5)])], 1);
+ assert_eq!(&expect.value(0), &actual.value(0));
+ }
+
#[test]
fn test_cast_list_containers() {
// large-list to list
