vortex_array/builders/
bool.rs

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

vortex_array/builders/
bool.rs
