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_ensure};
9use vortex_mask::Mask;
10use vortex_scalar::{BoolScalar, Scalar};
11
12use crate::arrays::BoolArray;
13use crate::builders::{ArrayBuilder, DEFAULT_BUILDER_CAPACITY, LazyNullBufferBuilder};
14use crate::canonical::{Canonical, ToCanonical};
15use crate::{Array, ArrayRef, IntoArray};
16
17pub struct BoolBuilder {
18    dtype: DType,
19    inner: BooleanBufferBuilder,
20    nulls: LazyNullBufferBuilder,
21}
22
23impl BoolBuilder {
24    pub fn new(nullability: Nullability) -> Self {
25        Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY)
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            dtype: DType::Bool(nullability),
33        }
34    }
35
36    /// Appends a boolean value to the builder.
37    pub fn append_value(&mut self, value: bool) {
38        self.append_values(value, 1)
39    }
40
41    /// Appends the same boolean value multiple times to the builder.
42    ///
43    /// This method appends the given boolean value `n` times.
44    pub fn append_values(&mut self, value: bool, n: usize) {
45        self.inner.append_n(n, value);
46        self.nulls.append_n_non_nulls(n)
47    }
48
49    /// Finishes the builder directly into a [`BoolArray`].
50    pub fn finish_into_bool(&mut self) -> BoolArray {
51        assert_eq!(
52            self.nulls.len(),
53            self.inner.len(),
54            "Null count and value count should match when calling BoolBuilder::finish."
55        );
56
57        BoolArray::from_bool_buffer(
58            self.inner.finish(),
59            self.nulls.finish_with_nullability(self.dtype.nullability()),
60        )
61    }
62}
63
64impl ArrayBuilder for BoolBuilder {
65    fn as_any(&self) -> &dyn Any {
66        self
67    }
68
69    fn as_any_mut(&mut self) -> &mut dyn Any {
70        self
71    }
72
73    fn dtype(&self) -> &DType {
74        &self.dtype
75    }
76
77    fn len(&self) -> usize {
78        self.inner.len()
79    }
80
81    fn append_zeros(&mut self, n: usize) {
82        self.append_values(false, n)
83    }
84
85    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
86        self.inner.append_n(n, false);
87        self.nulls.append_n_nulls(n)
88    }
89
90    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
91        vortex_ensure!(
92            scalar.dtype() == self.dtype(),
93            "BoolBuilder expected scalar with dtype {:?}, got {:?}",
94            self.dtype(),
95            scalar.dtype()
96        );
97
98        let bool_scalar = BoolScalar::try_from(scalar)?;
99        match bool_scalar.value() {
100            Some(value) => self.append_value(value),
101            None => self.append_null(),
102        }
103
104        Ok(())
105    }
106
107    unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) {
108        let bool_array = array.to_bool();
109
110        self.inner.append_buffer(bool_array.boolean_buffer());
111        self.nulls.append_validity_mask(bool_array.validity_mask());
112    }
113
114    fn ensure_capacity(&mut self, capacity: usize) {
115        if capacity > self.inner.capacity() {
116            self.inner.reserve(capacity - self.inner.capacity());
117            self.nulls.ensure_capacity(capacity);
118        }
119    }
120
121    fn set_validity(&mut self, validity: Mask) {
122        self.nulls = LazyNullBufferBuilder::new(validity.len());
123        self.nulls.append_validity_mask(validity);
124    }
125
126    fn finish(&mut self) -> ArrayRef {
127        self.finish_into_bool().into_array()
128    }
129
130    fn finish_into_canonical(&mut self) -> Canonical {
131        Canonical::Bool(self.finish_into_bool())
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use rand::prelude::StdRng;
138    use rand::{Rng, SeedableRng};
139    use vortex_dtype::{DType, Nullability};
140    use vortex_scalar::Scalar;
141
142    use crate::array::Array;
143    use crate::arrays::{BoolArray, ChunkedArray};
144    use crate::builders::{ArrayBuilder, BoolBuilder, builder_with_capacity};
145    use crate::canonical::ToCanonical;
146    use crate::vtable::ValidityHelper;
147    use crate::{ArrayRef, IntoArray};
148
149    fn make_opt_bool_chunks(len: usize, chunk_count: usize) -> ArrayRef {
150        let mut rng = StdRng::seed_from_u64(0);
151
152        (0..chunk_count)
153            .map(|_| {
154                BoolArray::from_iter((0..len).map(|_| match rng.random_range::<u8, _>(0..=2) {
155                    0 => Some(false),
156                    1 => Some(true),
157                    2 => None,
158                    _ => unreachable!(),
159                }))
160                .into_array()
161            })
162            .collect::<ChunkedArray>()
163            .into_array()
164    }
165
166    #[test]
167    fn tests() {
168        let len = 1000;
169        let chunk_count = 10;
170        let chunk = make_opt_bool_chunks(len, chunk_count);
171
172        let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
173        chunk.clone().append_to_builder(builder.as_mut());
174
175        let canon_into = builder.finish().to_bool();
176        let into_canon = chunk.to_bool();
177
178        assert_eq!(canon_into.validity(), into_canon.validity());
179        assert_eq!(canon_into.boolean_buffer(), into_canon.boolean_buffer());
180    }
181
182    #[test]
183    fn test_append_scalar() {
184        let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10);
185
186        // Test appending true value.
187        let true_scalar = Scalar::bool(true, Nullability::Nullable);
188        builder.append_scalar(&true_scalar).unwrap();
189
190        // Test appending false value.
191        let false_scalar = Scalar::bool(false, Nullability::Nullable);
192        builder.append_scalar(&false_scalar).unwrap();
193
194        // Test appending null value.
195        let null_scalar = Scalar::null(DType::Bool(Nullability::Nullable));
196        builder.append_scalar(&null_scalar).unwrap();
197
198        let array = builder.finish_into_bool();
199        assert_eq!(array.len(), 3);
200
201        // Check actual values.
202        assert!(array.boolean_buffer().value(0));
203        assert!(!array.boolean_buffer().value(1));
204        // The third value is null, but the buffer might have any value.
205
206        // Check validity - first two should be valid, third should be null.
207        assert!(array.validity().is_valid(0));
208        assert!(array.validity().is_valid(1));
209        assert!(!array.validity().is_valid(2));
210
211        // Test wrong dtype error.
212        let mut builder = BoolBuilder::with_capacity(Nullability::NonNullable, 10);
213        let wrong_scalar = Scalar::from(42i32);
214        assert!(builder.append_scalar(&wrong_scalar).is_err());
215    }
216}
vortex_array/builders/bool.rs

vortex_array/builders/
bool.rs
