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

vortex_array/builders/
bool.rs
