You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2024/05/18 11:36:49 UTC
(arrow-rs) branch master updated: Support casting a `FixedSizedList[1]` to `T` (#5779)
This is an automated email from the ASF dual-hosted git repository.
tustvold 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 a126d5097b7 Support casting a `FixedSizedList<T>[1]` to `T` (#5779)
a126d5097b7 is described below
commit a126d5097b71273428fba68d1c430f3d4beee684
Author: Bo Lin <bo...@dreamsphere.org>
AuthorDate: Sat May 18 07:36:43 2024 -0400
Support casting a `FixedSizedList<T>[1]` to `T` (#5779)
* Support casting a `FixedSizedList<T>[1]` to `T`
* Add FixedSizedList[1] => FixeSizedList[1] tests
---
arrow-array/src/array/fixed_size_list_array.rs | 3 +-
arrow-cast/src/cast/list.rs | 9 +++
arrow-cast/src/cast/mod.rs | 84 ++++++++++++++++++++++++++
3 files changed, 95 insertions(+), 1 deletion(-)
diff --git a/arrow-array/src/array/fixed_size_list_array.rs b/arrow-array/src/array/fixed_size_list_array.rs
index 1d07daf4cc7..7bb4e0c5ee5 100644
--- a/arrow-array/src/array/fixed_size_list_array.rs
+++ b/arrow-array/src/array/fixed_size_list_array.rs
@@ -183,7 +183,8 @@ impl FixedSizeListArray {
|| nulls
.as_ref()
.map(|n| n.expand(size as _).contains(&a))
- .unwrap_or_default();
+ .unwrap_or_default()
+ || (nulls.is_none() && a.null_count() == 0);
if !nulls_valid {
return Err(ArrowError::InvalidArgumentError(format!(
diff --git a/arrow-cast/src/cast/list.rs b/arrow-cast/src/cast/list.rs
index 2ceab1d0c8c..ec7a5c57d50 100644
--- a/arrow-cast/src/cast/list.rs
+++ b/arrow-cast/src/cast/list.rs
@@ -41,6 +41,15 @@ pub(crate) fn cast_values_to_fixed_size_list(
Ok(Arc::new(list))
}
+pub(crate) fn cast_single_element_fixed_size_list_to_values(
+ array: &dyn Array,
+ to: &DataType,
+ cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError> {
+ let values = array.as_fixed_size_list().values();
+ cast_with_options(values, to, cast_options)
+}
+
pub(crate) fn cast_fixed_size_list_to_list<OffsetSize>(
array: &dyn Array,
) -> Result<ArrayRef, ArrowError>
diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs
index ccbc8727f7d..f2fdfd04497 100644
--- a/arrow-cast/src/cast/mod.rs
+++ b/arrow-cast/src/cast/mod.rs
@@ -157,6 +157,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(_, 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())},
+ (FixedSizeList(list_from,size), _) if *size == 1 => {
+ can_cast_types(list_from.data_type(), to_type)},
// cast one decimal type to another decimal type
(Decimal128(_, _), Decimal128(_, _)) => true,
(Decimal256(_, _), Decimal256(_, _)) => true,
@@ -712,6 +714,9 @@ pub fn cast_with_options(
(_, FixedSizeList(ref to, size)) if *size == 1 => {
cast_values_to_fixed_size_list(array, to, *size, cast_options)
}
+ (FixedSizeList(_, size), _) if *size == 1 => {
+ cast_single_element_fixed_size_list_to_values(array, to_type, cast_options)
+ }
(Decimal128(_, s1), Decimal128(p2, s2)) => {
cast_decimal_to_decimal_same_type::<Decimal128Type>(
array.as_primitive(),
@@ -6895,6 +6900,85 @@ mod tests {
assert_eq!(&expect.value(0), &actual.value(0));
}
+ #[test]
+ fn test_cast_single_element_fixed_size_list() {
+ // FixedSizeList<T>[1] => T
+ let from_array = Arc::new(FixedSizeListArray::from_iter_primitive::<Int16Type, _, _>(
+ [(Some([Some(5)]))],
+ 1,
+ )) as ArrayRef;
+ let casted_array = cast(&from_array, &DataType::Int32).unwrap();
+ let actual: &Int32Array = casted_array.as_primitive();
+ let expected = Int32Array::from(vec![Some(5)]);
+ assert_eq!(&expected, actual);
+
+ // FixedSizeList<T>[1] => FixedSizeList<U>[1]
+ let from_array = Arc::new(FixedSizeListArray::from_iter_primitive::<Int16Type, _, _>(
+ [(Some([Some(5)]))],
+ 1,
+ )) as ArrayRef;
+ let to_field = Arc::new(Field::new("dummy", DataType::Float32, false));
+ let actual = cast(&from_array, &DataType::FixedSizeList(to_field.clone(), 1)).unwrap();
+ let expected = Arc::new(FixedSizeListArray::new(
+ to_field.clone(),
+ 1,
+ Arc::new(Float32Array::from(vec![Some(5.0)])) as ArrayRef,
+ None,
+ )) as ArrayRef;
+ assert_eq!(*expected, *actual);
+
+ // FixedSizeList<T>[1] => FixedSizeList<FixdSizedList<U>[1]>[1]
+ let from_array = Arc::new(FixedSizeListArray::from_iter_primitive::<Int16Type, _, _>(
+ [(Some([Some(5)]))],
+ 1,
+ )) as ArrayRef;
+ let to_field_inner = Arc::new(Field::new("item", DataType::Float32, false));
+ let to_field = Arc::new(Field::new(
+ "dummy",
+ DataType::FixedSizeList(to_field_inner.clone(), 1),
+ false,
+ ));
+ let actual = cast(&from_array, &DataType::FixedSizeList(to_field.clone(), 1)).unwrap();
+ let expected = Arc::new(FixedSizeListArray::new(
+ to_field.clone(),
+ 1,
+ Arc::new(FixedSizeListArray::new(
+ to_field_inner.clone(),
+ 1,
+ Arc::new(Float32Array::from(vec![Some(5.0)])) as ArrayRef,
+ None,
+ )) as ArrayRef,
+ None,
+ )) as ArrayRef;
+ assert_eq!(*expected, *actual);
+
+ // T => FixedSizeList<T>[1] (non-nullable)
+ let field = Arc::new(Field::new("dummy", DataType::Float32, false));
+ let from_array = Arc::new(Int8Array::from(vec![Some(5)])) as ArrayRef;
+ let casted_array = cast(&from_array, &DataType::FixedSizeList(field.clone(), 1)).unwrap();
+ let actual = casted_array.as_fixed_size_list();
+ let expected = Arc::new(FixedSizeListArray::new(
+ field.clone(),
+ 1,
+ Arc::new(Float32Array::from(vec![Some(5.0)])) as ArrayRef,
+ None,
+ )) as ArrayRef;
+ assert_eq!(expected.as_ref(), actual);
+
+ // T => FixedSizeList<T>[1] (nullable)
+ let field = Arc::new(Field::new("nullable", DataType::Float32, true));
+ let from_array = Arc::new(Int8Array::from(vec![None])) as ArrayRef;
+ let casted_array = cast(&from_array, &DataType::FixedSizeList(field.clone(), 1)).unwrap();
+ let actual = casted_array.as_fixed_size_list();
+ let expected = Arc::new(FixedSizeListArray::new(
+ field.clone(),
+ 1,
+ Arc::new(Float32Array::from(vec![None])) as ArrayRef,
+ None,
+ )) as ArrayRef;
+ assert_eq!(expected.as_ref(), actual);
+ }
+
#[test]
fn test_cast_list_containers() {
// large-list to list