diff --git a/src/compute/cast/mod.rs b/src/compute/cast/mod.rs index 14622f9b03..a9d27aa12d 100644 --- a/src/compute/cast/mod.rs +++ b/src/compute/cast/mod.rs @@ -104,6 +104,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool { (List(list_from), LargeList(list_to)) if list_from == list_to => true, (LargeList(list_from), List(list_to)) if list_from == list_to => true, (_, List(list_to)) => can_cast_types(from_type, &list_to.data_type), + (_, LargeList(list_to)) => can_cast_types(from_type, &list_to.data_type), (Dictionary(_, from_value_type, _), Dictionary(_, to_value_type, _)) => { can_cast_types(from_value_type, to_value_type) } @@ -509,6 +510,19 @@ pub fn cast(array: &dyn Array, to_type: &DataType, options: CastOptions) -> Resu Ok(Box::new(list_array)) } + (_, LargeList(to)) => { + // cast primitive to list's primitive + let values = cast(array, &to.data_type, options)?; + // create offsets, where if array.len() = 2, we have [0,1,2] + let offsets = (0..=array.len() as i64).collect::>(); + // Safety: offsets _are_ monotonically increasing + let offsets = unsafe { Offsets::new_unchecked(offsets) }; + + let list_array = ListArray::::new(to_type.clone(), offsets.into(), values, None); + + Ok(Box::new(list_array)) + } + (Dictionary(index_type, ..), _) => match_integer_type!(index_type, |$T| { dictionary_cast_dyn::<$T>(array, to_type, options) }), diff --git a/tests/it/compute/cast.rs b/tests/it/compute/cast.rs index 01cb31d2f2..d8a9ecfce1 100644 --- a/tests/it/compute/cast.rs +++ b/tests/it/compute/cast.rs @@ -1,5 +1,6 @@ use arrow2::array::*; use arrow2::compute::cast::{can_cast_types, cast, CastOptions}; +use arrow2::datatypes::DataType::LargeList; use arrow2::datatypes::*; use arrow2::types::{days_ms, months_days_ns, NativeType}; @@ -120,6 +121,28 @@ fn i32_to_i32() { assert_eq!(c, &expected); } +#[test] +fn i32_to_large_list_i32() { + let array = Int32Array::from_slice([5, 6, 7, 8, 9]); + let b = cast( + &array, + &LargeList(Box::new(Field::new("item", DataType::Int32, true))), + CastOptions::default(), + ) + .unwrap(); + + let arr = b.as_any().downcast_ref::>().unwrap(); + assert_eq!(&[0, 1, 2, 3, 4, 5], arr.offsets().as_slice()); + let values = arr.values(); + let c = values + .as_any() + .downcast_ref::>() + .unwrap(); + + let expected = Int32Array::from_slice([5, 6, 7, 8, 9]); + assert_eq!(c, &expected); +} + #[test] fn i32_to_list_i32() { let array = Int32Array::from_slice([5, 6, 7, 8, 9]);