arrow_array/builder/
boolean_builder.rs

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

arrow_array/builder/
boolean_builder.rs
