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

vortex_zigzag/
array.rs
