vortex_array/arrays/primitive/
serde.rs

1use vortex_buffer::{Alignment, Buffer};
2use vortex_dtype::{DType, PType, match_each_native_ptype};
3use vortex_error::{VortexResult, vortex_bail};
4
5use super::PrimitiveEncoding;
6use crate::arrays::PrimitiveArray;
7use crate::serde::ArrayParts;
8use crate::validity::Validity;
9use crate::vtable::EncodingVTable;
10use crate::{
11    Array, ArrayBufferVisitor, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl,
12    Canonical, EmptyMetadata, EncodingId,
13};
14
15impl EncodingVTable for PrimitiveEncoding {
16    fn id(&self) -> EncodingId {
17        EncodingId::new_ref("vortex.primitive")
18    }
19
20    fn decode(
21        &self,
22        parts: &ArrayParts,
23        ctx: &ArrayContext,
24        dtype: DType,
25        len: usize,
26    ) -> VortexResult<ArrayRef> {
27        if parts.nbuffers() != 1 {
28            vortex_bail!("Expected 1 buffer, got {}", parts.nbuffers());
29        }
30        let buffer = parts.buffer(0)?;
31
32        let validity = if parts.nchildren() == 0 {
33            Validity::from(dtype.nullability())
34        } else if parts.nchildren() == 1 {
35            let validity = parts.child(0).decode(ctx, Validity::DTYPE, len)?;
36            Validity::Array(validity)
37        } else {
38            vortex_bail!("Expected 0 or 1 child, got {}", parts.nchildren());
39        };
40
41        let ptype = PType::try_from(&dtype)?;
42
43        if !buffer.is_aligned(Alignment::new(ptype.byte_width())) {
44            vortex_bail!(
45                "Buffer is not aligned to {}-byte boundary",
46                ptype.byte_width()
47            );
48        }
49        if buffer.len() != ptype.byte_width() * len {
50            vortex_bail!(
51                "Buffer length {} does not match expected length {} for {}, {} in {:?}",
52                buffer.len(),
53                ptype.byte_width() * len,
54                ptype.byte_width(),
55                len,
56                parts,
57            );
58        }
59
60        match_each_native_ptype!(ptype, |$P| {
61            let buffer = Buffer::<$P>::from_byte_buffer(buffer);
62            Ok(PrimitiveArray::new(buffer, validity).into_array())
63        })
64    }
65
66    fn encode(
67        &self,
68        input: &Canonical,
69        _like: Option<&dyn Array>,
70    ) -> VortexResult<Option<ArrayRef>> {
71        Ok(Some(input.clone().into_primitive()?.into_array()))
72    }
73}
74
75impl ArrayVisitorImpl for PrimitiveArray {
76    fn _visit_buffers(&self, visitor: &mut dyn ArrayBufferVisitor) {
77        visitor.visit_buffer(self.byte_buffer());
78    }
79
80    fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
81        visitor.visit_validity(self.validity(), self.len());
82    }
83
84    fn _metadata(&self) -> EmptyMetadata {
85        EmptyMetadata
86    }
87}