vortex_array/builders/
bool.rs

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

vortex_array/builders/
bool.rs
