vortex_array/builders/
bool.rs

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

vortex_array/builders/
bool.rs
