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 mut ctx = SESSION.create_execution_ctx();
93        let list = ListArray::try_new(
94            buffer![1i32, 2, 3, 4].into_array(),
95            buffer![0, 2, 3].into_array(),
96            Validity::NonNullable,
97            &mut ctx,
98        )
99        .unwrap();
100
101        let target_dtype = DType::List(
102            Arc::new(DType::Primitive(PType::U64, Nullability::Nullable)),
103            Nullability::Nullable,
104        );
105
106        let result = list
107            .clone()
108            .into_array()
109            .cast(target_dtype.clone())
110            .unwrap();
111        assert_eq!(result.dtype(), &target_dtype);
112        assert_eq!(result.len(), list.len());
113        let expected = ListArray::new(
114            PrimitiveArray::from_option_iter([Some(1u64), Some(2), Some(3)]).into_array(),
115            buffer![0u32, 2, 3].into_array(),
116            Validity::AllValid,
117            &mut ctx,
118        );
119        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
120    }
121
122    #[test]
123    fn test_cast_to_wrong_type() {
124        let mut ctx = SESSION.create_execution_ctx();
125        let list = ListArray::try_new(
126            buffer![0i32, 2, 3, 4].into_array(),
127            buffer![0, 2, 3].into_array(),
128            Validity::NonNullable,
129            &mut ctx,
130        )
131        .unwrap();
132
133        let target_dtype = DType::Primitive(PType::U64, Nullability::NonNullable);
134        // can't cast list to u64
135
136        let result = list
137            .into_array()
138            .cast(target_dtype)
139            .and_then(|a| a.execute::<Canonical>(&mut SESSION.create_execution_ctx()))
140            .map(|c| c.into_array());
141        assert!(result.is_err());
142    }
143
144    #[test]
145    fn test_cant_cast_nulls_to_non_null() {
146        let mut ctx = SESSION.create_execution_ctx();
147        // Test that if list has nulls, the conversion will fail
148
149        // Nulls in the list itself
150        let list = ListArray::try_new(
151            buffer![0i32, 2, 3, 4].into_array(),
152            buffer![0, 2, 3].into_array(),
153            Validity::Array(BoolArray::from_iter(vec![false, true]).into_array()),
154            &mut ctx,
155        )
156        .unwrap();
157
158        let target_dtype = DType::List(
159            Arc::new(DType::Primitive(PType::U64, Nullability::Nullable)),
160            Nullability::NonNullable,
161        );
162
163        let result = list
164            .into_array()
165            .cast(target_dtype)
166            .and_then(|a| a.execute::<Canonical>(&mut SESSION.create_execution_ctx()))
167            .map(|c| c.into_array());
168        assert!(result.is_err());
169
170        // Nulls in list element array — the inner cast error is deferred until
171        // the elements are executed.
172        let list = ListArray::try_new(
173            PrimitiveArray::from_option_iter([Some(0i32), Some(2), None, None]).into_array(),
174            buffer![0, 2, 3].into_array(),
175            Validity::NonNullable,
176            &mut ctx,
177        )
178        .unwrap();
179
180        let target_dtype = DType::List(
181            Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)),
182            Nullability::NonNullable,
183        );
184
185        let result = list.into_array().cast(target_dtype).and_then(|a| {
186            a.execute::<RecursiveCanonical>(&mut SESSION.create_execution_ctx())
187                .map(|c| c.0.into_array())
188        });
189        assert!(result.is_err());
190    }
191
192    #[rstest]
193    #[case(create_simple_list())]
194    #[case(create_nullable_list())]
195    #[case(create_string_list())]
196    #[case(create_nested_list())]
197    #[case(create_empty_lists())]
198    fn test_cast_list_conformance(#[case] array: ListArray) {
199        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
200    }
201
202    fn create_simple_list() -> ListArray {
203        let mut ctx = SESSION.create_execution_ctx();
204        let data = buffer![1i32, 2, 3, 4, 5, 6].into_array();
205        let offsets = buffer![0i64, 2, 2, 5, 6].into_array();
206
207        ListArray::try_new(data, offsets, Validity::NonNullable, &mut ctx).unwrap()
208    }
209
210    fn create_nullable_list() -> ListArray {
211        let mut ctx = SESSION.create_execution_ctx();
212        let data = PrimitiveArray::from_option_iter([
213            Some(10i64),
214            None,
215            Some(20),
216            Some(30),
217            None,
218            Some(40),
219        ])
220        .into_array();
221        let offsets = buffer![0i64, 3, 6].into_array();
222        let validity = Validity::Array(BoolArray::from_iter(vec![true, false]).into_array());
223
224        ListArray::try_new(data, offsets, validity, &mut ctx).unwrap()
225    }
226
227    fn create_string_list() -> ListArray {
228        let mut ctx = SESSION.create_execution_ctx();
229        let data = VarBinArray::from_iter(
230            vec![Some("hello"), Some("world"), Some("foo"), Some("bar")],
231            DType::Utf8(Nullability::NonNullable),
232        )
233        .into_array();
234        let offsets = buffer![0i64, 2, 4].into_array();
235
236        ListArray::try_new(data, offsets, Validity::NonNullable, &mut ctx).unwrap()
237    }
238
239    fn create_nested_list() -> ListArray {
240        let mut ctx = SESSION.create_execution_ctx();
241        // Create inner lists: [[1, 2], [3], [4, 5, 6]]
242        let inner_data = buffer![1i32, 2, 3, 4, 5, 6].into_array();
243        let inner_offsets = buffer![0i64, 2, 3, 6].into_array();
244        let inner_list =
245            ListArray::try_new(inner_data, inner_offsets, Validity::NonNullable, &mut ctx)
246                .unwrap()
247                .into_array();
248
249        // Create outer list: [[[1, 2], [3]], [[4, 5, 6]]]
250        let outer_offsets = buffer![0i64, 2, 3].into_array();
251
252        ListArray::try_new(inner_list, outer_offsets, Validity::NonNullable, &mut ctx).unwrap()
253    }
254
255    fn create_empty_lists() -> ListArray {
256        let mut ctx = SESSION.create_execution_ctx();
257        let data = buffer![42u8].into_array();
258        let offsets = buffer![0i64, 0, 0, 1].into_array();
259
260        ListArray::try_new(data, offsets, Validity::NonNullable, &mut ctx).unwrap()
261    }
262}
vortex_array/arrays/list/compute/cast.rs

vortex_array/arrays/list/compute/
cast.rs
