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

vortex_array/builders/
bool.rs
