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}