vortex_zigzag/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::stats::{ArrayStats, StatsSetRef};
5use vortex_array::vtable::{
6    ArrayVTable, CanonicalVTable, NotSupported, OperationsVTable, VTable, ValidityChild,
7    ValidityVTableFromChild,
8};
9use vortex_array::{
10    Array, ArrayRef, Canonical, EncodingId, EncodingRef, IntoArray, ToCanonical, vtable,
11};
12use vortex_dtype::{DType, PType, match_each_unsigned_integer_ptype};
13use vortex_error::{VortexExpect, VortexResult, vortex_bail};
14use vortex_scalar::Scalar;
15use zigzag::ZigZag as ExternalZigZag;
16
17use crate::compute::ZigZagEncoded;
18use crate::zigzag_decode;
19
20vtable!(ZigZag);
21
22impl VTable for ZigZagVTable {
23    type Array = ZigZagArray;
24    type Encoding = ZigZagEncoding;
25
26    type ArrayVTable = Self;
27    type CanonicalVTable = Self;
28    type OperationsVTable = Self;
29    type ValidityVTable = ValidityVTableFromChild;
30    type VisitorVTable = Self;
31    type ComputeVTable = NotSupported;
32    type EncodeVTable = Self;
33    type SerdeVTable = Self;
34
35    fn id(_encoding: &Self::Encoding) -> EncodingId {
36        EncodingId::new_ref("vortex.zigzag")
37    }
38
39    fn encoding(_array: &Self::Array) -> EncodingRef {
40        EncodingRef::new_ref(ZigZagEncoding.as_ref())
41    }
42}
43
44#[derive(Clone, Debug)]
45pub struct ZigZagArray {
46    dtype: DType,
47    encoded: ArrayRef,
48    stats_set: ArrayStats,
49}
50
51#[derive(Clone, Debug)]
52pub struct ZigZagEncoding;
53
54impl ZigZagArray {
55    pub fn new(encoded: ArrayRef) -> Self {
56        Self::try_new(encoded).vortex_expect("ZigZigArray new")
57    }
58
59    pub fn try_new(encoded: ArrayRef) -> VortexResult<Self> {
60        let encoded_dtype = encoded.dtype().clone();
61        if !encoded_dtype.is_unsigned_int() {
62            vortex_bail!(MismatchedTypes: "unsigned int", encoded_dtype);
63        }
64
65        let dtype = DType::from(PType::try_from(&encoded_dtype)?.to_signed())
66            .with_nullability(encoded_dtype.nullability());
67
68        Ok(Self {
69            dtype,
70            encoded,
71            stats_set: Default::default(),
72        })
73    }
74
75    pub fn ptype(&self) -> PType {
76        self.dtype().as_ptype()
77    }
78
79    pub fn encoded(&self) -> &ArrayRef {
80        &self.encoded
81    }
82}
83
84impl ArrayVTable<ZigZagVTable> for ZigZagVTable {
85    fn len(array: &ZigZagArray) -> usize {
86        array.encoded.len()
87    }
88
89    fn dtype(array: &ZigZagArray) -> &DType {
90        &array.dtype
91    }
92
93    fn stats(array: &ZigZagArray) -> StatsSetRef<'_> {
94        array.stats_set.to_ref(array.as_ref())
95    }
96}
97
98impl CanonicalVTable<ZigZagVTable> for ZigZagVTable {
99    fn canonicalize(array: &ZigZagArray) -> VortexResult<Canonical> {
100        zigzag_decode(array.encoded().to_primitive()?).map(Canonical::Primitive)
101    }
102}
103
104impl OperationsVTable<ZigZagVTable> for ZigZagVTable {
105    fn slice(array: &ZigZagArray, start: usize, stop: usize) -> ArrayRef {
106        ZigZagArray::new(array.encoded().slice(start, stop)).into_array()
107    }
108
109    fn scalar_at(array: &ZigZagArray, index: usize) -> Scalar {
110        let scalar = array.encoded().scalar_at(index);
111        if scalar.is_null() {
112            return scalar.reinterpret_cast(array.ptype());
113        }
114
115        let pscalar = scalar.as_primitive();
116        match_each_unsigned_integer_ptype!(pscalar.ptype(), |P| {
117            Scalar::primitive(
118                <<P as ZigZagEncoded>::Int>::decode(
119                    pscalar
120                        .typed_value::<P>()
121                        .vortex_expect("zigzag corruption"),
122                ),
123                array.dtype().nullability(),
124            )
125        })
126    }
127}
128
129impl ValidityChild<ZigZagVTable> for ZigZagVTable {
130    fn validity_child(array: &ZigZagArray) -> &dyn Array {
131        array.encoded()
132    }
133}
134
135#[cfg(test)]
136mod test {
137    use vortex_array::IntoArray;
138    use vortex_buffer::buffer;
139    use vortex_scalar::Scalar;
140
141    use super::*;
142
143    #[test]
144    fn test_compute_statistics() {
145        let array = buffer![1i32, -5i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array();
146        let canonical = array.to_canonical().unwrap();
147        let zigzag = ZigZagEncoding.encode(&canonical, None).unwrap().unwrap();
148
149        assert_eq!(
150            zigzag.statistics().compute_max::<i32>(),
151            array.statistics().compute_max::<i32>()
152        );
153        assert_eq!(
154            zigzag.statistics().compute_null_count(),
155            array.statistics().compute_null_count()
156        );
157        assert_eq!(
158            zigzag.statistics().compute_is_constant(),
159            array.statistics().compute_is_constant()
160        );
161
162        let sliced = zigzag.slice(0, 2);
163        let sliced = sliced.as_::<ZigZagVTable>();
164        assert_eq!(sliced.scalar_at(sliced.len() - 1), Scalar::from(-5i32));
165
166        assert_eq!(
167            sliced.statistics().compute_min::<i32>(),
168            array.statistics().compute_min::<i32>()
169        );
170        assert_eq!(
171            sliced.statistics().compute_null_count(),
172            array.statistics().compute_null_count()
173        );
174        assert_eq!(
175            sliced.statistics().compute_is_constant(),
176            array.statistics().compute_is_constant()
177        );
178    }
179}
vortex_zigzag/array.rs

vortex_zigzag/
array.rs
