vortex_zigzag/
compute.rs

1use vortex_array::compute::{
2    FilterFn, ScalarAtFn, SliceFn, TakeFn, filter, scalar_at, slice, take,
3};
4use vortex_array::variants::PrimitiveArrayTrait;
5use vortex_array::vtable::ComputeVTable;
6use vortex_array::{Array, ArrayRef};
7use vortex_dtype::match_each_unsigned_integer_ptype;
8use vortex_error::{VortexResult, vortex_err};
9use vortex_mask::Mask;
10use vortex_scalar::{PrimitiveScalar, Scalar};
11use zigzag::{ZigZag as ExternalZigZag, ZigZag};
12
13use crate::{ZigZagArray, ZigZagEncoding};
14
15impl ComputeVTable for ZigZagEncoding {
16    fn filter_fn(&self) -> Option<&dyn FilterFn<&dyn Array>> {
17        Some(self)
18    }
19
20    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
21        Some(self)
22    }
23
24    fn slice_fn(&self) -> Option<&dyn SliceFn<&dyn Array>> {
25        Some(self)
26    }
27
28    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
29        Some(self)
30    }
31}
32
33impl FilterFn<&ZigZagArray> for ZigZagEncoding {
34    fn filter(&self, array: &ZigZagArray, mask: &Mask) -> VortexResult<ArrayRef> {
35        let encoded = filter(array.encoded(), mask)?;
36        Ok(ZigZagArray::try_new(encoded)?.into_array())
37    }
38}
39
40impl ScalarAtFn<&ZigZagArray> for ZigZagEncoding {
41    fn scalar_at(&self, array: &ZigZagArray, index: usize) -> VortexResult<Scalar> {
42        let scalar = scalar_at(array.encoded(), index)?;
43        if scalar.is_null() {
44            return Ok(scalar.reinterpret_cast(array.ptype()));
45        }
46
47        let pscalar = PrimitiveScalar::try_from(&scalar)?;
48        match_each_unsigned_integer_ptype!(pscalar.ptype(), |$P| {
49            Ok(Scalar::primitive(
50                <<$P as ZigZagEncoded>::Int>::decode(pscalar.typed_value::<$P>().ok_or_else(|| {
51                    vortex_err!(
52                        "Cannot decode provided scalar: expected {}, got ptype {}",
53                        std::any::type_name::<$P>(),
54                        pscalar.ptype()
55                    )
56                })?),
57                array.dtype().nullability(),
58            ))
59        })
60    }
61}
62
63impl SliceFn<&ZigZagArray> for ZigZagEncoding {
64    fn slice(&self, array: &ZigZagArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
65        Ok(ZigZagArray::try_new(slice(array.encoded(), start, stop)?)?.into_array())
66    }
67}
68
69impl TakeFn<&ZigZagArray> for ZigZagEncoding {
70    fn take(&self, array: &ZigZagArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
71        let encoded = take(array.encoded(), indices)?;
72        Ok(ZigZagArray::try_new(encoded)?.into_array())
73    }
74}
75
76trait ZigZagEncoded {
77    type Int: ZigZag;
78}
79
80impl ZigZagEncoded for u8 {
81    type Int = i8;
82}
83
84impl ZigZagEncoded for u16 {
85    type Int = i16;
86}
87
88impl ZigZagEncoded for u32 {
89    type Int = i32;
90}
91
92impl ZigZagEncoded for u64 {
93    type Int = i64;
94}
95
96#[cfg(test)]
97mod tests {
98    use vortex_array::arrays::{BooleanBuffer, PrimitiveArray};
99    use vortex_array::compute::{
100        SearchResult, SearchSortedSide, filter, scalar_at, search_sorted, take,
101    };
102    use vortex_array::validity::Validity;
103    use vortex_array::{IntoArray, ToCanonical};
104    use vortex_buffer::buffer;
105    use vortex_dtype::Nullability;
106    use vortex_scalar::Scalar;
107
108    use crate::ZigZagArray;
109
110    #[test]
111    pub fn search_sorted_uncompressed() {
112        let zigzag = ZigZagArray::encode(&buffer![-189, -160, 1].into_array()).unwrap();
113        assert_eq!(
114            search_sorted(&zigzag, -169, SearchSortedSide::Right).unwrap(),
115            SearchResult::NotFound(1)
116        );
117    }
118
119    #[test]
120    pub fn nullable_scalar_at() {
121        let zigzag = ZigZagArray::encode(&PrimitiveArray::new(
122            buffer![-189, -160, 1],
123            Validity::AllValid,
124        ))
125        .unwrap();
126        assert_eq!(
127            scalar_at(&zigzag, 1).unwrap(),
128            Scalar::primitive(-160, Nullability::Nullable)
129        );
130    }
131
132    #[test]
133    fn take_zigzag() {
134        let zigzag = ZigZagArray::encode(&buffer![-189, -160, 1].into_array()).unwrap();
135        let indices = buffer![0, 2].into_array();
136        let actual = take(&zigzag, &indices).unwrap().to_primitive().unwrap();
137        let expected = ZigZagArray::encode(&buffer![-189, 1].into_array())
138            .unwrap()
139            .to_primitive()
140            .unwrap();
141        assert_eq!(actual.as_slice::<i32>(), expected.as_slice::<i32>());
142    }
143
144    #[test]
145    fn filter_zigzag() {
146        let zigzag = ZigZagArray::encode(&buffer![-189, -160, 1].into_array()).unwrap();
147        let filter_mask = BooleanBuffer::from(vec![true, false, true]).into();
148        let actual = filter(&zigzag, &filter_mask)
149            .unwrap()
150            .to_primitive()
151            .unwrap();
152        let expected = ZigZagArray::encode(&buffer![-189, 1].into_array())
153            .unwrap()
154            .to_primitive()
155            .unwrap();
156        assert_eq!(actual.as_slice::<i32>(), expected.as_slice::<i32>());
157    }
158}
vortex_zigzag/compute.rs

vortex_zigzag/
compute.rs
