polars_arrow/array/fixed_size_binary/
builder.rs

1use polars_utils::IdxSize;
2
3use super::FixedSizeBinaryArray;
4use crate::array::builder::{ShareStrategy, StaticArrayBuilder};
5use crate::bitmap::OptBitmapBuilder;
6use crate::buffer::Buffer;
7use crate::datatypes::ArrowDataType;
8use crate::pushable::Pushable;
9
10pub struct FixedSizeBinaryArrayBuilder {
11    dtype: ArrowDataType,
12    size: usize,
13    length: usize,
14    values: Vec<u8>,
15    validity: OptBitmapBuilder,
16}
17
18impl FixedSizeBinaryArrayBuilder {
19    pub fn new(dtype: ArrowDataType) -> Self {
20        Self {
21            size: FixedSizeBinaryArray::get_size(&dtype),
22            length: 0,
23            dtype,
24            values: Vec::new(),
25            validity: OptBitmapBuilder::default(),
26        }
27    }
28}
29
30impl StaticArrayBuilder for FixedSizeBinaryArrayBuilder {
31    type Array = FixedSizeBinaryArray;
32
33    fn dtype(&self) -> &ArrowDataType {
34        &self.dtype
35    }
36
37    fn reserve(&mut self, additional: usize) {
38        let bytes = additional * self.size;
39        self.values.reserve(bytes);
40        self.validity.reserve(additional);
41    }
42
43    fn freeze(self) -> FixedSizeBinaryArray {
44        // TODO: FixedSizeBinaryArray should track its own length to be correct
45        // for size-0 inner.
46        let values = Buffer::from(self.values);
47        let validity = self.validity.into_opt_validity();
48        FixedSizeBinaryArray::new(self.dtype, values, validity)
49    }
50
51    fn freeze_reset(&mut self) -> Self::Array {
52        // TODO: FixedSizeBinaryArray should track its own length to be correct
53        // for size-0 inner.
54        let values = Buffer::from(core::mem::take(&mut self.values));
55        let validity = core::mem::take(&mut self.validity).into_opt_validity();
56        let out = FixedSizeBinaryArray::new(self.dtype.clone(), values, validity);
57        self.length = 0;
58        out
59    }
60
61    fn len(&self) -> usize {
62        self.length
63    }
64
65    fn extend_nulls(&mut self, length: usize) {
66        self.values.extend_constant(length * self.size, 0);
67        self.validity.extend_constant(length, false);
68        self.length += length;
69    }
70
71    fn subslice_extend(
72        &mut self,
73        other: &FixedSizeBinaryArray,
74        start: usize,
75        length: usize,
76        _share: ShareStrategy,
77    ) {
78        let other_slice = other.values().as_slice();
79        self.values
80            .extend_from_slice(&other_slice[start * self.size..(start + length) * self.size]);
81        self.validity
82            .subslice_extend_from_opt_validity(other.validity(), start, length);
83        self.length += length.min(other.len().saturating_sub(start));
84    }
85
86    fn subslice_extend_each_repeated(
87        &mut self,
88        other: &FixedSizeBinaryArray,
89        start: usize,
90        length: usize,
91        repeats: usize,
92        _share: ShareStrategy,
93    ) {
94        let other_slice = other.values().as_slice();
95        for outer_idx in start..start + length {
96            for _ in 0..repeats {
97                self.values.extend_from_slice(
98                    &other_slice[outer_idx * self.size..(outer_idx + 1) * self.size],
99                );
100            }
101        }
102
103        self.validity
104            .subslice_extend_each_repeated_from_opt_validity(
105                other.validity(),
106                start,
107                length,
108                repeats,
109            );
110        self.length += repeats * length.min(other.len().saturating_sub(start));
111    }
112
113    unsafe fn gather_extend(
114        &mut self,
115        other: &FixedSizeBinaryArray,
116        idxs: &[IdxSize],
117        _share: ShareStrategy,
118    ) {
119        let other_slice = other.values().as_slice();
120        self.values.reserve(idxs.len() * self.size);
121        for idx in idxs {
122            let idx = *idx as usize;
123            let subslice = other_slice.get_unchecked(idx * self.size..(idx + 1) * self.size);
124            self.values.extend_from_slice(subslice);
125        }
126        self.validity
127            .gather_extend_from_opt_validity(other.validity(), idxs);
128        self.length += idxs.len();
129    }
130
131    fn opt_gather_extend(
132        &mut self,
133        other: &FixedSizeBinaryArray,
134        idxs: &[IdxSize],
135        _share: ShareStrategy,
136    ) {
137        let other_slice = other.values().as_slice();
138        self.values.reserve(idxs.len() * self.size);
139        for idx in idxs {
140            let idx = *idx as usize;
141            if let Some(subslice) = other_slice.get(idx * self.size..(idx + 1) * self.size) {
142                self.values.extend_from_slice(subslice);
143            } else {
144                self.values.extend_constant(self.size, 0);
145            }
146        }
147        self.validity
148            .opt_gather_extend_from_opt_validity(other.validity(), idxs, other.len());
149        self.length += idxs.len();
150    }
151}
polars_arrow/array/fixed_size_binary/builder.rs

polars_arrow/array/fixed_size_binary/
builder.rs
