polars_arrow/array/boolean/
builder.rs

1use polars_utils::IdxSize;
2
3use super::BooleanArray;
4use crate::array::builder::{ShareStrategy, StaticArrayBuilder};
5use crate::bitmap::{BitmapBuilder, OptBitmapBuilder};
6use crate::datatypes::ArrowDataType;
7
8pub struct BooleanArrayBuilder {
9    dtype: ArrowDataType,
10    values: BitmapBuilder,
11    validity: OptBitmapBuilder,
12}
13
14impl BooleanArrayBuilder {
15    pub fn new(dtype: ArrowDataType) -> Self {
16        Self {
17            dtype,
18            values: BitmapBuilder::new(),
19            validity: OptBitmapBuilder::default(),
20        }
21    }
22}
23
24impl StaticArrayBuilder for BooleanArrayBuilder {
25    type Array = BooleanArray;
26
27    fn dtype(&self) -> &ArrowDataType {
28        &self.dtype
29    }
30
31    fn reserve(&mut self, additional: usize) {
32        self.values.reserve(additional);
33        self.validity.reserve(additional);
34    }
35
36    fn freeze(self) -> BooleanArray {
37        let values = self.values.freeze();
38        let validity = self.validity.into_opt_validity();
39        BooleanArray::try_new(self.dtype, values, validity).unwrap()
40    }
41
42    fn freeze_reset(&mut self) -> Self::Array {
43        let values = core::mem::take(&mut self.values).freeze();
44        let validity = core::mem::take(&mut self.validity).into_opt_validity();
45        BooleanArray::try_new(self.dtype.clone(), values, validity).unwrap()
46    }
47
48    fn len(&self) -> usize {
49        self.values.len()
50    }
51
52    fn extend_nulls(&mut self, length: usize) {
53        self.values.extend_constant(length, false);
54        self.validity.extend_constant(length, false);
55    }
56
57    fn subslice_extend(
58        &mut self,
59        other: &BooleanArray,
60        start: usize,
61        length: usize,
62        _share: ShareStrategy,
63    ) {
64        self.values
65            .subslice_extend_from_bitmap(other.values(), start, length);
66        self.validity
67            .subslice_extend_from_opt_validity(other.validity(), start, length);
68    }
69
70    fn subslice_extend_each_repeated(
71        &mut self,
72        other: &BooleanArray,
73        start: usize,
74        length: usize,
75        repeats: usize,
76        _share: ShareStrategy,
77    ) {
78        self.values.subslice_extend_each_repeated_from_bitmap(
79            other.values(),
80            start,
81            length,
82            repeats,
83        );
84        self.validity
85            .subslice_extend_each_repeated_from_opt_validity(
86                other.validity(),
87                start,
88                length,
89                repeats,
90            );
91    }
92
93    unsafe fn gather_extend(
94        &mut self,
95        other: &BooleanArray,
96        idxs: &[IdxSize],
97        _share: ShareStrategy,
98    ) {
99        self.values.reserve(idxs.len());
100        for idx in idxs {
101            self.values
102                .push_unchecked(other.value_unchecked(*idx as usize));
103        }
104        self.validity
105            .gather_extend_from_opt_validity(other.validity(), idxs);
106    }
107
108    fn opt_gather_extend(&mut self, other: &BooleanArray, idxs: &[IdxSize], _share: ShareStrategy) {
109        self.values.reserve(idxs.len());
110        unsafe {
111            for idx in idxs {
112                let val = if (*idx as usize) < other.len() {
113                    // We don't use get here as that double-checks the validity
114                    // which we don't care about here.
115                    other.value_unchecked(*idx as usize)
116                } else {
117                    false
118                };
119                self.values.push_unchecked(val);
120            }
121        }
122        self.validity
123            .opt_gather_extend_from_opt_validity(other.validity(), idxs, other.len());
124    }
125}