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        // SAFETY: casting all chunks retains all chunks have same DType
19        unsafe {
20            Ok(Some(
21                ChunkedArray::new_unchecked(cast_chunks, dtype.clone()).into_array(),
22            ))
23        }
24    }
25}
26
27register_kernel!(CastKernelAdapter(ChunkedVTable).lift());
28
29#[cfg(test)]
30mod test {
31    use rstest::rstest;
32    use vortex_buffer::buffer;
33    use vortex_dtype::{DType, Nullability, PType};
34
35    use crate::IntoArray;
36    use crate::arrays::PrimitiveArray;
37    use crate::arrays::chunked::ChunkedArray;
38    use crate::canonical::ToCanonical;
39    use crate::compute::cast;
40    use crate::compute::conformance::cast::test_cast_conformance;
41
42    #[test]
43    fn test_cast_chunked() {
44        let arr0 = buffer![0u32, 1].into_array();
45        let arr1 = buffer![2u32, 3].into_array();
46
47        let chunked = ChunkedArray::try_new(
48            vec![arr0, arr1],
49            DType::Primitive(PType::U32, Nullability::NonNullable),
50        )
51        .unwrap()
52        .into_array();
53
54        // Two levels of chunking, just to be fancy.
55        let root = ChunkedArray::try_new(
56            vec![chunked],
57            DType::Primitive(PType::U32, Nullability::NonNullable),
58        )
59        .unwrap()
60        .into_array();
61
62        assert_eq!(
63            cast(
64                &root,
65                &DType::Primitive(PType::U64, Nullability::NonNullable)
66            )
67            .unwrap()
68            .to_primitive()
69            .unwrap()
70            .as_slice::<u64>(),
71            &[0u64, 1, 2, 3],
72        );
73    }
74
75    #[rstest]
76    #[case(ChunkedArray::try_new(
77        vec![buffer![0u32, 1, 2].into_array(), buffer![3u32, 4].into_array()],
78        DType::Primitive(PType::U32, Nullability::NonNullable)
79    ).unwrap().into_array())]
80    #[case(ChunkedArray::try_new(
81        vec![
82            buffer![-10i32, -5, 0].into_array(),
83            buffer![5i32, 10].into_array()
84        ],
85        DType::Primitive(PType::I32, Nullability::NonNullable)
86    ).unwrap().into_array())]
87    #[case(ChunkedArray::try_new(
88        vec![
89            PrimitiveArray::from_option_iter([Some(1.5f32), None, Some(2.5)]).into_array(),
90            PrimitiveArray::from_option_iter([Some(3.5f32), Some(4.5)]).into_array()
91        ],
92        DType::Primitive(PType::F32, Nullability::Nullable)
93    ).unwrap().into_array())]
94    #[case(ChunkedArray::try_new(
95        vec![buffer![42u8].into_array()],
96        DType::Primitive(PType::U8, Nullability::NonNullable)
97    ).unwrap().into_array())]
98    fn test_cast_chunked_conformance(#[case] array: crate::ArrayRef) {
99        test_cast_conformance(array.as_ref());
100    }
101}
vortex_array/arrays/chunked/compute/cast.rs

vortex_array/arrays/chunked/compute/
cast.rs
