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;
7use crate::arrays::BoolVTable;
8use crate::serde::ArrayChildren;
9use crate::validity::Validity;
10use crate::vtable::{SerdeVTable, VTable, VisitorVTable};
11use crate::{ArrayBufferVisitor, ArrayChildVisitor, ProstMetadata};
12
13#[derive(prost::Message)]
14pub struct BoolMetadata {
15    // The offset in bits must be <8
16    #[prost(uint32, tag = "1")]
17    pub offset: u32,
18}
19
20impl SerdeVTable<BoolVTable> for BoolVTable {
21    type Metadata = ProstMetadata<BoolMetadata>;
22
23    fn metadata(array: &BoolArray) -> VortexResult<Option<Self::Metadata>> {
24        let bit_offset = array.boolean_buffer().offset();
25        assert!(bit_offset < 8, "Offset must be <8, got {bit_offset}");
26        Ok(Some(ProstMetadata(BoolMetadata {
27            offset: u32::try_from(bit_offset).vortex_expect("checked"),
28        })))
29    }
30
31    fn build(
32        _encoding: &<BoolVTable as VTable>::Encoding,
33        dtype: &DType,
34        len: usize,
35        metadata: &BoolMetadata,
36        buffers: &[ByteBuffer],
37        children: &dyn ArrayChildren,
38    ) -> VortexResult<BoolArray> {
39        if buffers.len() != 1 {
40            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
41        }
42        let buffer = BooleanBuffer::new(
43            buffers[0].clone().into_arrow_buffer(),
44            metadata.offset as usize,
45            len,
46        );
47
48        let validity = if children.is_empty() {
49            Validity::from(dtype.nullability())
50        } else if children.len() == 1 {
51            let validity = children.get(0, &Validity::DTYPE, len)?;
52            Validity::Array(validity)
53        } else {
54            vortex_bail!("Expected 0 or 1 child, got {}", children.len());
55        };
56
57        Ok(BoolArray::new(buffer, validity))
58    }
59}
60
61impl VisitorVTable<BoolVTable> for BoolVTable {
62    fn visit_buffers(array: &BoolArray, visitor: &mut dyn ArrayBufferVisitor) {
63        visitor.visit_buffer(&ByteBuffer::from_arrow_buffer(
64            array.boolean_buffer().clone().into_inner(),
65            Alignment::none(),
66        ))
67    }
68
69    fn visit_children(array: &BoolArray, visitor: &mut dyn ArrayChildVisitor) {
70        visitor.visit_validity(&array.validity, array.len());
71    }
72}