polars_arrow/array/primitive/
builder.rs

1use polars_utils::IdxSize;
2use polars_utils::vec::PushUnchecked;
3
4use super::PrimitiveArray;
5use crate::array::builder::{ShareStrategy, StaticArrayBuilder};
6use crate::bitmap::OptBitmapBuilder;
7use crate::buffer::Buffer;
8use crate::datatypes::ArrowDataType;
9use crate::types::NativeType;
10
11pub struct PrimitiveArrayBuilder<T> {
12    dtype: ArrowDataType,
13    values: Vec<T>,
14    validity: OptBitmapBuilder,
15}
16
17impl<T: NativeType> PrimitiveArrayBuilder<T> {
18    pub fn new(dtype: ArrowDataType) -> Self {
19        Self {
20            dtype,
21            values: Vec::new(),
22            validity: OptBitmapBuilder::default(),
23        }
24    }
25}
26
27impl<T: NativeType> StaticArrayBuilder for PrimitiveArrayBuilder<T> {
28    type Array = PrimitiveArray<T>;
29
30    fn dtype(&self) -> &ArrowDataType {
31        &self.dtype
32    }
33
34    fn reserve(&mut self, additional: usize) {
35        self.values.reserve(additional);
36        self.validity.reserve(additional);
37    }
38
39    fn freeze(self) -> PrimitiveArray<T> {
40        let values = Buffer::from(self.values);
41        let validity = self.validity.into_opt_validity();
42        PrimitiveArray::new(self.dtype, values, validity)
43    }
44
45    fn freeze_reset(&mut self) -> Self::Array {
46        let values = Buffer::from(core::mem::take(&mut self.values));
47        let validity = core::mem::take(&mut self.validity).into_opt_validity();
48        PrimitiveArray::new(self.dtype.clone(), values, validity)
49    }
50
51    fn len(&self) -> usize {
52        self.values.len()
53    }
54
55    fn extend_nulls(&mut self, length: usize) {
56        self.values.resize(self.values.len() + length, T::zeroed());
57        self.validity.extend_constant(length, false);
58    }
59
60    fn subslice_extend(
61        &mut self,
62        other: &PrimitiveArray<T>,
63        start: usize,
64        length: usize,
65        _share: ShareStrategy,
66    ) {
67        self.values
68            .extend_from_slice(&other.values()[start..start + length]);
69        self.validity
70            .subslice_extend_from_opt_validity(other.validity(), start, length);
71    }
72
73    fn subslice_extend_each_repeated(
74        &mut self,
75        other: &PrimitiveArray<T>,
76        start: usize,
77        length: usize,
78        repeats: usize,
79        _share: ShareStrategy,
80    ) {
81        self.values.reserve(length * repeats);
82
83        for value in other.values()[start..start + length].iter() {
84            unsafe {
85                for _ in 0..repeats {
86                    self.values.push_unchecked(*value);
87                }
88            }
89        }
90
91        self.validity
92            .subslice_extend_each_repeated_from_opt_validity(
93                other.validity(),
94                start,
95                length,
96                repeats,
97            );
98    }
99
100    unsafe fn gather_extend(
101        &mut self,
102        other: &PrimitiveArray<T>,
103        idxs: &[IdxSize],
104        _share: ShareStrategy,
105    ) {
106        // TODO: SIMD gather kernels?
107        let other_values_slice = other.values().as_slice();
108        self.values.extend(
109            idxs.iter()
110                .map(|idx| *other_values_slice.get_unchecked(*idx as usize)),
111        );
112        self.validity
113            .gather_extend_from_opt_validity(other.validity(), idxs);
114    }
115
116    fn opt_gather_extend(
117        &mut self,
118        other: &PrimitiveArray<T>,
119        idxs: &[IdxSize],
120        _share: ShareStrategy,
121    ) {
122        self.values.reserve(idxs.len());
123        unsafe {
124            for idx in idxs {
125                let val = if (*idx as usize) < other.len() {
126                    other.value_unchecked(*idx as usize)
127                } else {
128                    T::zeroed()
129                };
130                self.values.push_unchecked(val);
131            }
132        }
133        self.validity
134            .opt_gather_extend_from_opt_validity(other.validity(), idxs, other.len());
135    }
136}
polars_arrow/array/primitive/builder.rs

polars_arrow/array/primitive/
builder.rs
