Skip to content

Commit

Permalink
Minor: support cast values to fixedsizelist (#5340)
Browse files Browse the repository at this point in the history
* 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]>
  • Loading branch information
3 people authored Feb 4, 2024
1 parent f303c9e commit 79721ec
Showing 1 changed file with 41 additions and 0 deletions.
41 changes: 41 additions & 0 deletions arrow-cast/src/cast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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>(
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 79721ec

Please sign in to comment.