vortex_zigzag/
compute.rs

1use vortex_array::compute::{
2    FilterKernel, FilterKernelAdapter, TakeKernel, TakeKernelAdapter, filter, take,
3};
4use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
5use vortex_error::VortexResult;
6use vortex_mask::Mask;
7
8use crate::{ZigZagArray, ZigZagVTable};
9
10impl FilterKernel for ZigZagVTable {
11    fn filter(&self, array: &ZigZagArray, mask: &Mask) -> VortexResult<ArrayRef> {
12        let encoded = filter(array.encoded(), mask)?;
13        Ok(ZigZagArray::try_new(encoded)?.into_array())
14    }
15}
16
17register_kernel!(FilterKernelAdapter(ZigZagVTable).lift());
18
19impl TakeKernel for ZigZagVTable {
20    fn take(&self, array: &ZigZagArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
21        let encoded = take(array.encoded(), indices)?;
22        Ok(ZigZagArray::try_new(encoded)?.into_array())
23    }
24}
25
26register_kernel!(TakeKernelAdapter(ZigZagVTable).lift());
27
28pub(crate) trait ZigZagEncoded {
29    type Int: zigzag::ZigZag;
30}
31
32impl ZigZagEncoded for u8 {
33    type Int = i8;
34}
35
36impl ZigZagEncoded for u16 {
37    type Int = i16;
38}
39
40impl ZigZagEncoded for u32 {
41    type Int = i32;
42}
43
44impl ZigZagEncoded for u64 {
45    type Int = i64;
46}
47
48#[cfg(test)]
49mod tests {
50    use vortex_array::arrays::{BooleanBuffer, PrimitiveArray};
51    use vortex_array::compute::{filter, take};
52    use vortex_array::validity::Validity;
53    use vortex_array::{Array, IntoArray, ToCanonical};
54    use vortex_buffer::buffer;
55    use vortex_dtype::Nullability;
56    use vortex_scalar::Scalar;
57
58    use crate::ZigZagEncoding;
59
60    #[test]
61    pub fn nullable_scalar_at() {
62        let zigzag = ZigZagEncoding
63            .encode(
64                &PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid)
65                    .to_canonical()
66                    .unwrap(),
67                None,
68            )
69            .unwrap()
70            .unwrap();
71        assert_eq!(
72            zigzag.scalar_at(1).unwrap(),
73            Scalar::primitive(-160, Nullability::Nullable)
74        );
75    }
76
77    #[test]
78    fn take_zigzag() {
79        let zigzag = ZigZagEncoding
80            .encode(
81                &buffer![-189, -160, 1].into_array().to_canonical().unwrap(),
82                None,
83            )
84            .unwrap()
85            .unwrap();
86
87        let indices = buffer![0, 2].into_array();
88        let actual = take(&zigzag, &indices).unwrap().to_primitive().unwrap();
89        let expected = ZigZagEncoding
90            .encode(&buffer![-189, 1].into_array().to_canonical().unwrap(), None)
91            .unwrap()
92            .unwrap()
93            .to_primitive()
94            .unwrap();
95        assert_eq!(actual.as_slice::<i32>(), expected.as_slice::<i32>());
96    }
97
98    #[test]
99    fn filter_zigzag() {
100        let zigzag = ZigZagEncoding
101            .encode(
102                &buffer![-189, -160, 1].into_array().to_canonical().unwrap(),
103                None,
104            )
105            .unwrap()
106            .unwrap();
107        let filter_mask = BooleanBuffer::from(vec![true, false, true]).into();
108        let actual = filter(&zigzag, &filter_mask)
109            .unwrap()
110            .to_primitive()
111            .unwrap();
112        let expected = ZigZagEncoding
113            .encode(&buffer![-189, 1].into_array().to_canonical().unwrap(), None)
114            .unwrap()
115            .unwrap()
116            .to_primitive()
117            .unwrap();
118        assert_eq!(actual.as_slice::<i32>(), expected.as_slice::<i32>());
119    }
120}
vortex_zigzag/compute.rs

vortex_zigzag/
compute.rs
