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            .as_slice::<u64>(),
70            &[0u64, 1, 2, 3],
71        );
72    }
73
74    #[rstest]
75    #[case(ChunkedArray::try_new(
76        vec![buffer![0u32, 1, 2].into_array(), buffer![3u32, 4].into_array()],
77        DType::Primitive(PType::U32, Nullability::NonNullable)
78    ).unwrap().into_array())]
79    #[case(ChunkedArray::try_new(
80        vec![
81            buffer![-10i32, -5, 0].into_array(),
82            buffer![5i32, 10].into_array()
83        ],
84        DType::Primitive(PType::I32, Nullability::NonNullable)
85    ).unwrap().into_array())]
86    #[case(ChunkedArray::try_new(
87        vec![
88            PrimitiveArray::from_option_iter([Some(1.5f32), None, Some(2.5)]).into_array(),
89            PrimitiveArray::from_option_iter([Some(3.5f32), Some(4.5)]).into_array()
90        ],
91        DType::Primitive(PType::F32, Nullability::Nullable)
92    ).unwrap().into_array())]
93    #[case(ChunkedArray::try_new(
94        vec![buffer![42u8].into_array()],
95        DType::Primitive(PType::U8, Nullability::NonNullable)
96    ).unwrap().into_array())]
97    fn test_cast_chunked_conformance(#[case] array: crate::ArrayRef) {
98        test_cast_conformance(array.as_ref());
99    }
100}
vortex_array/arrays/chunked/compute/cast.rs

vortex_array/arrays/chunked/compute/
cast.rs
