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

vortex_zigzag/
array.rs
