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, EncodingId, RkyvMetadata,
13};
14
15impl EncodingVTable for BoolEncoding {
16    fn id(&self) -> EncodingId {
17        EncodingId::new_ref("vortex.bool")
18    }
19
20    fn decode(
21        &self,
22        parts: &ArrayParts,
23        ctx: &ArrayContext,
24        dtype: DType,
25        len: usize,
26    ) -> VortexResult<ArrayRef> {
27        let metadata = RkyvMetadata::<BoolMetadata>::deserialize(parts.metadata())?;
28
29        if parts.nbuffers() != 1 {
30            vortex_bail!("Expected 1 buffer, got {}", parts.nbuffers());
31        }
32        let buffer = BooleanBuffer::new(
33            parts.buffer(0)?.into_arrow_buffer(),
34            metadata.offset as usize,
35            len,
36        );
37
38        let validity = if parts.nchildren() == 0 {
39            Validity::from(dtype.nullability())
40        } else if parts.nchildren() == 1 {
41            let validity = parts.child(0).decode(ctx, Validity::DTYPE, len)?;
42            Validity::Array(validity)
43        } else {
44            vortex_bail!("Expected 0 or 1 child, got {}", parts.nchildren());
45        };
46
47        Ok(BoolArray::new(buffer, validity).into_array())
48    }
49}
50
51#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
52pub struct BoolMetadata {
53    // We know the offset in bits must be <8
54    offset: u8,
55}
56
57impl ArrayVisitorImpl<RkyvMetadata<BoolMetadata>> for BoolArray {
58    fn _visit_buffers(&self, visitor: &mut dyn ArrayBufferVisitor) {
59        visitor.visit_buffer(&ByteBuffer::from_arrow_buffer(
60            self.boolean_buffer().clone().into_inner(),
61            Alignment::none(),
62        ))
63    }
64
65    fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
66        visitor.visit_validity(&self.validity, self.len());
67    }
68
69    fn _metadata(&self) -> RkyvMetadata<BoolMetadata> {
70        let bit_offset = self.boolean_buffer().offset();
71        assert!(bit_offset < 8, "Offset must be <8, got {}", bit_offset);
72        RkyvMetadata(BoolMetadata {
73            offset: u8::try_from(bit_offset).vortex_expect("checked"),
74        })
75    }
76}