vortex_array/arrays/bool/
serde.rs

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