vortex_zigzag/
serde.rs

1use vortex_array::serde::ArrayParts;
2use vortex_array::variants::PrimitiveArrayTrait;
3use vortex_array::vtable::EncodingVTable;
4use vortex_array::{
5    Array, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl, Canonical, EmptyMetadata,
6    EncodingId,
7};
8use vortex_dtype::{DType, PType};
9use vortex_error::{VortexResult, vortex_bail};
10
11use crate::{ZigZagArray, ZigZagEncoding, zigzag_encode};
12
13impl EncodingVTable for ZigZagEncoding {
14    fn id(&self) -> EncodingId {
15        EncodingId::new_ref("vortex.zigzag")
16    }
17
18    fn decode(
19        &self,
20        parts: &ArrayParts,
21        ctx: &ArrayContext,
22        dtype: DType,
23        len: usize,
24    ) -> VortexResult<ArrayRef> {
25        if parts.nchildren() != 1 {
26            vortex_bail!("Expected 1 child, got {}", parts.nchildren());
27        }
28
29        let ptype = PType::try_from(&dtype)?;
30        let encoded_type = DType::Primitive(ptype.to_unsigned(), dtype.nullability());
31
32        let encoded = parts.child(0).decode(ctx, encoded_type, len)?;
33        Ok(ZigZagArray::try_new(encoded)?.into_array())
34    }
35
36    fn encode(
37        &self,
38        input: &Canonical,
39        _like: Option<&dyn Array>,
40    ) -> VortexResult<Option<ArrayRef>> {
41        let parray = input.clone().into_primitive()?;
42
43        if !parray.ptype().is_signed_int() {
44            vortex_bail!(
45                "only signed integers can be encoded into {}, got {}",
46                self.id(),
47                parray.ptype()
48            )
49        }
50
51        Ok(Some(zigzag_encode(parray)?.into_array()))
52    }
53}
54
55impl ArrayVisitorImpl for ZigZagArray {
56    fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
57        visitor.visit_child("encoded", self.encoded())
58    }
59
60    fn _metadata(&self) -> EmptyMetadata {
61        EmptyMetadata
62    }
63}