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_eq;
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_eq!(
129            scalar.dtype(),
130            self.dtype(),
131            "BoolBuilder received a scalar with the wrong dtype"
132        );
133
134        match scalar.as_bool().value() {
135            Some(value) => self.append_value(value),
136            None => self.append_null(),
137        }
138
139        Ok(())
140    }
141
142    fn reserve_exact(&mut self, additional: usize) {
143        self.inner.reserve(additional);
144        self.nulls.reserve_exact(additional);
145    }
146
147    fn finish(&mut self) -> ArrayRef {
148        self.finish_into_bool().into_array()
149    }
150
151    fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
152        Canonical::Bool(self.finish_into_bool())
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use rand::RngExt;
159    use rand::SeedableRng;
160    use rand::prelude::StdRng;
161    use vortex_buffer::BufferAllocatorRef;
162    use vortex_error::VortexResult;
163
164    use crate::ArrayRef;
165    use crate::IntoArray;
166    use crate::VortexSessionExecute;
167    use crate::array_session;
168    use crate::arrays::ChunkedArray;
169    use crate::arrays::bool::BoolArrayExt;
170    use crate::assert_arrays_eq;
171    use crate::builders::ArrayBuilder;
172    use crate::builders::BoolBuilder;
173    use crate::builders::bool::BoolArray;
174    use crate::builders::builder_with_capacity_in;
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_in(
204            chunk.dtype(),
205            len * chunk_count,
206            BufferAllocatorRef::static_ref(),
207        );
208        chunk
209            .clone()
210            .append_to_builder(builder.as_mut(), &mut ctx)?;
211
212        let canon_into = builder.finish().execute::<BoolArray>(&mut ctx)?;
213        let into_canon = chunk.clone().execute::<BoolArray>(&mut ctx)?;
214
215        assert!(canon_into.validity()?.mask_eq(
216            &into_canon.validity()?,
217            canon_into.len(),
218            &mut ctx
219        )?);
220        assert_eq!(canon_into.to_bit_buffer(), into_canon.to_bit_buffer());
221        Ok(())
222    }
223
224    #[test]
225    fn test_append_scalar() {
226        let mut ctx = array_session().create_execution_ctx();
227        let mut builder = BoolBuilder::with_capacity_in(
228            Nullability::Nullable,
229            10,
230            BufferAllocatorRef::static_ref(),
231        );
232
233        // Test appending true value.
234        let true_scalar = Scalar::bool(true, Nullability::Nullable);
235        builder.append_scalar(&true_scalar).unwrap();
236
237        // Test appending false value.
238        let false_scalar = Scalar::bool(false, Nullability::Nullable);
239        builder.append_scalar(&false_scalar).unwrap();
240
241        // Test appending null value.
242        let null_scalar = Scalar::null(DType::Bool(Nullability::Nullable));
243        builder.append_scalar(&null_scalar).unwrap();
244
245        let array = builder.finish_into_bool();
246        let expected = BoolArray::from_iter([Some(true), Some(false), None]);
247        assert_arrays_eq!(&array, &expected, &mut ctx);
248
249        // Test wrong dtype error.
250        let mut builder = BoolBuilder::with_capacity_in(
251            Nullability::NonNullable,
252            10,
253            BufferAllocatorRef::static_ref(),
254        );
255        let wrong_scalar = Scalar::from(42i32);
256        assert!(builder.append_scalar(&wrong_scalar).is_err());
257    }
258}
vortex_array/builders/bool.rs

vortex_array/builders/
bool.rs
