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

vortex_array/builders/
bool.rs
