vortex_array/arrays/bool/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BitBuffer;
5use vortex_dtype::DType;
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_ensure;
10use vortex_vector::bool::BoolVector;
11
12use crate::ArrayRef;
13use crate::DeserializeMetadata;
14use crate::ProstMetadata;
15use crate::SerializeMetadata;
16use crate::arrays::BoolArray;
17use crate::buffer::BufferHandle;
18use crate::executor::ExecutionCtx;
19use crate::serde::ArrayChildren;
20use crate::validity::Validity;
21use crate::vtable;
22use crate::vtable::ArrayVTableExt;
23use crate::vtable::NotSupported;
24use crate::vtable::VTable;
25use crate::vtable::ValidityVTableFromValidityHelper;
26
27mod array;
28mod canonical;
29mod operations;
30pub mod rules;
31mod validity;
32mod visitor;
33
34pub use rules::BoolMaskedValidityRule;
35use vortex_vector::Vector;
36
37use crate::arrays::bool::vtable::rules::RULES;
38use crate::vtable::ArrayId;
39use crate::vtable::ArrayVTable;
40
41vtable!(Bool);
42
43#[derive(prost::Message)]
44pub struct BoolMetadata {
45    // The offset in bits must be <8
46    #[prost(uint32, tag = "1")]
47    pub offset: u32,
48}
49
50impl VTable for BoolVTable {
51    type Array = BoolArray;
52
53    type Metadata = ProstMetadata<BoolMetadata>;
54
55    type ArrayVTable = Self;
56    type CanonicalVTable = Self;
57    type OperationsVTable = Self;
58    type ValidityVTable = ValidityVTableFromValidityHelper;
59    type VisitorVTable = Self;
60    type ComputeVTable = NotSupported;
61    type EncodeVTable = NotSupported;
62
63    fn id(&self) -> ArrayId {
64        ArrayId::new_ref("vortex.bool")
65    }
66
67    fn encoding(_array: &Self::Array) -> ArrayVTable {
68        BoolVTable.as_vtable()
69    }
70
71    fn metadata(array: &BoolArray) -> VortexResult<Self::Metadata> {
72        let bit_offset = array.bit_buffer().offset();
73        assert!(bit_offset < 8, "Offset must be <8, got {bit_offset}");
74        Ok(ProstMetadata(BoolMetadata {
75            offset: u32::try_from(bit_offset).vortex_expect("checked"),
76        }))
77    }
78
79    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
80        Ok(Some(metadata.serialize()))
81    }
82
83    fn deserialize(bytes: &[u8]) -> VortexResult<Self::Metadata> {
84        let metadata = <Self::Metadata as DeserializeMetadata>::deserialize(bytes)?;
85        Ok(ProstMetadata(metadata))
86    }
87
88    fn build(
89        &self,
90        dtype: &DType,
91        len: usize,
92        metadata: &Self::Metadata,
93        buffers: &[BufferHandle],
94        children: &dyn ArrayChildren,
95    ) -> VortexResult<BoolArray> {
96        if buffers.len() != 1 {
97            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
98        }
99
100        let validity = if children.is_empty() {
101            Validity::from(dtype.nullability())
102        } else if children.len() == 1 {
103            let validity = children.get(0, &Validity::DTYPE, len)?;
104            Validity::Array(validity)
105        } else {
106            vortex_bail!("Expected 0 or 1 child, got {}", children.len());
107        };
108
109        let buffer = buffers[0].clone().try_to_bytes()?;
110        let bits = BitBuffer::new_with_offset(buffer, len, metadata.offset as usize);
111
112        BoolArray::try_new(bits, validity)
113    }
114
115    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
116        vortex_ensure!(
117            children.len() <= 1,
118            "BoolArray can have at most 1 child (validity), got {}",
119            children.len()
120        );
121
122        array.validity = if children.is_empty() {
123            Validity::from(array.dtype().nullability())
124        } else {
125            Validity::Array(children.into_iter().next().vortex_expect("checked"))
126        };
127
128        Ok(())
129    }
130
131    fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
132        Ok(BoolVector::new(array.bit_buffer().clone(), array.validity_mask()).into())
133    }
134
135    fn reduce_parent(
136        array: &Self::Array,
137        parent: &ArrayRef,
138        child_idx: usize,
139    ) -> VortexResult<Option<ArrayRef>> {
140        RULES.evaluate(array, parent, child_idx)
141    }
142}
143
144#[derive(Debug)]
145pub struct BoolVTable;