vortex_array/arrays/chunked/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::DType;
5use vortex_error::VortexResult;
6
7use crate::arrays::{ChunkedArray, ChunkedVTable};
8use crate::compute::{CastKernel, CastKernelAdapter, cast};
9use crate::{ArrayRef, IntoArray, register_kernel};
10
11impl CastKernel for ChunkedVTable {
12    fn cast(&self, array: &ChunkedArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
13        let mut cast_chunks = Vec::new();
14        for chunk in array.chunks() {
15            cast_chunks.push(cast(chunk, dtype)?);
16        }
17
18        Ok(Some(
19            ChunkedArray::new_unchecked(cast_chunks, dtype.clone()).into_array(),
20        ))
21    }
22}
23
24register_kernel!(CastKernelAdapter(ChunkedVTable).lift());
25
26#[cfg(test)]
27mod test {
28    use rstest::rstest;
29    use vortex_buffer::buffer;
30    use vortex_dtype::{DType, Nullability, PType};
31
32    use crate::IntoArray;
33    use crate::arrays::PrimitiveArray;
34    use crate::arrays::chunked::ChunkedArray;
35    use crate::canonical::ToCanonical;
36    use crate::compute::cast;
37    use crate::compute::conformance::cast::test_cast_conformance;
38
39    #[test]
40    fn test_cast_chunked() {
41        let arr0 = buffer![0u32, 1].into_array();
42        let arr1 = buffer![2u32, 3].into_array();
43
44        let chunked = ChunkedArray::try_new(
45            vec![arr0, arr1],
46            DType::Primitive(PType::U32, Nullability::NonNullable),
47        )
48        .unwrap()
49        .into_array();
50
51        // Two levels of chunking, just to be fancy.
52        let root = ChunkedArray::try_new(
53            vec![chunked],
54            DType::Primitive(PType::U32, Nullability::NonNullable),
55        )
56        .unwrap()
57        .into_array();
58
59        assert_eq!(
60            cast(
61                &root,
62                &DType::Primitive(PType::U64, Nullability::NonNullable)
63            )
64            .unwrap()
65            .to_primitive()
66            .unwrap()
67            .as_slice::<u64>(),
68            &[0u64, 1, 2, 3],
69        );
70    }
71
72    #[rstest]
73    #[case(ChunkedArray::try_new(
74        vec![buffer![0u32, 1, 2].into_array(), buffer![3u32, 4].into_array()],
75        DType::Primitive(PType::U32, Nullability::NonNullable)
76    ).unwrap().into_array())]
77    #[case(ChunkedArray::try_new(
78        vec![
79            buffer![-10i32, -5, 0].into_array(),
80            buffer![5i32, 10].into_array()
81        ],
82        DType::Primitive(PType::I32, Nullability::NonNullable)
83    ).unwrap().into_array())]
84    #[case(ChunkedArray::try_new(
85        vec![
86            PrimitiveArray::from_option_iter([Some(1.5f32), None, Some(2.5)]).into_array(),
87            PrimitiveArray::from_option_iter([Some(3.5f32), Some(4.5)]).into_array()
88        ],
89        DType::Primitive(PType::F32, Nullability::Nullable)
90    ).unwrap().into_array())]
91    #[case(ChunkedArray::try_new(
92        vec![buffer![42u8].into_array()],
93        DType::Primitive(PType::U8, Nullability::NonNullable)
94    ).unwrap().into_array())]
95    fn test_cast_chunked_conformance(#[case] array: crate::ArrayRef) {
96        test_cast_conformance(array.as_ref());
97    }
98}
vortex_array/arrays/chunked/compute/cast.rs

vortex_array/arrays/chunked/compute/
cast.rs
