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}