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::{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    /// Builds the [BooleanArray] and reset this builder.
150    pub fn finish(&mut self) -> BooleanArray {
151        let len = self.len();
152        let null_bit_buffer = self.null_buffer_builder.finish();
153        let builder = ArrayData::builder(DataType::Boolean)
154            .len(len)
155            .add_buffer(self.values_builder.finish().into_inner())
156            .nulls(null_bit_buffer);
157
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        let array_data = unsafe { builder.build_unchecked() };
173        BooleanArray::from(array_data)
174    }
175
176    /// Returns the current values buffer as a slice
177    ///
178    /// Boolean values are bit-packed into bytes. To extract the i-th boolean
179    /// from the bytes, you can use `arrow_buffer::bit_util::get_bit()`.
180    pub fn values_slice(&self) -> &[u8] {
181        self.values_builder.as_slice()
182    }
183
184    /// Returns the current null buffer as a slice
185    pub fn validity_slice(&self) -> Option<&[u8]> {
186        self.null_buffer_builder.as_slice()
187    }
188}
189
190impl ArrayBuilder for BooleanBuilder {
191    /// Returns the builder as a non-mutable `Any` reference.
192    fn as_any(&self) -> &dyn Any {
193        self
194    }
195
196    /// Returns the builder as a mutable `Any` reference.
197    fn as_any_mut(&mut self) -> &mut dyn Any {
198        self
199    }
200
201    /// Returns the boxed builder as a box of `Any`.
202    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
203        self
204    }
205
206    /// Returns the number of array slots in the builder
207    fn len(&self) -> usize {
208        self.values_builder.len()
209    }
210
211    /// Builds the array and reset this builder.
212    fn finish(&mut self) -> ArrayRef {
213        Arc::new(self.finish())
214    }
215
216    /// Builds the array without resetting the builder.
217    fn finish_cloned(&self) -> ArrayRef {
218        Arc::new(self.finish_cloned())
219    }
220}
221
222impl Extend<Option<bool>> for BooleanBuilder {
223    #[inline]
224    fn extend<T: IntoIterator<Item = Option<bool>>>(&mut self, iter: T) {
225        for v in iter {
226            self.append_option(v)
227        }
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234    use crate::Array;
235
236    #[test]
237    fn test_boolean_array_builder() {
238        // 00000010 01001000
239        let buf = Buffer::from([72_u8, 2_u8]);
240        let mut builder = BooleanArray::builder(10);
241        for i in 0..10 {
242            if i == 3 || i == 6 || i == 9 {
243                builder.append_value(true);
244            } else {
245                builder.append_value(false);
246            }
247        }
248
249        let arr = builder.finish();
250        assert_eq!(&buf, arr.values().inner());
251        assert_eq!(10, arr.len());
252        assert_eq!(0, arr.offset());
253        assert_eq!(0, arr.null_count());
254        for i in 0..10 {
255            assert!(!arr.is_null(i));
256            assert!(arr.is_valid(i));
257            assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {i}")
258        }
259    }
260
261    #[test]
262    fn test_boolean_array_builder_append_slice() {
263        let arr1 = BooleanArray::from(vec![Some(true), Some(false), None, None, Some(false)]);
264
265        let mut builder = BooleanArray::builder(0);
266        builder.append_slice(&[true, false]);
267        builder.append_null();
268        builder.append_null();
269        builder.append_value(false);
270        let arr2 = builder.finish();
271
272        assert_eq!(arr1, arr2);
273    }
274
275    #[test]
276    fn test_boolean_array_builder_append_slice_large() {
277        let arr1 = BooleanArray::from(vec![true; 513]);
278
279        let mut builder = BooleanArray::builder(512);
280        builder.append_slice(&[true; 513]);
281        let arr2 = builder.finish();
282
283        assert_eq!(arr1, arr2);
284    }
285
286    #[test]
287    fn test_boolean_array_builder_no_null() {
288        let mut builder = BooleanArray::builder(0);
289        builder.append_option(Some(true));
290        builder.append_value(false);
291        builder.append_slice(&[true, false, true]);
292        builder
293            .append_values(&[false, false, true], &[true, true, true])
294            .unwrap();
295
296        let array = builder.finish();
297        assert_eq!(0, array.null_count());
298        assert!(array.nulls().is_none());
299    }
300
301    #[test]
302    fn test_boolean_array_builder_finish_cloned() {
303        let mut builder = BooleanArray::builder(16);
304        builder.append_option(Some(true));
305        builder.append_value(false);
306        builder.append_slice(&[true, false, true]);
307        let mut array = builder.finish_cloned();
308        assert_eq!(3, array.true_count());
309        assert_eq!(2, array.false_count());
310
311        builder
312            .append_values(&[false, false, true], &[true, true, true])
313            .unwrap();
314
315        array = builder.finish();
316        assert_eq!(4, array.true_count());
317        assert_eq!(4, array.false_count());
318
319        assert_eq!(0, array.null_count());
320        assert!(array.nulls().is_none());
321    }
322
323    #[test]
324    fn test_extend() {
325        let mut builder = BooleanBuilder::new();
326        builder.extend([false, false, true, false, false].into_iter().map(Some));
327        builder.extend([true, true, false].into_iter().map(Some));
328        let array = builder.finish();
329        let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
330        assert_eq!(
331            &values,
332            &[false, false, true, false, false, true, true, false]
333        )
334    }
335
336    #[test]
337    fn test_boolean_array_builder_append_n() {
338        let mut builder = BooleanBuilder::new();
339        builder.append_n(3, true);
340        builder.append_n(2, false);
341        let array = builder.finish();
342        assert_eq!(3, array.true_count());
343        assert_eq!(2, array.false_count());
344        assert_eq!(0, array.null_count());
345
346        let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
347        assert_eq!(&values, &[true, true, true, false, false])
348    }
349}
arrow_array/builder/boolean_builder.rs

arrow_array/builder/
boolean_builder.rs
