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 vortex_buffer::buffer;
29    use vortex_dtype::{DType, Nullability, PType};
30
31    use crate::IntoArray;
32    use crate::arrays::chunked::ChunkedArray;
33    use crate::canonical::ToCanonical;
34    use crate::compute::cast;
35
36    #[test]
37    fn test_cast_chunked() {
38        let arr0 = buffer![0u32, 1].into_array();
39        let arr1 = buffer![2u32, 3].into_array();
40
41        let chunked = ChunkedArray::try_new(
42            vec![arr0, arr1],
43            DType::Primitive(PType::U32, Nullability::NonNullable),
44        )
45        .unwrap()
46        .into_array();
47
48        // Two levels of chunking, just to be fancy.
49        let root = ChunkedArray::try_new(
50            vec![chunked],
51            DType::Primitive(PType::U32, Nullability::NonNullable),
52        )
53        .unwrap()
54        .into_array();
55
56        assert_eq!(
57            cast(
58                &root,
59                &DType::Primitive(PType::U64, Nullability::NonNullable)
60            )
61            .unwrap()
62            .to_primitive()
63            .unwrap()
64            .as_slice::<u64>(),
65            &[0u64, 1, 2, 3],
66        );
67    }
68}
vortex_array/arrays/chunked/compute/cast.rs

vortex_array/arrays/chunked/compute/
cast.rs
