arrow_array/builder/
boolean_builder.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::builder::{ArrayBuilder, BooleanBufferBuilder};
19use crate::{Array, ArrayRef, BooleanArray};
20use arrow_buffer::Buffer;
21use arrow_buffer::NullBufferBuilder;
22use arrow_data::ArrayData;
23use arrow_schema::{ArrowError, DataType};
24use std::any::Any;
25use std::sync::Arc;
26
27/// Builder for [`BooleanArray`]
28///
29/// # Example
30///
31/// Create a `BooleanArray` from a `BooleanBuilder`
32///
33/// ```
34///
35/// # use arrow_array::{Array, BooleanArray, builder::BooleanBuilder};
36///
37/// let mut b = BooleanBuilder::new();
38/// b.append_value(true);
39/// b.append_null();
40/// b.append_value(false);
41/// b.append_value(true);
42/// let arr = b.finish();
43///
44/// assert_eq!(4, arr.len());
45/// assert_eq!(1, arr.null_count());
46/// assert_eq!(true, arr.value(0));
47/// assert!(arr.is_valid(0));
48/// assert!(!arr.is_null(0));
49/// assert!(!arr.is_valid(1));
50/// assert!(arr.is_null(1));
51/// assert_eq!(false, arr.value(2));
52/// assert!(arr.is_valid(2));
53/// assert!(!arr.is_null(2));
54/// assert_eq!(true, arr.value(3));
55/// assert!(arr.is_valid(3));
56/// assert!(!arr.is_null(3));
57/// ```
58#[derive(Debug)]
59pub struct BooleanBuilder {
60    values_builder: BooleanBufferBuilder,
61    null_buffer_builder: NullBufferBuilder,
62}
63
64impl Default for BooleanBuilder {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70impl BooleanBuilder {
71    /// Creates a new boolean builder
72    pub fn new() -> Self {
73        Self::with_capacity(1024)
74    }
75
76    /// Creates a new boolean builder with space for `capacity` elements without re-allocating
77    pub fn with_capacity(capacity: usize) -> Self {
78        Self {
79            values_builder: BooleanBufferBuilder::new(capacity),
80            null_buffer_builder: NullBufferBuilder::new(capacity),
81        }
82    }
83
84    /// Returns the capacity of this builder measured in slots of type `T`
85    pub fn capacity(&self) -> usize {
86        self.values_builder.capacity()
87    }
88
89    /// Appends a value of type `T` into the builder
90    #[inline]
91    pub fn append_value(&mut self, v: bool) {
92        self.values_builder.append(v);
93        self.null_buffer_builder.append_non_null();
94    }
95
96    /// Appends a null slot into the builder
97    #[inline]
98    pub fn append_null(&mut self) {
99        self.null_buffer_builder.append_null();
100        self.values_builder.advance(1);
101    }
102
103    /// Appends `n` `null`s into the builder.
104    #[inline]
105    pub fn append_nulls(&mut self, n: usize) {
106        self.null_buffer_builder.append_n_nulls(n);
107        self.values_builder.advance(n);
108    }
109
110    /// Appends an `Option<T>` into the builder
111    #[inline]
112    pub fn append_option(&mut self, v: Option<bool>) {
113        match v {
114            None => self.append_null(),
115            Some(v) => self.append_value(v),
116        }
117    }
118
119    /// Appends a slice of type `T` into the builder
120    #[inline]
121    pub fn append_slice(&mut self, v: &[bool]) {
122        self.values_builder.append_slice(v);
123        self.null_buffer_builder.append_n_non_nulls(v.len());
124    }
125
126    /// Appends n `additional` bits of value `v` into the buffer
127    #[inline]
128    pub fn append_n(&mut self, additional: usize, v: bool) {
129        self.values_builder.append_n(additional, v);
130        self.null_buffer_builder.append_n_non_nulls(additional);
131    }
132
133    /// Appends values from a slice of type `T` and a validity boolean slice.
134    ///
135    /// Returns an error if the slices are of different lengths
136    #[inline]
137    pub fn append_values(&mut self, values: &[bool], is_valid: &[bool]) -> Result<(), ArrowError> {
138        if values.len() != is_valid.len() {
139            Err(ArrowError::InvalidArgumentError(
140                "Value and validity lengths must be equal".to_string(),
141            ))
142        } else {
143            self.null_buffer_builder.append_slice(is_valid);
144            self.values_builder.append_slice(values);
145            Ok(())
146        }
147    }
148
149    /// Appends array values and null to this builder as is
150    /// (this means that underlying null values are copied as is).
151    #[inline]
152    pub fn append_array(&mut self, array: &BooleanArray) {
153        self.values_builder.append_buffer(array.values());
154        if let Some(null_buffer) = array.nulls() {
155            self.null_buffer_builder.append_buffer(null_buffer);
156        } else {
157            self.null_buffer_builder.append_n_non_nulls(array.len());
158        }
159    }
160
161    /// Builds the [BooleanArray] and reset this builder.
162    pub fn finish(&mut self) -> BooleanArray {
163        let len = self.len();
164        let null_bit_buffer = self.null_buffer_builder.finish();
165        let builder = ArrayData::builder(DataType::Boolean)
166            .len(len)
167            .add_buffer(self.values_builder.finish().into_inner())
168            .nulls(null_bit_buffer);
169
170        // SAFETY: values buffer and nulls have matching lengths
171        let array_data = unsafe { builder.build_unchecked() };
172        BooleanArray::from(array_data)
173    }
174
175    /// Builds the [BooleanArray] without resetting the builder.
176    pub fn finish_cloned(&self) -> BooleanArray {
177        let len = self.len();
178        let nulls = self.null_buffer_builder.finish_cloned();
179        let value_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
180        let builder = ArrayData::builder(DataType::Boolean)
181            .len(len)
182            .add_buffer(value_buffer)
183            .nulls(nulls);
184
185        // SAFETY: values buffer and nulls have matching lengths
186        let array_data = unsafe { builder.build_unchecked() };
187        BooleanArray::from(array_data)
188    }
189
190    /// Returns the current values buffer as a slice
191    ///
192    /// Boolean values are bit-packed into bytes. To extract the i-th boolean
193    /// from the bytes, you can use `arrow_buffer::bit_util::get_bit()`.
194    pub fn values_slice(&self) -> &[u8] {
195        self.values_builder.as_slice()
196    }
197
198    /// Returns the current null buffer as a slice
199    pub fn validity_slice(&self) -> Option<&[u8]> {
200        self.null_buffer_builder.as_slice()
201    }
202}
203
204impl ArrayBuilder for BooleanBuilder {
205    /// Returns the builder as a non-mutable `Any` reference.
206    fn as_any(&self) -> &dyn Any {
207        self
208    }
209
210    /// Returns the builder as a mutable `Any` reference.
211    fn as_any_mut(&mut self) -> &mut dyn Any {
212        self
213    }
214
215    /// Returns the boxed builder as a box of `Any`.
216    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
217        self
218    }
219
220    /// Returns the number of array slots in the builder
221    fn len(&self) -> usize {
222        self.values_builder.len()
223    }
224
225    /// Builds the array and reset this builder.
226    fn finish(&mut self) -> ArrayRef {
227        Arc::new(self.finish())
228    }
229
230    /// Builds the array without resetting the builder.
231    fn finish_cloned(&self) -> ArrayRef {
232        Arc::new(self.finish_cloned())
233    }
234}
235
236impl Extend<Option<bool>> for BooleanBuilder {
237    #[inline]
238    fn extend<T: IntoIterator<Item = Option<bool>>>(&mut self, iter: T) {
239        let buffered = iter.into_iter().collect::<Vec<_>>();
240        let array = unsafe {
241            // SAFETY: std::vec::IntoIter implements TrustedLen
242            BooleanArray::from_trusted_len_iter(buffered.into_iter())
243        };
244        self.append_array(&array)
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251    use crate::Array;
252    use arrow_buffer::{BooleanBuffer, NullBuffer};
253
254    #[test]
255    fn test_boolean_array_builder() {
256        // 00000010 01001000
257        let buf = Buffer::from([72_u8, 2_u8]);
258        let mut builder = BooleanArray::builder(10);
259        for i in 0..10 {
260            if i == 3 || i == 6 || i == 9 {
261                builder.append_value(true);
262            } else {
263                builder.append_value(false);
264            }
265        }
266
267        let arr = builder.finish();
268        assert_eq!(&buf, arr.values().inner());
269        assert_eq!(10, arr.len());
270        assert_eq!(0, arr.offset());
271        assert_eq!(0, arr.null_count());
272        for i in 0..10 {
273            assert!(!arr.is_null(i));
274            assert!(arr.is_valid(i));
275            assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {i}")
276        }
277    }
278
279    #[test]
280    fn test_boolean_array_builder_append_slice() {
281        let arr1 = BooleanArray::from(vec![Some(true), Some(false), None, None, Some(false)]);
282
283        let mut builder = BooleanArray::builder(0);
284        builder.append_slice(&[true, false]);
285        builder.append_null();
286        builder.append_null();
287        builder.append_value(false);
288        let arr2 = builder.finish();
289
290        assert_eq!(arr1, arr2);
291    }
292
293    #[test]
294    fn test_boolean_array_builder_append_slice_large() {
295        let arr1 = BooleanArray::from(vec![true; 513]);
296
297        let mut builder = BooleanArray::builder(512);
298        builder.append_slice(&[true; 513]);
299        let arr2 = builder.finish();
300
301        assert_eq!(arr1, arr2);
302    }
303
304    #[test]
305    fn test_boolean_array_builder_no_null() {
306        let mut builder = BooleanArray::builder(0);
307        builder.append_option(Some(true));
308        builder.append_value(false);
309        builder.append_slice(&[true, false, true]);
310        builder
311            .append_values(&[false, false, true], &[true, true, true])
312            .unwrap();
313
314        let array = builder.finish();
315        assert_eq!(0, array.null_count());
316        assert!(array.nulls().is_none());
317    }
318
319    #[test]
320    fn test_boolean_array_builder_finish_cloned() {
321        let mut builder = BooleanArray::builder(16);
322        builder.append_option(Some(true));
323        builder.append_value(false);
324        builder.append_slice(&[true, false, true]);
325        let mut array = builder.finish_cloned();
326        assert_eq!(3, array.true_count());
327        assert_eq!(2, array.false_count());
328
329        builder
330            .append_values(&[false, false, true], &[true, true, true])
331            .unwrap();
332
333        array = builder.finish();
334        assert_eq!(4, array.true_count());
335        assert_eq!(4, array.false_count());
336
337        assert_eq!(0, array.null_count());
338        assert!(array.nulls().is_none());
339    }
340
341    #[test]
342    fn test_extend() {
343        let mut builder = BooleanBuilder::new();
344        builder.extend([false, false, true, false, false].into_iter().map(Some));
345        builder.extend([true, true, false].into_iter().map(Some));
346        let array = builder.finish();
347        let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
348        assert_eq!(
349            &values,
350            &[false, false, true, false, false, true, true, false]
351        )
352    }
353
354    #[test]
355    fn test_boolean_array_builder_append_n() {
356        let mut builder = BooleanBuilder::new();
357        builder.append_n(3, true);
358        builder.append_n(2, false);
359        let array = builder.finish();
360        assert_eq!(3, array.true_count());
361        assert_eq!(2, array.false_count());
362        assert_eq!(0, array.null_count());
363
364        let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
365        assert_eq!(&values, &[true, true, true, false, false])
366    }
367
368    #[test]
369    fn test_append_array() {
370        let input = vec![
371            Some(true),
372            None,
373            Some(true),
374            None,
375            Some(false),
376            None,
377            None,
378            None,
379            Some(false),
380            Some(false),
381            Some(false),
382            Some(true),
383            Some(false),
384        ];
385        let arr1 = BooleanArray::from(input[..5].to_vec());
386        let arr2 = BooleanArray::from(input[5..8].to_vec());
387        let arr3 = BooleanArray::from(input[8..].to_vec());
388
389        let mut builder = BooleanBuilder::new();
390        builder.append_array(&arr1);
391        builder.append_array(&arr2);
392        builder.append_array(&arr3);
393        let actual = builder.finish();
394        let expected = BooleanArray::from(input);
395
396        assert_eq!(actual, expected);
397    }
398
399    #[test]
400    fn test_append_array_add_underlying_null_values() {
401        let array = BooleanArray::new(
402            BooleanBuffer::from(vec![true, false, true, false]),
403            Some(NullBuffer::from(&[true, true, false, false])),
404        );
405
406        let mut builder = BooleanBuilder::new();
407        builder.append_array(&array);
408        let actual = builder.finish();
409
410        assert_eq!(actual, array);
411        assert_eq!(actual.values(), array.values())
412    }
413}
arrow_array/builder/boolean_builder.rs

arrow_array/builder/
boolean_builder.rs
