Skip to content

Commit

Permalink
Support casting a FixedSizedList<T>[1] to T (#5779)
Browse files Browse the repository at this point in the history
* Support casting a `FixedSizedList<T>[1]` to `T`

* Add FixedSizedList[1] => FixeSizedList[1] tests
  • Loading branch information
sadboy committed May 18, 2024
1 parent 30762e8 commit a126d50
Show file tree
Hide file tree
Showing 3 changed files with 95 additions and 1 deletion.
3 changes: 2 additions & 1 deletion arrow-array/src/array/fixed_size_list_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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!(
Expand Down
9 changes: 9 additions & 0 deletions arrow-cast/src/cast/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
84 changes: 84 additions & 0 deletions arrow-cast/src/cast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit a126d50

Please sign in to comment.