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

vortex_zigzag/compute/
mod.rs
