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

vortex_array/builders/
bool.rs
