Skip to main content

vortex_array/arrays/listview/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::ListView;
11use crate::arrays::ListViewArray;
12use crate::arrays::listview::ListViewArrayExt;
13use crate::builtins::ArrayBuiltins;
14use crate::dtype::DType;
15use crate::scalar_fn::fns::cast::CastKernel;
16use crate::scalar_fn::fns::cast::CastReduce;
17use crate::validity::Validity;
18
19fn build_with_validity(
20    array: ArrayView<'_, ListView>,
21    new_elements: ArrayRef,
22    validity: Validity,
23) -> ArrayRef {
24    // SAFETY: Since `cast` is length-preserving, all of the invariants remain the same.
25    unsafe {
26        ListViewArray::new_unchecked(
27            new_elements,
28            array.offsets().clone(),
29            array.sizes().clone(),
30            validity,
31        )
32        .with_zero_copy_to_list(array.is_zero_copy_to_list())
33    }
34    .into_array()
35}
36
37impl CastReduce for ListView {
38    fn cast(array: ArrayView<'_, ListView>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
39        // Check if we're casting to a `List` type.
40        let Some(target_element_type) = dtype.as_list_element_opt() else {
41            return Ok(None);
42        };
43        let Some(validity) = array
44            .validity()?
45            .trivial_cast_nullability(dtype.nullability(), array.len())?
46        else {
47            return Ok(None);
48        };
49
50        // Cast the elements to the target element type.
51        let new_elements = array.elements().cast((**target_element_type).clone())?;
52        Ok(Some(build_with_validity(array, new_elements, validity)))
53    }
54}
55
56impl CastKernel for ListView {
57    fn cast(
58        array: ArrayView<'_, ListView>,
59        dtype: &DType,
60        ctx: &mut ExecutionCtx,
61    ) -> VortexResult<Option<ArrayRef>> {
62        let Some(target_element_type) = dtype.as_list_element_opt() else {
63            return Ok(None);
64        };
65
66        let validity = array
67            .validity()?
68            .cast_nullability(dtype.nullability(), array.len(), ctx)?;
69        let new_elements = array.elements().cast((**target_element_type).clone())?;
70
71        Ok(Some(build_with_validity(array, new_elements, validity)))
72    }
73}