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

vortex_array/builders/
bool.rs
