From 3428ad528bec6b2240e0f68113d544eea948621c Mon Sep 17 00:00:00 2001 From: Bo Lin Date: Thu, 16 May 2024 22:46:58 -0400 Subject: [PATCH] Support casting a `FixedSizedList[1]` to `T` --- .../src/array/fixed_size_list_array.rs | 3 +- arrow-cast/src/cast/list.rs | 9 ++++ arrow-cast/src/cast/mod.rs | 44 +++++++++++++++++++ 3 files changed, 55 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 1d07daf4cc75..7bb4e0c5ee5a 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 2ceab1d0c8c2..ec7a5c57d504 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 { + let values = array.as_fixed_size_list().values(); + cast_with_options(values, to, cast_options) +} + pub(crate) fn cast_fixed_size_list_to_list( array: &dyn Array, ) -> Result diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 8b7579c4cfc0..ba29e85f0b3e 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -153,6 +153,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, @@ -703,6 +705,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::( array.as_primitive(), @@ -6660,6 +6665,45 @@ mod tests { assert_eq!(&expect.value(0), &actual.value(0)); } + #[test] + fn test_cast_single_element_fixed_size_list() { + // FixedSizeList[1] => T + let from_array = Arc::new(FixedSizeListArray::from_iter_primitive::( + [(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); + + // T => FixedSizeList[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[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