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

vortex_array/builders/
bool.rs
