vortex_zigzag/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::IntoArray;
6use vortex_array::builtins::ArrayBuiltins;
7use vortex_array::dtype::DType;
8use vortex_array::scalar_fn::fns::cast::CastReduce;
9use vortex_error::VortexResult;
10
11use crate::ZigZagArray;
12use crate::ZigZagVTable;
13
14impl CastReduce for ZigZagVTable {
15    fn cast(array: &ZigZagArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
16        if !dtype.is_signed_int() {
17            return Ok(None);
18        }
19
20        let new_encoded_dtype =
21            DType::Primitive(dtype.as_ptype().to_unsigned(), dtype.nullability());
22        let new_encoded = array.encoded().cast(new_encoded_dtype)?;
23        Ok(Some(ZigZagArray::try_new(new_encoded)?.into_array()))
24    }
25}
26
27#[cfg(test)]
28mod tests {
29    use rstest::rstest;
30    use vortex_array::Array;
31    use vortex_array::arrays::PrimitiveArray;
32    use vortex_array::assert_arrays_eq;
33    use vortex_array::builtins::ArrayBuiltins;
34    use vortex_array::compute::conformance::cast::test_cast_conformance;
35    use vortex_array::dtype::DType;
36    use vortex_array::dtype::Nullability;
37    use vortex_array::dtype::PType;
38
39    use crate::ZigZagArray;
40    use crate::zigzag_encode;
41
42    #[test]
43    fn test_cast_zigzag_i32_to_i64() {
44        let values = PrimitiveArray::from_iter([-100i32, -1, 0, 1, 100]);
45        let zigzag = zigzag_encode(values).unwrap();
46
47        let casted = zigzag
48            .to_array()
49            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
50            .unwrap();
51        assert_eq!(
52            casted.dtype(),
53            &DType::Primitive(PType::I64, Nullability::NonNullable)
54        );
55
56        // Verify the result is still a ZigZagArray (not decoded)
57        // Note: The result might be wrapped, so let's check the encoding ID
58        assert_eq!(
59            casted.encoding_id().as_ref(),
60            "vortex.zigzag",
61            "Cast should preserve ZigZag encoding"
62        );
63
64        assert_arrays_eq!(casted, PrimitiveArray::from_iter([-100i64, -1, 0, 1, 100]));
65    }
66
67    #[test]
68    fn test_cast_zigzag_width_changes() {
69        // Test i32 to i16 (narrowing)
70        let values = PrimitiveArray::from_iter([100i32, -50, 0, 25, -100]);
71        let zigzag = zigzag_encode(values).unwrap();
72
73        let casted = zigzag
74            .to_array()
75            .cast(DType::Primitive(PType::I16, Nullability::NonNullable))
76            .unwrap();
77        assert_eq!(
78            casted.encoding_id().as_ref(),
79            "vortex.zigzag",
80            "Should remain ZigZag encoded"
81        );
82
83        assert_arrays_eq!(
84            casted,
85            PrimitiveArray::from_iter([100i16, -50, 0, 25, -100])
86        );
87
88        // Test i16 to i64 (widening)
89        let values16 = PrimitiveArray::from_iter([1000i16, -500, 0, 250, -1000]);
90        let zigzag16 = zigzag_encode(values16).unwrap();
91
92        let casted64 = zigzag16
93            .to_array()
94            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
95            .unwrap();
96        assert_eq!(
97            casted64.encoding_id().as_ref(),
98            "vortex.zigzag",
99            "Should remain ZigZag encoded"
100        );
101
102        assert_arrays_eq!(
103            casted64,
104            PrimitiveArray::from_iter([1000i64, -500, 0, 250, -1000])
105        );
106    }
107
108    #[test]
109    fn test_cast_zigzag_nullable() {
110        let values =
111            PrimitiveArray::from_option_iter([Some(-10i32), None, Some(0), Some(10), None]);
112        let zigzag = zigzag_encode(values).unwrap();
113
114        let casted = zigzag
115            .to_array()
116            .cast(DType::Primitive(PType::I64, Nullability::Nullable))
117            .unwrap();
118        assert_eq!(
119            casted.dtype(),
120            &DType::Primitive(PType::I64, Nullability::Nullable)
121        );
122    }
123
124    #[rstest]
125    #[case(zigzag_encode(PrimitiveArray::from_iter([-100i32, -50, -1, 0, 1, 50, 100])).unwrap())]
126    #[case(zigzag_encode(PrimitiveArray::from_iter([-1000i64, -1, 0, 1, 1000])).unwrap())]
127    #[case(zigzag_encode(PrimitiveArray::from_option_iter([Some(-5i16), None, Some(0), Some(5), None])).unwrap())]
128    #[case(zigzag_encode(PrimitiveArray::from_iter([i32::MIN, -1, 0, 1, i32::MAX])).unwrap())]
129    fn test_cast_zigzag_conformance(#[case] array: ZigZagArray) {
130        test_cast_conformance(&array.to_array());
131    }
132}
vortex_zigzag/compute/cast.rs

vortex_zigzag/compute/
cast.rs
