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

vortex_zigzag/compute/
mod.rs
