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_mask::Mask;
9
10use crate::arrays::BoolArray;
11use crate::builders::{ArrayBuilder, DEFAULT_BUILDER_CAPACITY, LazyNullBufferBuilder};
12use crate::canonical::{Canonical, ToCanonical};
13use crate::{Array, ArrayRef, IntoArray};
14
15pub struct BoolBuilder {
16    dtype: DType,
17    inner: BooleanBufferBuilder,
18    nulls: LazyNullBufferBuilder,
19}
20
21impl BoolBuilder {
22    pub fn new(nullability: Nullability) -> Self {
23        Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY)
24    }
25
26    pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {
27        Self {
28            inner: BooleanBufferBuilder::new(capacity),
29            nulls: LazyNullBufferBuilder::new(capacity),
30            dtype: DType::Bool(nullability),
31        }
32    }
33
34    /// Appends a boolean value to the builder.
35    pub fn append_value(&mut self, value: bool) {
36        self.append_values(value, 1)
37    }
38
39    /// Appends the same boolean value multiple times to the builder.
40    ///
41    /// This method appends the given boolean value `n` times.
42    pub fn append_values(&mut self, value: bool, n: usize) {
43        self.inner.append_n(n, value);
44        self.nulls.append_n_non_nulls(n)
45    }
46
47    /// Appends an optional boolean value to the builder.
48    ///
49    /// If the value is `Some`, it appends the boolean value. If the value is `None`, it appends a
50    /// null.
51    ///
52    /// # Panics
53    ///
54    /// This method will panic if the input is `None` and the builder is non-nullable.
55    pub fn append_option(&mut self, value: Option<bool>) {
56        match value {
57            Some(value) => self.append_value(value),
58            None => self.append_null(),
59        }
60    }
61
62    /// Finishes the builder directly into a [`BoolArray`].
63    pub fn finish_into_bool(&mut self) -> BoolArray {
64        assert_eq!(
65            self.nulls.len(),
66            self.inner.len(),
67            "Null count and value count should match when calling BoolBuilder::finish."
68        );
69
70        BoolArray::new(
71            self.inner.finish(),
72            self.nulls.finish_with_nullability(self.dtype.nullability()),
73        )
74    }
75}
76
77impl ArrayBuilder for BoolBuilder {
78    fn as_any(&self) -> &dyn Any {
79        self
80    }
81
82    fn as_any_mut(&mut self) -> &mut dyn Any {
83        self
84    }
85
86    fn dtype(&self) -> &DType {
87        &self.dtype
88    }
89
90    fn len(&self) -> usize {
91        self.inner.len()
92    }
93
94    fn append_zeros(&mut self, n: usize) {
95        self.append_values(false, n)
96    }
97
98    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
99        self.inner.append_n(n, false);
100        self.nulls.append_n_nulls(n)
101    }
102
103    unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) {
104        let bool_array = array.to_bool();
105
106        self.inner.append_buffer(bool_array.boolean_buffer());
107        self.nulls.append_validity_mask(bool_array.validity_mask());
108    }
109
110    fn ensure_capacity(&mut self, capacity: usize) {
111        if capacity > self.inner.capacity() {
112            self.inner.reserve(capacity - self.inner.capacity());
113            self.nulls.ensure_capacity(capacity);
114        }
115    }
116
117    fn set_validity(&mut self, validity: Mask) {
118        self.nulls = LazyNullBufferBuilder::new(validity.len());
119        self.nulls.append_validity_mask(validity);
120    }
121
122    fn finish(&mut self) -> ArrayRef {
123        self.finish_into_bool().into_array()
124    }
125
126    fn finish_into_canonical(&mut self) -> Canonical {
127        Canonical::Bool(self.finish_into_bool())
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use rand::prelude::StdRng;
134    use rand::{Rng, SeedableRng};
135
136    use crate::array::Array;
137    use crate::arrays::{BoolArray, ChunkedArray};
138    use crate::builders::builder_with_capacity;
139    use crate::canonical::ToCanonical;
140    use crate::vtable::ValidityHelper;
141    use crate::{ArrayRef, IntoArray};
142    fn make_opt_bool_chunks(len: usize, chunk_count: usize) -> ArrayRef {
143        let mut rng = StdRng::seed_from_u64(0);
144
145        (0..chunk_count)
146            .map(|_| {
147                BoolArray::from_iter((0..len).map(|_| match rng.random_range::<u8, _>(0..=2) {
148                    0 => Some(false),
149                    1 => Some(true),
150                    2 => None,
151                    _ => unreachable!(),
152                }))
153                .into_array()
154            })
155            .collect::<ChunkedArray>()
156            .into_array()
157    }
158
159    #[test]
160    fn tests() {
161        let len = 1000;
162        let chunk_count = 10;
163        let chunk = make_opt_bool_chunks(len, chunk_count);
164
165        let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
166        chunk.clone().append_to_builder(builder.as_mut());
167
168        let canon_into = builder.finish().to_bool();
169        let into_canon = chunk.to_bool();
170
171        assert_eq!(canon_into.validity(), into_canon.validity());
172        assert_eq!(canon_into.boolean_buffer(), into_canon.boolean_buffer());
173    }
174}
vortex_array/builders/bool.rs

vortex_array/builders/
bool.rs
