vortex_array/arrays/bool/
serde.rs

1use arrow_buffer::BooleanBuffer;
2use vortex_buffer::{Alignment, ByteBuffer};
3use vortex_dtype::DType;
4use vortex_error::{VortexExpect, VortexResult, vortex_bail};
5
6use super::{BoolArray, BoolEncoding};
7use crate::serde::ArrayParts;
8use crate::validity::Validity;
9use crate::vtable::EncodingVTable;
10use crate::{
11    Array, ArrayBufferVisitor, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl,
12    DeserializeMetadata, Encoding, EncodingId, ProstMetadata,
13};
14
15#[derive(prost::Message)]
16pub struct BoolMetadata {
17    // The offset in bits must be <8
18    #[prost(uint32, tag = "1")]
19    pub offset: u32,
20}
21
22impl EncodingVTable for BoolEncoding {
23    fn id(&self) -> EncodingId {
24        EncodingId::new_ref("vortex.bool")
25    }
26
27    fn decode(
28        &self,
29        parts: &ArrayParts,
30        ctx: &ArrayContext,
31        dtype: DType,
32        len: usize,
33    ) -> VortexResult<ArrayRef> {
34        let metadata = <Self as Encoding>::Metadata::deserialize(parts.metadata())?;
35
36        if parts.nbuffers() != 1 {
37            vortex_bail!("Expected 1 buffer, got {}", parts.nbuffers());
38        }
39        let buffer = BooleanBuffer::new(
40            parts.buffer(0)?.into_arrow_buffer(),
41            metadata.offset as usize,
42            len,
43        );
44
45        let validity = if parts.nchildren() == 0 {
46            Validity::from(dtype.nullability())
47        } else if parts.nchildren() == 1 {
48            let validity = parts.child(0).decode(ctx, Validity::DTYPE, len)?;
49            Validity::Array(validity)
50        } else {
51            vortex_bail!("Expected 0 or 1 child, got {}", parts.nchildren());
52        };
53
54        Ok(BoolArray::new(buffer, validity).into_array())
55    }
56}
57
58impl ArrayVisitorImpl<ProstMetadata<BoolMetadata>> for BoolArray {
59    fn _visit_buffers(&self, visitor: &mut dyn ArrayBufferVisitor) {
60        visitor.visit_buffer(&ByteBuffer::from_arrow_buffer(
61            self.boolean_buffer().clone().into_inner(),
62            Alignment::none(),
63        ))
64    }
65
66    fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
67        visitor.visit_validity(&self.validity, self.len());
68    }
69
70    fn _metadata(&self) -> ProstMetadata<BoolMetadata> {
71        let bit_offset = self.boolean_buffer().offset();
72        assert!(bit_offset < 8, "Offset must be <8, got {}", bit_offset);
73        ProstMetadata(BoolMetadata {
74            offset: u32::try_from(bit_offset).vortex_expect("checked"),
75        })
76    }
77}