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

vortex_array/builders/
bool.rs
