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::ArrayView;
6use vortex_array::IntoArray;
7use vortex_array::builtins::ArrayBuiltins;
8use vortex_array::dtype::DType;
9use vortex_array::scalar_fn::fns::cast::CastReduce;
10use vortex_error::VortexResult;
11
12use crate::ZigZag;
13use crate::array::ZigZagArrayExt;
14impl CastReduce for ZigZag {
15    fn cast(array: ArrayView<'_, Self>, 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(ZigZag::try_new(new_encoded)?.into_array()))
24    }
25}
26
27#[cfg(test)]
28mod tests {
29    use std::sync::LazyLock;
30
31    use rstest::rstest;
32    use vortex_array::IntoArray;
33    use vortex_array::VortexSessionExecute;
34    use vortex_array::arrays::PrimitiveArray;
35    use vortex_array::assert_arrays_eq;
36    use vortex_array::builtins::ArrayBuiltins;
37    use vortex_array::compute::conformance::cast::test_cast_conformance;
38    use vortex_array::dtype::DType;
39    use vortex_array::dtype::Nullability;
40    use vortex_array::dtype::PType;
41    use vortex_session::VortexSession;
42
43    use crate::ZigZagArray;
44    use crate::zigzag_encode;
45
46    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
47        let session = vortex_array::array_session();
48        crate::initialize(&session);
49        session
50    });
51
52    #[test]
53    fn test_cast_zigzag_i32_to_i64() {
54        let values = PrimitiveArray::from_iter([-100i32, -1, 0, 1, 100]);
55        let zigzag = zigzag_encode(values.as_view()).unwrap();
56
57        let casted = zigzag
58            .into_array()
59            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
60            .unwrap();
61        assert_eq!(
62            casted.dtype(),
63            &DType::Primitive(PType::I64, Nullability::NonNullable)
64        );
65
66        // Verify the result is still a ZigZagArray (not decoded)
67        // Note: The result might be wrapped, so let's check the encoding ID
68        assert_eq!(
69            casted.encoding_id().as_ref(),
70            "vortex.zigzag",
71            "Cast should preserve ZigZag encoding"
72        );
73
74        assert_arrays_eq!(
75            casted,
76            PrimitiveArray::from_iter([-100i64, -1, 0, 1, 100]),
77            &mut SESSION.create_execution_ctx()
78        );
79    }
80
81    #[test]
82    fn test_cast_zigzag_width_changes() {
83        // Test i32 to i16 (narrowing)
84        let values = PrimitiveArray::from_iter([100i32, -50, 0, 25, -100]);
85        let zigzag = zigzag_encode(values.as_view()).unwrap();
86
87        let casted = zigzag
88            .into_array()
89            .cast(DType::Primitive(PType::I16, Nullability::NonNullable))
90            .unwrap();
91        assert_eq!(
92            casted.encoding_id().as_ref(),
93            "vortex.zigzag",
94            "Should remain ZigZag encoded"
95        );
96
97        assert_arrays_eq!(
98            casted,
99            PrimitiveArray::from_iter([100i16, -50, 0, 25, -100]),
100            &mut SESSION.create_execution_ctx()
101        );
102
103        // Test i16 to i64 (widening)
104        let values16 = PrimitiveArray::from_iter([1000i16, -500, 0, 250, -1000]);
105        let zigzag16 = zigzag_encode(values16.as_view()).unwrap();
106
107        let casted64 = zigzag16
108            .into_array()
109            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
110            .unwrap();
111        assert_eq!(
112            casted64.encoding_id().as_ref(),
113            "vortex.zigzag",
114            "Should remain ZigZag encoded"
115        );
116
117        assert_arrays_eq!(
118            casted64,
119            PrimitiveArray::from_iter([1000i64, -500, 0, 250, -1000]),
120            &mut SESSION.create_execution_ctx()
121        );
122    }
123
124    #[test]
125    fn test_cast_zigzag_nullable() {
126        let values =
127            PrimitiveArray::from_option_iter([Some(-10i32), None, Some(0), Some(10), None]);
128        let zigzag = zigzag_encode(values.as_view()).unwrap();
129
130        let casted = zigzag
131            .into_array()
132            .cast(DType::Primitive(PType::I64, Nullability::Nullable))
133            .unwrap();
134        assert_eq!(
135            casted.dtype(),
136            &DType::Primitive(PType::I64, Nullability::Nullable)
137        );
138    }
139
140    #[rstest]
141    #[case(zigzag_encode(PrimitiveArray::from_iter([-100i32, -50, -1, 0, 1, 50, 100]).as_view()).unwrap())]
142    #[case(zigzag_encode(PrimitiveArray::from_iter([-1000i64, -1, 0, 1, 1000]).as_view()).unwrap())]
143    #[case(zigzag_encode(PrimitiveArray::from_option_iter([Some(-5i16), None, Some(0), Some(5), None]).as_view()).unwrap())]
144    #[case(zigzag_encode(PrimitiveArray::from_iter([i32::MIN, -1, 0, 1, i32::MAX]).as_view()).unwrap())]
145    fn test_cast_zigzag_conformance(#[case] array: ZigZagArray) {
146        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
147    }
148}
vortex_zigzag/compute/cast.rs

vortex_zigzag/compute/
cast.rs
