vortex_zigzag/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6use std::hash::Hasher;
7
8use vortex_array::Array;
9use vortex_array::ArrayEq;
10use vortex_array::ArrayHash;
11use vortex_array::ArrayId;
12use vortex_array::ArrayParts;
13use vortex_array::ArrayRef;
14use vortex_array::ArrayView;
15use vortex_array::EqMode;
16use vortex_array::ExecutionCtx;
17use vortex_array::ExecutionResult;
18use vortex_array::IntoArray;
19use vortex_array::TypedArrayRef;
20use vortex_array::array_slots;
21use vortex_array::buffer::BufferHandle;
22use vortex_array::dtype::DType;
23use vortex_array::dtype::PType;
24use vortex_array::match_each_unsigned_integer_ptype;
25use vortex_array::scalar::Scalar;
26use vortex_array::serde::ArrayChildren;
27use vortex_array::smallvec::smallvec;
28use vortex_array::vtable::OperationsVTable;
29use vortex_array::vtable::VTable;
30use vortex_array::vtable::ValidityChild;
31use vortex_array::vtable::ValidityVTableFromChild;
32use vortex_error::VortexExpect;
33use vortex_error::VortexResult;
34use vortex_error::vortex_bail;
35use vortex_error::vortex_ensure;
36use vortex_error::vortex_panic;
37use vortex_session::VortexSession;
38use vortex_session::registry::CachedId;
39use zigzag::ZigZag as ExternalZigZag;
40
41use crate::compute::ZigZagEncoded;
42use crate::rules::RULES;
43use crate::zigzag_decode;
44
45/// A [`ZigZag`]-encoded Vortex array.
46pub type ZigZagArray = Array<ZigZag>;
47
48impl VTable for ZigZag {
49    type TypedArrayData = ZigZagData;
50
51    type OperationsVTable = Self;
52    type ValidityVTable = ValidityVTableFromChild;
53
54    fn id(&self) -> ArrayId {
55        static ID: CachedId = CachedId::new("vortex.zigzag");
56        *ID
57    }
58
59    fn validate(
60        &self,
61        _data: &Self::TypedArrayData,
62        dtype: &DType,
63        len: usize,
64        slots: &[Option<ArrayRef>],
65    ) -> VortexResult<()> {
66        let encoded = ZigZagSlotsView::from_slots(slots).encoded;
67        let expected_dtype = ZigZagData::dtype_from_encoded_dtype(encoded.dtype())?;
68        vortex_ensure!(
69            dtype == &expected_dtype,
70            "expected dtype {expected_dtype}, got {dtype}"
71        );
72        vortex_ensure!(
73            encoded.len() == len,
74            "expected len {len}, got {}",
75            encoded.len()
76        );
77        Ok(())
78    }
79
80    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
81        0
82    }
83
84    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
85        vortex_panic!("ZigZagArray buffer index {idx} out of bounds")
86    }
87
88    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
89        vortex_panic!("ZigZagArray buffer_name index {idx} out of bounds")
90    }
91
92    fn with_buffers(
93        &self,
94        array: ArrayView<'_, Self>,
95        buffers: &[BufferHandle],
96    ) -> VortexResult<ArrayParts<Self>> {
97        vortex_array::vtable::with_empty_buffers(self, array, buffers)
98    }
99
100    fn serialize(
101        _array: ArrayView<'_, Self>,
102        _session: &VortexSession,
103    ) -> VortexResult<Option<Vec<u8>>> {
104        Ok(Some(vec![]))
105    }
106
107    fn deserialize(
108        &self,
109        dtype: &DType,
110        len: usize,
111        metadata: &[u8],
112        _buffers: &[BufferHandle],
113        children: &dyn ArrayChildren,
114        _session: &VortexSession,
115    ) -> VortexResult<ArrayParts<Self>> {
116        if !metadata.is_empty() {
117            vortex_bail!(
118                "ZigZagArray expects empty metadata, got {} bytes",
119                metadata.len()
120            );
121        }
122        if children.len() != 1 {
123            vortex_bail!("Expected 1 child, got {}", children.len());
124        }
125
126        let ptype = PType::try_from(dtype)?;
127        let encoded_type = DType::Primitive(ptype.to_unsigned(), dtype.nullability());
128
129        let encoded = children.get(0, &encoded_type, len)?;
130        let slots = smallvec![Some(encoded.clone())];
131        let data = ZigZagData::try_new(encoded.dtype())?;
132        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
133    }
134
135    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
136        ZigZagSlots::NAMES[idx].to_string()
137    }
138
139    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
140        Ok(ExecutionResult::done(
141            zigzag_decode(array.encoded().clone().execute(ctx)?).into_array(),
142        ))
143    }
144
145    fn reduce_parent(
146        array: ArrayView<'_, Self>,
147        parent: &ArrayRef,
148        child_idx: usize,
149    ) -> VortexResult<Option<ArrayRef>> {
150        RULES.evaluate(array, parent, child_idx)
151    }
152}
153
154impl ArrayHash for ZigZagData {
155    fn array_hash<H: Hasher>(&self, _state: &mut H, _accuracy: EqMode) {}
156}
157
158impl ArrayEq for ZigZagData {
159    fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool {
160        true
161    }
162}
163
164#[array_slots(ZigZag)]
165pub struct ZigZagSlots {
166    /// The zigzag-encoded values (signed integers mapped to unsigned).
167    #[slot(0)]
168    pub encoded: ArrayRef,
169}
170
171#[derive(Clone, Debug)]
172pub struct ZigZagData {}
173
174impl Display for ZigZagData {
175    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
176        Ok(())
177    }
178}
179
180pub trait ZigZagArrayExt: ZigZagArraySlotsExt {
181    fn ptype(&self) -> PType {
182        PType::try_from(self.encoded().dtype())
183            .vortex_expect("ZigZagArray encoded dtype")
184            .to_signed()
185    }
186}
187
188impl<T: TypedArrayRef<ZigZag>> ZigZagArrayExt for T {}
189
190#[derive(Clone, Debug)]
191pub struct ZigZag;
192
193impl ZigZag {
194    /// Construct a new [`ZigZagArray`] from an encoded unsigned integer array.
195    pub fn try_new(encoded: ArrayRef) -> VortexResult<ZigZagArray> {
196        let dtype = ZigZagData::dtype_from_encoded_dtype(encoded.dtype())?;
197        let len = encoded.len();
198        let slots = smallvec![Some(encoded.clone())];
199        let data = ZigZagData::try_new(encoded.dtype())?;
200        Ok(unsafe {
201            Array::from_parts_unchecked(ArrayParts::new(ZigZag, dtype, len, data).with_slots(slots))
202        })
203    }
204}
205
206impl ZigZagData {
207    fn dtype_from_encoded_dtype(encoded_dtype: &DType) -> VortexResult<DType> {
208        Ok(DType::from(PType::try_from(encoded_dtype)?.to_signed())
209            .with_nullability(encoded_dtype.nullability()))
210    }
211
212    pub fn new() -> Self {
213        Self {}
214    }
215
216    pub fn try_new(encoded_dtype: &DType) -> VortexResult<Self> {
217        if !encoded_dtype.is_unsigned_int() {
218            vortex_bail!(MismatchedTypes: "unsigned int", encoded_dtype);
219        }
220
221        Self::dtype_from_encoded_dtype(encoded_dtype)?;
222
223        Ok(Self {})
224    }
225}
226
227impl Default for ZigZagData {
228    fn default() -> Self {
229        Self::new()
230    }
231}
232
233impl OperationsVTable<ZigZag> for ZigZag {
234    fn scalar_at(
235        array: ArrayView<'_, ZigZag>,
236        index: usize,
237        ctx: &mut ExecutionCtx,
238    ) -> VortexResult<Scalar> {
239        let scalar = array.encoded().execute_scalar(index, ctx)?;
240        if scalar.is_null() {
241            return scalar.primitive_reinterpret_cast(ZigZagArrayExt::ptype(&array));
242        }
243
244        let pscalar = scalar.as_primitive();
245        Ok(match_each_unsigned_integer_ptype!(pscalar.ptype(), |P| {
246            Scalar::primitive(
247                <<P as ZigZagEncoded>::Int>::decode(
248                    pscalar
249                        .typed_value::<P>()
250                        .vortex_expect("zigzag corruption"),
251                ),
252                array.dtype().nullability(),
253            )
254        }))
255    }
256}
257
258impl ValidityChild<ZigZag> for ZigZag {
259    fn validity_child(array: ArrayView<'_, ZigZag>) -> ArrayRef {
260        array.encoded().clone()
261    }
262}
263
264#[cfg(test)]
265mod test {
266    use vortex_array::IntoArray;
267    use vortex_array::VortexSessionExecute;
268    use vortex_array::array_session;
269    use vortex_array::arrays::PrimitiveArray;
270    use vortex_array::scalar::Scalar;
271    use vortex_buffer::buffer;
272
273    use super::*;
274    use crate::zigzag_encode;
275
276    #[test]
277    fn test_compute_statistics() -> VortexResult<()> {
278        let mut ctx = array_session().create_execution_ctx();
279        let array = buffer![1i32, -5i32, 2, 3, 4, 5, 6, 7, 8, 9, 10]
280            .into_array()
281            .execute::<PrimitiveArray>(&mut ctx)?;
282        let zigzag = zigzag_encode(array.as_view())?;
283
284        assert_eq!(
285            zigzag.statistics().compute_max::<i32>(&mut ctx),
286            array.statistics().compute_max::<i32>(&mut ctx)
287        );
288        assert_eq!(
289            zigzag.statistics().compute_null_count(&mut ctx),
290            array.statistics().compute_null_count(&mut ctx)
291        );
292        assert_eq!(
293            zigzag.statistics().compute_is_constant(&mut ctx),
294            array.statistics().compute_is_constant(&mut ctx)
295        );
296
297        let sliced = zigzag.slice(0..2)?;
298        let sliced = sliced.as_::<ZigZag>();
299        assert_eq!(
300            sliced.array().execute_scalar(sliced.len() - 1, &mut ctx,)?,
301            Scalar::from(-5i32)
302        );
303
304        assert_eq!(
305            sliced.statistics().compute_min::<i32>(&mut ctx),
306            array.statistics().compute_min::<i32>(&mut ctx)
307        );
308        assert_eq!(
309            sliced.statistics().compute_null_count(&mut ctx),
310            array.statistics().compute_null_count(&mut ctx)
311        );
312        assert_eq!(
313            sliced.statistics().compute_is_constant(&mut ctx),
314            array.statistics().compute_is_constant(&mut ctx)
315        );
316        Ok(())
317    }
318}
vortex_zigzag/array.rs

vortex_zigzag/
array.rs
