vortex_array/arrays/list/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::ExecutionCtx;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::List;
11use crate::arrays::ListArray;
12use crate::arrays::list::ListArraySlotsExt;
13use crate::builtins::ArrayBuiltins;
14use crate::dtype::DType;
15use crate::scalar_fn::fns::cast::CastKernel;
16use crate::scalar_fn::fns::cast::CastReduce;
17
18impl CastReduce for List {
19    fn cast(array: ArrayView<'_, List>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
20        let Some(target_element_type) = dtype.as_list_element_opt() else {
21            return Ok(None);
22        };
23
24        let Some(validity) = array
25            .validity()?
26            .trivially_cast_nullability(dtype.nullability(), array.len())?
27        else {
28            return Ok(None);
29        };
30
31        let new_elements = array.elements().cast((**target_element_type).clone())?;
32
33        Ok(Some(
34            unsafe { ListArray::new_unchecked(new_elements, array.offsets().clone(), validity) }
35                .into_array(),
36        ))
37    }
38}
39
40impl CastKernel for List {
41    fn cast(
42        array: ArrayView<'_, List>,
43        dtype: &DType,
44        ctx: &mut ExecutionCtx,
45    ) -> VortexResult<Option<ArrayRef>> {
46        let Some(target_element_type) = dtype.as_list_element_opt() else {
47            return Ok(None);
48        };
49
50        let validity = array
51            .validity()?
52            .cast_nullability(dtype.nullability(), array.len(), ctx)?;
53
54        let new_elements = array.elements().cast((**target_element_type).clone())?;
55
56        Ok(Some(
57            unsafe { ListArray::new_unchecked(new_elements, array.offsets().clone(), validity) }
58                .into_array(),
59        ))
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use std::sync::Arc;
66    use std::sync::LazyLock;
67
68    use rstest::rstest;
69    use vortex_buffer::buffer;
70    use vortex_session::VortexSession;
71
72    use crate::Canonical;
73    use crate::IntoArray;
74    use crate::RecursiveCanonical;
75    use crate::VortexSessionExecute;
76    use crate::arrays::BoolArray;
77    use crate::arrays::ListArray;
78    use crate::arrays::PrimitiveArray;
79    use crate::arrays::VarBinArray;
80    use crate::assert_arrays_eq;
81    use crate::builtins::ArrayBuiltins;
82    use crate::compute::conformance::cast::test_cast_conformance;
83    use crate::dtype::DType;
84    use crate::dtype::Nullability;
85    use crate::dtype::PType;
86    use crate::validity::Validity;
87
88    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
89
90    #[test]
91    fn test_cast_list_success() {
92        let list = ListArray::try_new(
93            buffer![1i32, 2, 3, 4].into_array(),
94            buffer![0, 2, 3].into_array(),
95            Validity::NonNullable,
96        )
97        .unwrap();
98
99        let target_dtype = DType::List(
100            Arc::new(DType::Primitive(PType::U64, Nullability::Nullable)),
101            Nullability::Nullable,
102        );
103
104        let result = list
105            .clone()
106            .into_array()
107            .cast(target_dtype.clone())
108            .unwrap();
109        assert_eq!(result.dtype(), &target_dtype);
110        assert_eq!(result.len(), list.len());
111        let expected = ListArray::new(
112            PrimitiveArray::from_option_iter([Some(1u64), Some(2), Some(3)]).into_array(),
113            buffer![0u32, 2, 3].into_array(),
114            Validity::AllValid,
115        );
116        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
117    }
118
119    #[test]
120    fn test_cast_to_wrong_type() {
121        let list = ListArray::try_new(
122            buffer![0i32, 2, 3, 4].into_array(),
123            buffer![0, 2, 3].into_array(),
124            Validity::NonNullable,
125        )
126        .unwrap();
127
128        let target_dtype = DType::Primitive(PType::U64, Nullability::NonNullable);
129        // can't cast list to u64
130
131        let result = list
132            .into_array()
133            .cast(target_dtype)
134            .and_then(|a| a.execute::<Canonical>(&mut SESSION.create_execution_ctx()))
135            .map(|c| c.into_array());
136        assert!(result.is_err());
137    }
138
139    #[test]
140    fn test_cant_cast_nulls_to_non_null() {
141        // Test that if list has nulls, the conversion will fail
142
143        // Nulls in the list itself
144        let list = ListArray::try_new(
145            buffer![0i32, 2, 3, 4].into_array(),
146            buffer![0, 2, 3].into_array(),
147            Validity::Array(BoolArray::from_iter(vec![false, true]).into_array()),
148        )
149        .unwrap();
150
151        let target_dtype = DType::List(
152            Arc::new(DType::Primitive(PType::U64, Nullability::Nullable)),
153            Nullability::NonNullable,
154        );
155
156        let result = list
157            .into_array()
158            .cast(target_dtype)
159            .and_then(|a| a.execute::<Canonical>(&mut SESSION.create_execution_ctx()))
160            .map(|c| c.into_array());
161        assert!(result.is_err());
162
163        // Nulls in list element array — the inner cast error is deferred until
164        // the elements are executed.
165        let list = ListArray::try_new(
166            PrimitiveArray::from_option_iter([Some(0i32), Some(2), None, None]).into_array(),
167            buffer![0, 2, 3].into_array(),
168            Validity::NonNullable,
169        )
170        .unwrap();
171
172        let target_dtype = DType::List(
173            Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)),
174            Nullability::NonNullable,
175        );
176
177        let result = list.into_array().cast(target_dtype).and_then(|a| {
178            a.execute::<RecursiveCanonical>(&mut SESSION.create_execution_ctx())
179                .map(|c| c.0.into_array())
180        });
181        assert!(result.is_err());
182    }
183
184    #[rstest]
185    #[case(create_simple_list())]
186    #[case(create_nullable_list())]
187    #[case(create_string_list())]
188    #[case(create_nested_list())]
189    #[case(create_empty_lists())]
190    fn test_cast_list_conformance(#[case] array: ListArray) {
191        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
192    }
193
194    fn create_simple_list() -> ListArray {
195        let data = buffer![1i32, 2, 3, 4, 5, 6].into_array();
196        let offsets = buffer![0i64, 2, 2, 5, 6].into_array();
197
198        ListArray::try_new(data, offsets, Validity::NonNullable).unwrap()
199    }
200
201    fn create_nullable_list() -> ListArray {
202        let data = PrimitiveArray::from_option_iter([
203            Some(10i64),
204            None,
205            Some(20),
206            Some(30),
207            None,
208            Some(40),
209        ])
210        .into_array();
211        let offsets = buffer![0i64, 3, 6].into_array();
212        let validity = Validity::Array(BoolArray::from_iter(vec![true, false]).into_array());
213
214        ListArray::try_new(data, offsets, validity).unwrap()
215    }
216
217    fn create_string_list() -> ListArray {
218        let data = VarBinArray::from_iter(
219            vec![Some("hello"), Some("world"), Some("foo"), Some("bar")],
220            DType::Utf8(Nullability::NonNullable),
221        )
222        .into_array();
223        let offsets = buffer![0i64, 2, 4].into_array();
224
225        ListArray::try_new(data, offsets, Validity::NonNullable).unwrap()
226    }
227
228    fn create_nested_list() -> ListArray {
229        // Create inner lists: [[1, 2], [3], [4, 5, 6]]
230        let inner_data = buffer![1i32, 2, 3, 4, 5, 6].into_array();
231        let inner_offsets = buffer![0i64, 2, 3, 6].into_array();
232        let inner_list = ListArray::try_new(inner_data, inner_offsets, Validity::NonNullable)
233            .unwrap()
234            .into_array();
235
236        // Create outer list: [[[1, 2], [3]], [[4, 5, 6]]]
237        let outer_offsets = buffer![0i64, 2, 3].into_array();
238
239        ListArray::try_new(inner_list, outer_offsets, Validity::NonNullable).unwrap()
240    }
241
242    fn create_empty_lists() -> ListArray {
243        let data = buffer![42u8].into_array();
244        let offsets = buffer![0i64, 0, 0, 1].into_array();
245
246        ListArray::try_new(data, offsets, Validity::NonNullable).unwrap()
247    }
248}
vortex_array/arrays/list/compute/cast.rs

vortex_array/arrays/list/compute/
cast.rs
