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

vortex_zigzag/
array.rs
