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        let array_data = unsafe { builder.build_unchecked() };
171        BooleanArray::from(array_data)
172    }
173
174    /// Builds the [BooleanArray] without resetting the builder.
175    pub fn finish_cloned(&self) -> BooleanArray {
176        let len = self.len();
177        let nulls = self.null_buffer_builder.finish_cloned();
178        let value_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
179        let builder = ArrayData::builder(DataType::Boolean)
180            .len(len)
181            .add_buffer(value_buffer)
182            .nulls(nulls);
183
184        let array_data = unsafe { builder.build_unchecked() };
185        BooleanArray::from(array_data)
186    }
187
188    /// Returns the current values buffer as a slice
189    ///
190    /// Boolean values are bit-packed into bytes. To extract the i-th boolean
191    /// from the bytes, you can use `arrow_buffer::bit_util::get_bit()`.
192    pub fn values_slice(&self) -> &[u8] {
193        self.values_builder.as_slice()
194    }
195
196    /// Returns the current null buffer as a slice
197    pub fn validity_slice(&self) -> Option<&[u8]> {
198        self.null_buffer_builder.as_slice()
199    }
200}
201
202impl ArrayBuilder for BooleanBuilder {
203    /// Returns the builder as a non-mutable `Any` reference.
204    fn as_any(&self) -> &dyn Any {
205        self
206    }
207
208    /// Returns the builder as a mutable `Any` reference.
209    fn as_any_mut(&mut self) -> &mut dyn Any {
210        self
211    }
212
213    /// Returns the boxed builder as a box of `Any`.
214    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
215        self
216    }
217
218    /// Returns the number of array slots in the builder
219    fn len(&self) -> usize {
220        self.values_builder.len()
221    }
222
223    /// Builds the array and reset this builder.
224    fn finish(&mut self) -> ArrayRef {
225        Arc::new(self.finish())
226    }
227
228    /// Builds the array without resetting the builder.
229    fn finish_cloned(&self) -> ArrayRef {
230        Arc::new(self.finish_cloned())
231    }
232}
233
234impl Extend<Option<bool>> for BooleanBuilder {
235    #[inline]
236    fn extend<T: IntoIterator<Item = Option<bool>>>(&mut self, iter: T) {
237        let buffered = iter.into_iter().collect::<Vec<_>>();
238        let array = unsafe {
239            // SAFETY: std::vec::IntoIter implements TrustedLen
240            BooleanArray::from_trusted_len_iter(buffered.into_iter())
241        };
242        self.append_array(&array)
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use crate::Array;
250    use arrow_buffer::{BooleanBuffer, NullBuffer};
251
252    #[test]
253    fn test_boolean_array_builder() {
254        // 00000010 01001000
255        let buf = Buffer::from([72_u8, 2_u8]);
256        let mut builder = BooleanArray::builder(10);
257        for i in 0..10 {
258            if i == 3 || i == 6 || i == 9 {
259                builder.append_value(true);
260            } else {
261                builder.append_value(false);
262            }
263        }
264
265        let arr = builder.finish();
266        assert_eq!(&buf, arr.values().inner());
267        assert_eq!(10, arr.len());
268        assert_eq!(0, arr.offset());
269        assert_eq!(0, arr.null_count());
270        for i in 0..10 {
271            assert!(!arr.is_null(i));
272            assert!(arr.is_valid(i));
273            assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {i}")
274        }
275    }
276
277    #[test]
278    fn test_boolean_array_builder_append_slice() {
279        let arr1 = BooleanArray::from(vec![Some(true), Some(false), None, None, Some(false)]);
280
281        let mut builder = BooleanArray::builder(0);
282        builder.append_slice(&[true, false]);
283        builder.append_null();
284        builder.append_null();
285        builder.append_value(false);
286        let arr2 = builder.finish();
287
288        assert_eq!(arr1, arr2);
289    }
290
291    #[test]
292    fn test_boolean_array_builder_append_slice_large() {
293        let arr1 = BooleanArray::from(vec![true; 513]);
294
295        let mut builder = BooleanArray::builder(512);
296        builder.append_slice(&[true; 513]);
297        let arr2 = builder.finish();
298
299        assert_eq!(arr1, arr2);
300    }
301
302    #[test]
303    fn test_boolean_array_builder_no_null() {
304        let mut builder = BooleanArray::builder(0);
305        builder.append_option(Some(true));
306        builder.append_value(false);
307        builder.append_slice(&[true, false, true]);
308        builder
309            .append_values(&[false, false, true], &[true, true, true])
310            .unwrap();
311
312        let array = builder.finish();
313        assert_eq!(0, array.null_count());
314        assert!(array.nulls().is_none());
315    }
316
317    #[test]
318    fn test_boolean_array_builder_finish_cloned() {
319        let mut builder = BooleanArray::builder(16);
320        builder.append_option(Some(true));
321        builder.append_value(false);
322        builder.append_slice(&[true, false, true]);
323        let mut array = builder.finish_cloned();
324        assert_eq!(3, array.true_count());
325        assert_eq!(2, array.false_count());
326
327        builder
328            .append_values(&[false, false, true], &[true, true, true])
329            .unwrap();
330
331        array = builder.finish();
332        assert_eq!(4, array.true_count());
333        assert_eq!(4, array.false_count());
334
335        assert_eq!(0, array.null_count());
336        assert!(array.nulls().is_none());
337    }
338
339    #[test]
340    fn test_extend() {
341        let mut builder = BooleanBuilder::new();
342        builder.extend([false, false, true, false, false].into_iter().map(Some));
343        builder.extend([true, true, false].into_iter().map(Some));
344        let array = builder.finish();
345        let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
346        assert_eq!(
347            &values,
348            &[false, false, true, false, false, true, true, false]
349        )
350    }
351
352    #[test]
353    fn test_boolean_array_builder_append_n() {
354        let mut builder = BooleanBuilder::new();
355        builder.append_n(3, true);
356        builder.append_n(2, false);
357        let array = builder.finish();
358        assert_eq!(3, array.true_count());
359        assert_eq!(2, array.false_count());
360        assert_eq!(0, array.null_count());
361
362        let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
363        assert_eq!(&values, &[true, true, true, false, false])
364    }
365
366    #[test]
367    fn test_append_array() {
368        let input = vec![
369            Some(true),
370            None,
371            Some(true),
372            None,
373            Some(false),
374            None,
375            None,
376            None,
377            Some(false),
378            Some(false),
379            Some(false),
380            Some(true),
381            Some(false),
382        ];
383        let arr1 = BooleanArray::from(input[..5].to_vec());
384        let arr2 = BooleanArray::from(input[5..8].to_vec());
385        let arr3 = BooleanArray::from(input[8..].to_vec());
386
387        let mut builder = BooleanBuilder::new();
388        builder.append_array(&arr1);
389        builder.append_array(&arr2);
390        builder.append_array(&arr3);
391        let actual = builder.finish();
392        let expected = BooleanArray::from(input);
393
394        assert_eq!(actual, expected);
395    }
396
397    #[test]
398    fn test_append_array_add_underlying_null_values() {
399        let array = BooleanArray::new(
400            BooleanBuffer::from(vec![true, false, true, false]),
401            Some(NullBuffer::from(&[true, true, false, false])),
402        );
403
404        let mut builder = BooleanBuilder::new();
405        builder.append_array(&array);
406        let actual = builder.finish();
407
408        assert_eq!(actual, array);
409        assert_eq!(actual.values(), array.values())
410    }
411}
arrow_array/builder/boolean_builder.rs

arrow_array/builder/
boolean_builder.rs
