vortex_zstd/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::compute::{CastKernel, CastKernelAdapter};
5use vortex_array::{ArrayRef, IntoArray, register_kernel};
6use vortex_dtype::DType;
7use vortex_error::VortexResult;
8
9use crate::{ZstdArray, ZstdVTable};
10
11impl CastKernel for ZstdVTable {
12    fn cast(&self, array: &ZstdArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
13        // ZstdArray is a general-purpose compression encoding using Zstandard compression.
14        // It can handle nullability changes without decompression by updating the validity
15        // bitmap, but type changes require decompression since the compressed data is
16        // type-specific and Zstd operates on raw bytes.
17        if array.dtype().eq_ignore_nullability(dtype) {
18            // Create a new validity with the target nullability
19            let new_validity = array
20                .unsliced_validity
21                .clone()
22                .cast_nullability(dtype.nullability())?;
23
24            return Ok(Some(
25                ZstdArray::new(
26                    array.dictionary.clone(),
27                    array.frames.clone(),
28                    dtype.clone(),
29                    array.metadata.clone(),
30                    array.unsliced_n_rows(),
31                    new_validity,
32                )
33                ._slice(array.slice_start(), array.slice_stop())
34                .into_array(),
35            ));
36        }
37
38        // For other casts (e.g., type changes), decode to canonical and let the underlying array handle it
39        Ok(None)
40    }
41}
42
43register_kernel!(CastKernelAdapter(ZstdVTable).lift());
44
45#[cfg(test)]
46mod tests {
47    use rstest::rstest;
48    use vortex_array::arrays::PrimitiveArray;
49    use vortex_array::compute::cast;
50    use vortex_array::compute::conformance::cast::test_cast_conformance;
51    use vortex_buffer::Buffer;
52    use vortex_dtype::{DType, Nullability, PType};
53
54    use crate::ZstdArray;
55
56    #[test]
57    fn test_cast_zstd_i32_to_i64() {
58        let values = PrimitiveArray::new(
59            Buffer::copy_from(vec![1i32, 2, 3, 4, 5]),
60            vortex_array::validity::Validity::NonNullable,
61        );
62        let zstd = ZstdArray::from_primitive(&values, 0, 0).unwrap();
63
64        let casted = cast(
65            zstd.as_ref(),
66            &DType::Primitive(PType::I64, Nullability::NonNullable),
67        )
68        .unwrap();
69        assert_eq!(
70            casted.dtype(),
71            &DType::Primitive(PType::I64, Nullability::NonNullable)
72        );
73
74        let decoded = casted.to_canonical().unwrap().into_primitive().unwrap();
75        assert_eq!(decoded.as_slice::<i64>(), &[1i64, 2, 3, 4, 5]);
76    }
77
78    #[test]
79    fn test_cast_zstd_nullability_change() {
80        let values = PrimitiveArray::new(
81            Buffer::copy_from(vec![10u32, 20, 30, 40]),
82            vortex_array::validity::Validity::NonNullable,
83        );
84        let zstd = ZstdArray::from_primitive(&values, 0, 0).unwrap();
85
86        let casted = cast(
87            zstd.as_ref(),
88            &DType::Primitive(PType::U32, Nullability::Nullable),
89        )
90        .unwrap();
91        assert_eq!(
92            casted.dtype(),
93            &DType::Primitive(PType::U32, Nullability::Nullable)
94        );
95    }
96
97    #[rstest]
98    #[case::i32(PrimitiveArray::new(
99        Buffer::copy_from(vec![100i32, 200, 300, 400, 500]),
100        vortex_array::validity::Validity::NonNullable,
101    ))]
102    #[case::f64(PrimitiveArray::new(
103        Buffer::copy_from(vec![1.1f64, 2.2, 3.3, 4.4, 5.5]),
104        vortex_array::validity::Validity::NonNullable,
105    ))]
106    #[case::single(PrimitiveArray::new(
107        Buffer::copy_from(vec![42i64]),
108        vortex_array::validity::Validity::NonNullable,
109    ))]
110    #[case::large(PrimitiveArray::new(
111        Buffer::copy_from((0..1000).map(|i| i as u32).collect::<Vec<_>>()),
112        vortex_array::validity::Validity::NonNullable,
113    ))]
114    fn test_cast_zstd_conformance(#[case] values: PrimitiveArray) {
115        let zstd = ZstdArray::from_primitive(&values, 0, 0).unwrap();
116        test_cast_conformance(zstd.as_ref());
117    }
118}
vortex_zstd/compute/cast.rs

vortex_zstd/compute/
cast.rs
