vortex_zigzag/
array.rs

1use vortex_array::stats::{ArrayStats, StatsSetRef};
2use vortex_array::variants::PrimitiveArrayTrait;
3use vortex_array::vtable::VTableRef;
4use vortex_array::{
5    Array, ArrayCanonicalImpl, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayValidityImpl,
6    ArrayVariantsImpl, Canonical, EmptyMetadata, Encoding, ToCanonical, try_from_array_ref,
7};
8use vortex_dtype::{DType, PType};
9use vortex_error::{VortexResult, vortex_bail};
10use vortex_mask::Mask;
11
12use crate::zigzag_decode;
13
14#[derive(Clone, Debug)]
15pub struct ZigZagArray {
16    dtype: DType,
17    encoded: ArrayRef,
18    stats_set: ArrayStats,
19}
20
21try_from_array_ref!(ZigZagArray);
22
23pub struct ZigZagEncoding;
24impl Encoding for ZigZagEncoding {
25    type Array = ZigZagArray;
26    type Metadata = EmptyMetadata;
27}
28
29impl ZigZagArray {
30    pub fn try_new(encoded: ArrayRef) -> VortexResult<Self> {
31        let encoded_dtype = encoded.dtype().clone();
32        if !encoded_dtype.is_unsigned_int() {
33            vortex_bail!(MismatchedTypes: "unsigned int", encoded_dtype);
34        }
35
36        let dtype = DType::from(PType::try_from(&encoded_dtype)?.to_signed())
37            .with_nullability(encoded_dtype.nullability());
38
39        Ok(Self {
40            dtype,
41            encoded,
42            stats_set: Default::default(),
43        })
44    }
45
46    pub fn encoded(&self) -> &ArrayRef {
47        &self.encoded
48    }
49}
50
51impl ArrayImpl for ZigZagArray {
52    type Encoding = ZigZagEncoding;
53
54    fn _len(&self) -> usize {
55        self.encoded.len()
56    }
57
58    fn _dtype(&self) -> &DType {
59        &self.dtype
60    }
61
62    fn _vtable(&self) -> VTableRef {
63        VTableRef::new_ref(&ZigZagEncoding)
64    }
65
66    fn _with_children(&self, children: &[ArrayRef]) -> VortexResult<Self> {
67        let encoded = children[0].clone();
68
69        Self::try_new(encoded)
70    }
71}
72
73impl ArrayCanonicalImpl for ZigZagArray {
74    fn _to_canonical(&self) -> VortexResult<Canonical> {
75        zigzag_decode(self.encoded().to_primitive()?).map(Canonical::Primitive)
76    }
77}
78
79impl ArrayStatisticsImpl for ZigZagArray {
80    fn _stats_ref(&self) -> StatsSetRef<'_> {
81        self.stats_set.to_ref(self)
82    }
83}
84
85impl ArrayValidityImpl for ZigZagArray {
86    fn _is_valid(&self, index: usize) -> VortexResult<bool> {
87        self.encoded.is_valid(index)
88    }
89
90    fn _all_valid(&self) -> VortexResult<bool> {
91        self.encoded.all_valid()
92    }
93
94    fn _all_invalid(&self) -> VortexResult<bool> {
95        self.encoded.all_invalid()
96    }
97
98    fn _validity_mask(&self) -> VortexResult<Mask> {
99        self.encoded.validity_mask()
100    }
101}
102
103impl ArrayVariantsImpl for ZigZagArray {
104    fn _as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
105        Some(self)
106    }
107}
108
109impl PrimitiveArrayTrait for ZigZagArray {}
110
111#[cfg(test)]
112mod test {
113    use vortex_array::IntoArray;
114    use vortex_array::compute::{scalar_at, slice};
115    use vortex_array::vtable::EncodingVTable;
116    use vortex_buffer::buffer;
117    use vortex_scalar::Scalar;
118
119    use super::*;
120
121    #[test]
122    fn test_compute_statistics() {
123        let array = buffer![1i32, -5i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array();
124        let canonical = array.to_canonical().unwrap();
125        let zigzag = ZigZagEncoding.encode(&canonical, None).unwrap().unwrap();
126
127        assert_eq!(
128            zigzag.statistics().compute_max::<i32>(),
129            array.statistics().compute_max::<i32>()
130        );
131        assert_eq!(
132            zigzag.statistics().compute_null_count(),
133            array.statistics().compute_null_count()
134        );
135        assert_eq!(
136            zigzag.statistics().compute_is_constant(),
137            array.statistics().compute_is_constant()
138        );
139
140        let sliced = ZigZagArray::try_from(slice(&zigzag, 0, 2).unwrap()).unwrap();
141        assert_eq!(
142            scalar_at(&sliced, sliced.len() - 1).unwrap(),
143            Scalar::from(-5i32)
144        );
145
146        assert_eq!(
147            sliced.statistics().compute_min::<i32>(),
148            array.statistics().compute_min::<i32>()
149        );
150        assert_eq!(
151            sliced.statistics().compute_null_count(),
152            array.statistics().compute_null_count()
153        );
154        assert_eq!(
155            sliced.statistics().compute_is_constant(),
156            array.statistics().compute_is_constant()
157        );
158    }
159}
vortex_zigzag/array.rs

vortex_zigzag/
array.rs
