vortex_array/builders/
bool.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use arrow_buffer::BooleanBufferBuilder;
7use vortex_dtype::{DType, Nullability};
8use vortex_error::{VortexResult, vortex_bail};
9use vortex_mask::Mask;
10
11use crate::arrays::BoolArray;
12use crate::builders::ArrayBuilder;
13use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
14use crate::{Array, ArrayRef, Canonical, IntoArray};
15
16pub struct BoolBuilder {
17    inner: BooleanBufferBuilder,
18    nulls: LazyNullBufferBuilder,
19    nullability: Nullability,
20    dtype: DType,
21}
22
23impl BoolBuilder {
24    pub fn new(nullability: Nullability) -> Self {
25        Self::with_capacity(nullability, 1024) // Same as Arrow builders
26    }
27
28    pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {
29        Self {
30            inner: BooleanBufferBuilder::new(capacity),
31            nulls: LazyNullBufferBuilder::new(capacity),
32            nullability,
33            dtype: DType::Bool(nullability),
34        }
35    }
36
37    pub fn append_value(&mut self, value: bool) {
38        self.append_values(value, 1)
39    }
40
41    pub fn append_values(&mut self, value: bool, n: usize) {
42        self.inner.append_n(n, value);
43        self.nulls.append_n_non_nulls(n)
44    }
45
46    pub fn append_option(&mut self, value: Option<bool>) {
47        match value {
48            Some(value) => self.append_value(value),
49            None => self.append_null(),
50        }
51    }
52}
53
54impl ArrayBuilder for BoolBuilder {
55    fn as_any(&self) -> &dyn Any {
56        self
57    }
58
59    fn as_any_mut(&mut self) -> &mut dyn Any {
60        self
61    }
62
63    fn dtype(&self) -> &DType {
64        &self.dtype
65    }
66
67    fn len(&self) -> usize {
68        self.inner.len()
69    }
70
71    fn append_zeros(&mut self, n: usize) {
72        self.append_values(false, n)
73    }
74
75    fn append_nulls(&mut self, n: usize) {
76        self.inner.append_n(n, false);
77        self.nulls.append_n_nulls(n)
78    }
79
80    fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {
81        let array = array.to_canonical()?;
82        let Canonical::Bool(array) = array else {
83            vortex_bail!("Expected Canonical::Bool, found {:?}", array);
84        };
85
86        self.inner.append_buffer(array.boolean_buffer());
87        self.nulls.append_validity_mask(array.validity_mask());
88
89        Ok(())
90    }
91
92    fn ensure_capacity(&mut self, capacity: usize) {
93        if capacity > self.inner.capacity() {
94            self.nulls.ensure_capacity(capacity);
95            self.inner.reserve(capacity - self.inner.capacity());
96        }
97    }
98
99    fn set_validity(&mut self, validity: Mask) {
100        self.nulls = LazyNullBufferBuilder::new(validity.len());
101        self.nulls.append_validity_mask(validity);
102    }
103
104    fn finish(&mut self) -> ArrayRef {
105        assert_eq!(
106            self.nulls.len(),
107            self.inner.len(),
108            "Null count and value count should match when calling BoolBuilder::finish."
109        );
110
111        BoolArray::new(
112            self.inner.finish(),
113            self.nulls.finish_with_nullability(self.nullability),
114        )
115        .into_array()
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use rand::prelude::StdRng;
122    use rand::{Rng, SeedableRng};
123
124    use crate::array::Array;
125    use crate::arrays::{BoolArray, ChunkedArray};
126    use crate::builders::builder_with_capacity;
127    use crate::canonical::ToCanonical;
128    use crate::vtable::ValidityHelper;
129    use crate::{ArrayRef, IntoArray};
130    fn make_opt_bool_chunks(len: usize, chunk_count: usize) -> ArrayRef {
131        let mut rng = StdRng::seed_from_u64(0);
132
133        (0..chunk_count)
134            .map(|_| {
135                BoolArray::from_iter((0..len).map(|_| match rng.random_range::<u8, _>(0..=2) {
136                    0 => Some(false),
137                    1 => Some(true),
138                    2 => None,
139                    _ => unreachable!(),
140                }))
141                .into_array()
142            })
143            .collect::<ChunkedArray>()
144            .into_array()
145    }
146
147    #[test]
148    fn tests() {
149        let len = 1000;
150        let chunk_count = 10;
151        let chunk = make_opt_bool_chunks(len, chunk_count);
152
153        let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
154        chunk.clone().append_to_builder(builder.as_mut()).unwrap();
155        let canon_into = builder
156            .finish()
157            .to_canonical()
158            .unwrap()
159            .into_bool()
160            .unwrap();
161
162        let into_canon = chunk.to_bool().unwrap();
163
164        assert_eq!(canon_into.validity(), into_canon.validity());
165        assert_eq!(canon_into.boolean_buffer(), into_canon.boolean_buffer());
166    }
167}
vortex_array/builders/bool.rs

vortex_array/builders/
bool.rs
