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

vortex_array/builders/
bool.rs
