vortex_zigzag/compute/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod cast;
5
6use vortex_array::ArrayRef;
7use vortex_array::ArrayView;
8use vortex_array::ExecutionCtx;
9use vortex_array::IntoArray;
10use vortex_array::arrays::dict::TakeExecute;
11use vortex_array::arrays::filter::FilterReduce;
12use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr;
13use vortex_array::scalar_fn::fns::mask::MaskReduce;
14use vortex_error::VortexResult;
15use vortex_mask::Mask;
16
17use crate::ZigZag;
18use crate::array::ZigZagArraySlotsExt;
19
20impl FilterReduce for ZigZag {
21    fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
22        let encoded = array.encoded().filter(mask.clone())?;
23        Ok(Some(ZigZag::try_new(encoded)?.into_array()))
24    }
25}
26
27impl TakeExecute for ZigZag {
28    fn take(
29        array: ArrayView<'_, Self>,
30        indices: &ArrayRef,
31        _ctx: &mut ExecutionCtx,
32    ) -> VortexResult<Option<ArrayRef>> {
33        let encoded = array.encoded().take(indices.clone())?;
34        Ok(Some(ZigZag::try_new(encoded)?.into_array()))
35    }
36}
37
38impl MaskReduce for ZigZag {
39    fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
40        let masked_encoded = MaskExpr::try_new(array.encoded().clone(), mask.clone())?.into_array();
41        Ok(Some(ZigZag::try_new(masked_encoded)?.into_array()))
42    }
43}
44
45pub(crate) trait ZigZagEncoded {
46    type Int: zigzag::ZigZag;
47}
48
49impl ZigZagEncoded for u8 {
50    type Int = i8;
51}
52
53impl ZigZagEncoded for u16 {
54    type Int = i16;
55}
56
57impl ZigZagEncoded for u32 {
58    type Int = i32;
59}
60
61impl ZigZagEncoded for u64 {
62    type Int = i64;
63}
64
65#[cfg(test)]
66mod tests {
67    use std::sync::LazyLock;
68
69    use rstest::rstest;
70    use vortex_array::ArrayRef;
71    use vortex_array::IntoArray;
72    use vortex_array::VortexSessionExecute;
73    use vortex_array::arrays::PrimitiveArray;
74    use vortex_array::assert_arrays_eq;
75    use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array;
76    use vortex_array::compute::conformance::consistency::test_array_consistency;
77    use vortex_array::dtype::Nullability;
78    use vortex_array::scalar::Scalar;
79    use vortex_array::validity::Validity;
80    use vortex_buffer::BitBuffer;
81    use vortex_buffer::buffer;
82    use vortex_error::VortexResult;
83    use vortex_session::VortexSession;
84
85    use crate::ZigZagArray;
86    use crate::zigzag_encode;
87
88    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
89        let session = vortex_array::array_session();
90        crate::initialize(&session);
91        session
92    });
93
94    #[test]
95    pub fn nullable_scalar_at() -> VortexResult<()> {
96        let zigzag = zigzag_encode(
97            PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid).as_view(),
98        )?;
99        assert_eq!(
100            zigzag.execute_scalar(1, &mut SESSION.create_execution_ctx())?,
101            Scalar::primitive(-160, Nullability::Nullable)
102        );
103        Ok(())
104    }
105
106    #[test]
107    fn take_zigzag() -> VortexResult<()> {
108        let zigzag = zigzag_encode(
109            PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid).as_view(),
110        )?;
111
112        let indices = buffer![0, 2].into_array();
113        let actual = zigzag.take(indices)?;
114        let expected =
115            zigzag_encode(PrimitiveArray::new(buffer![-189, 1], Validity::AllValid).as_view())?
116                .into_array();
117        assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx());
118        Ok(())
119    }
120
121    #[test]
122    fn filter_zigzag() -> VortexResult<()> {
123        let zigzag = zigzag_encode(
124            PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid).as_view(),
125        )?;
126
127        let filter_mask = BitBuffer::from(vec![true, false, true]).into();
128        let actual = zigzag.filter(filter_mask)?;
129        let expected =
130            zigzag_encode(PrimitiveArray::new(buffer![-189, 1], Validity::AllValid).as_view())?
131                .into_array();
132        assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx());
133        Ok(())
134    }
135
136    #[test]
137    fn test_filter_conformance() -> VortexResult<()> {
138        use vortex_array::compute::conformance::filter::test_filter_conformance;
139
140        // Test with i32 values
141        let zigzag = zigzag_encode(
142            PrimitiveArray::new(buffer![-189i32, -160, 1, 42, -73], Validity::AllValid).as_view(),
143        )?;
144        test_filter_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx());
145
146        // Test with i64 values
147        let zigzag = zigzag_encode(
148            PrimitiveArray::new(
149                buffer![1000i64, -2000, 3000, -4000, 5000],
150                Validity::AllValid,
151            )
152            .as_view(),
153        )?;
154        test_filter_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx());
155
156        // Test with nullable values
157        let array =
158            PrimitiveArray::from_option_iter([Some(-10i16), None, Some(20), Some(-30), None]);
159        let zigzag = zigzag_encode(array.as_view())?;
160        test_filter_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx());
161        Ok(())
162    }
163
164    #[test]
165    fn test_mask_conformance() -> VortexResult<()> {
166        use vortex_array::compute::conformance::mask::test_mask_conformance;
167
168        // Test with i32 values
169        let zigzag = zigzag_encode(
170            PrimitiveArray::new(buffer![-100i32, 200, -300, 400, -500], Validity::AllValid)
171                .as_view(),
172        )?;
173        test_mask_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx());
174
175        // Test with i8 values
176        let zigzag = zigzag_encode(
177            PrimitiveArray::new(buffer![-127i8, 0, 127, -1, 1], Validity::AllValid).as_view(),
178        )?;
179        test_mask_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx());
180        Ok(())
181    }
182
183    #[rstest]
184    #[case(buffer![-189i32, -160, 1, 42, -73].into_array())]
185    #[case(buffer![1000i64, -2000, 3000, -4000, 5000].into_array())]
186    #[case(PrimitiveArray::from_option_iter([Some(-10i16), None, Some(20), Some(-30), None]).into_array()
187    )]
188    #[case(buffer![42i32].into_array())]
189    fn test_take_zigzag_conformance(#[case] array: ArrayRef) -> VortexResult<()> {
190        use vortex_array::compute::conformance::take::test_take_conformance;
191
192        let mut ctx = SESSION.create_execution_ctx();
193        let array_primitive = array.execute::<PrimitiveArray>(&mut ctx)?;
194        let zigzag = zigzag_encode(array_primitive.as_view())?;
195        test_take_conformance(&zigzag.into_array(), &mut ctx);
196        Ok(())
197    }
198
199    #[rstest]
200    // Basic ZigZag arrays
201    #[case::zigzag_i8(zigzag_encode(PrimitiveArray::from_iter([-128i8, -1, 0, 1, 127]).as_view()).unwrap())]
202    #[case::zigzag_i16(zigzag_encode(PrimitiveArray::from_iter([-1000i16, -100, 0, 100, 1000]).as_view()).unwrap())]
203    #[case::zigzag_i32(zigzag_encode(PrimitiveArray::from_iter([-100000i32, -1000, 0, 1000, 100000]).as_view()).unwrap())]
204    #[case::zigzag_i64(zigzag_encode(PrimitiveArray::from_iter([-1000000i64, -10000, 0, 10000, 1000000]).as_view()).unwrap())]
205    // Nullable arrays
206    #[case::zigzag_nullable_i32(zigzag_encode(PrimitiveArray::from_option_iter([Some(-100i32), None, Some(0), Some(100), None]).as_view()).unwrap())]
207    #[case::zigzag_nullable_i64(zigzag_encode(PrimitiveArray::from_option_iter([Some(-1000i64), None, Some(0), Some(1000), None]).as_view()).unwrap())]
208    // Edge cases
209    #[case::zigzag_single(zigzag_encode(PrimitiveArray::from_iter([-42i32]).as_view()).unwrap())]
210    #[case::zigzag_alternating(zigzag_encode(PrimitiveArray::from_iter([-1i32, 1, -2, 2, -3, 3]).as_view()).unwrap())]
211    // Large arrays
212    #[case::zigzag_large_i32(zigzag_encode(PrimitiveArray::from_iter(-500..500).as_view()).unwrap())]
213    #[case::zigzag_large_i64(zigzag_encode(PrimitiveArray::from_iter((-1000..1000).map(|i| i as i64 * 100)).as_view()).unwrap())]
214    fn test_zigzag_consistency(#[case] array: ZigZagArray) {
215        test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx());
216    }
217
218    #[rstest]
219    #[case::zigzag_i8_basic(zigzag_encode(PrimitiveArray::from_iter([-10i8, -5, 0, 5, 10]).as_view()).unwrap())]
220    #[case::zigzag_i16_basic(zigzag_encode(PrimitiveArray::from_iter([-100i16, -50, 0, 50, 100]).as_view()).unwrap())]
221    #[case::zigzag_i32_basic(zigzag_encode(PrimitiveArray::from_iter([-1000i32, -500, 0, 500, 1000]).as_view()).unwrap())]
222    #[case::zigzag_i64_basic(zigzag_encode(PrimitiveArray::from_iter([-10000i64, -5000, 0, 5000, 10000]).as_view()).unwrap())]
223    #[case::zigzag_i32_large(zigzag_encode(PrimitiveArray::from_iter((-50..50).map(|i| i * 10)).as_view()).unwrap())]
224    fn test_zigzag_binary_numeric(#[case] array: ZigZagArray) {
225        test_binary_numeric_array(&array.into_array(), &mut SESSION.create_execution_ctx());
226    }
227}
vortex_zigzag/compute/mod.rs

vortex_zigzag/compute/
mod.rs
