polars_arrow/
pushable.rs

1use crate::array::{
2    BinaryViewArrayGeneric, BooleanArray, MutableBinaryViewArray, MutableBooleanArray,
3    MutablePrimitiveArray, PrimitiveArray, ViewType,
4};
5use crate::bitmap::{Bitmap, MutableBitmap};
6use crate::offset::{Offset, Offsets, OffsetsBuffer};
7use crate::types::NativeType;
8
9/// A private trait representing structs that can receive elements.
10pub trait Pushable<T>: Sized + Default {
11    type Freeze;
12
13    fn with_capacity(capacity: usize) -> Self {
14        let mut new = Self::default();
15        new.reserve(capacity);
16        new
17    }
18    fn reserve(&mut self, additional: usize);
19    fn push(&mut self, value: T);
20    fn len(&self) -> usize;
21    fn push_null(&mut self);
22    #[inline]
23    fn extend_n(&mut self, n: usize, iter: impl Iterator<Item = T>) {
24        for item in iter.take(n) {
25            self.push(item);
26        }
27    }
28    fn extend_constant(&mut self, additional: usize, value: T);
29    fn extend_null_constant(&mut self, additional: usize);
30    fn freeze(self) -> Self::Freeze;
31}
32
33impl Pushable<bool> for MutableBitmap {
34    type Freeze = Bitmap;
35
36    #[inline]
37    fn reserve(&mut self, additional: usize) {
38        MutableBitmap::reserve(self, additional)
39    }
40
41    #[inline]
42    fn len(&self) -> usize {
43        self.len()
44    }
45
46    #[inline]
47    fn push(&mut self, value: bool) {
48        self.push(value)
49    }
50
51    #[inline]
52    fn push_null(&mut self) {
53        self.push(false)
54    }
55
56    #[inline]
57    fn extend_constant(&mut self, additional: usize, value: bool) {
58        self.extend_constant(additional, value)
59    }
60
61    #[inline]
62    fn extend_null_constant(&mut self, additional: usize) {
63        self.extend_constant(additional, false)
64    }
65
66    fn freeze(self) -> Self::Freeze {
67        self.into()
68    }
69}
70
71impl<T: Copy + Default> Pushable<T> for Vec<T> {
72    type Freeze = Vec<T>;
73    #[inline]
74    fn reserve(&mut self, additional: usize) {
75        Vec::reserve(self, additional)
76    }
77    #[inline]
78    fn len(&self) -> usize {
79        self.len()
80    }
81
82    #[inline]
83    fn push_null(&mut self) {
84        self.push(T::default())
85    }
86
87    #[inline]
88    fn push(&mut self, value: T) {
89        self.push(value)
90    }
91
92    #[inline]
93    fn extend_n(&mut self, n: usize, iter: impl Iterator<Item = T>) {
94        self.extend(iter.take(n));
95    }
96
97    #[inline]
98    fn extend_constant(&mut self, additional: usize, value: T) {
99        self.resize(self.len() + additional, value);
100    }
101
102    #[inline]
103    fn extend_null_constant(&mut self, additional: usize) {
104        self.extend_constant(additional, T::default())
105    }
106    fn freeze(self) -> Self::Freeze {
107        self
108    }
109}
110impl<O: Offset> Pushable<usize> for Offsets<O> {
111    type Freeze = OffsetsBuffer<O>;
112    fn reserve(&mut self, additional: usize) {
113        self.reserve(additional)
114    }
115    #[inline]
116    fn len(&self) -> usize {
117        self.len_proxy()
118    }
119
120    #[inline]
121    fn push(&mut self, value: usize) {
122        self.try_push(value).unwrap()
123    }
124
125    #[inline]
126    fn push_null(&mut self) {
127        self.extend_constant(1);
128    }
129
130    #[inline]
131    fn extend_constant(&mut self, additional: usize, _: usize) {
132        self.extend_constant(additional)
133    }
134
135    #[inline]
136    fn extend_null_constant(&mut self, additional: usize) {
137        self.extend_constant(additional)
138    }
139    fn freeze(self) -> Self::Freeze {
140        self.into()
141    }
142}
143
144impl<T: NativeType> Pushable<Option<T>> for MutablePrimitiveArray<T> {
145    type Freeze = PrimitiveArray<T>;
146
147    #[inline]
148    fn reserve(&mut self, additional: usize) {
149        MutablePrimitiveArray::reserve(self, additional)
150    }
151
152    #[inline]
153    fn push(&mut self, value: Option<T>) {
154        MutablePrimitiveArray::push(self, value)
155    }
156
157    #[inline]
158    fn len(&self) -> usize {
159        self.values().len()
160    }
161
162    #[inline]
163    fn push_null(&mut self) {
164        self.push(None)
165    }
166
167    #[inline]
168    fn extend_constant(&mut self, additional: usize, value: Option<T>) {
169        MutablePrimitiveArray::extend_constant(self, additional, value)
170    }
171
172    #[inline]
173    fn extend_null_constant(&mut self, additional: usize) {
174        MutablePrimitiveArray::extend_constant(self, additional, None)
175    }
176    fn freeze(self) -> Self::Freeze {
177        self.into()
178    }
179}
180
181pub trait NoOption {}
182impl NoOption for &str {}
183impl NoOption for &[u8] {}
184impl NoOption for Vec<u8> {}
185
186impl<T, K> Pushable<T> for MutableBinaryViewArray<K>
187where
188    T: AsRef<K> + NoOption,
189    K: ViewType + ?Sized,
190{
191    type Freeze = BinaryViewArrayGeneric<K>;
192
193    #[inline]
194    fn reserve(&mut self, additional: usize) {
195        MutableBinaryViewArray::reserve(self, additional)
196    }
197
198    #[inline]
199    fn push(&mut self, value: T) {
200        MutableBinaryViewArray::push_value(self, value)
201    }
202
203    #[inline]
204    fn len(&self) -> usize {
205        MutableBinaryViewArray::len(self)
206    }
207
208    fn push_null(&mut self) {
209        MutableBinaryViewArray::push_null(self)
210    }
211
212    fn extend_constant(&mut self, additional: usize, value: T) {
213        MutableBinaryViewArray::extend_constant(self, additional, Some(value));
214    }
215
216    #[inline]
217    fn extend_null_constant(&mut self, additional: usize) {
218        self.extend_null(additional);
219    }
220    fn freeze(self) -> Self::Freeze {
221        self.into()
222    }
223}
224
225impl<T, K> Pushable<Option<T>> for MutableBinaryViewArray<K>
226where
227    T: AsRef<K>,
228    K: ViewType + ?Sized,
229{
230    type Freeze = BinaryViewArrayGeneric<K>;
231    #[inline]
232    fn reserve(&mut self, additional: usize) {
233        MutableBinaryViewArray::reserve(self, additional)
234    }
235
236    #[inline]
237    fn push(&mut self, value: Option<T>) {
238        MutableBinaryViewArray::push(self, value.as_ref())
239    }
240
241    #[inline]
242    fn len(&self) -> usize {
243        MutableBinaryViewArray::len(self)
244    }
245
246    fn push_null(&mut self) {
247        MutableBinaryViewArray::push_null(self)
248    }
249
250    fn extend_constant(&mut self, additional: usize, value: Option<T>) {
251        MutableBinaryViewArray::extend_constant(self, additional, value);
252    }
253
254    #[inline]
255    fn extend_null_constant(&mut self, additional: usize) {
256        self.extend_null(additional);
257    }
258    fn freeze(self) -> Self::Freeze {
259        self.into()
260    }
261}
262
263impl Pushable<bool> for MutableBooleanArray {
264    type Freeze = BooleanArray;
265    #[inline]
266    fn reserve(&mut self, additional: usize) {
267        MutableBooleanArray::reserve(self, additional)
268    }
269
270    #[inline]
271    fn push(&mut self, value: bool) {
272        MutableBooleanArray::push_value(self, value)
273    }
274
275    #[inline]
276    fn len(&self) -> usize {
277        self.values().len()
278    }
279
280    #[inline]
281    fn push_null(&mut self) {
282        unimplemented!()
283    }
284
285    #[inline]
286    fn extend_constant(&mut self, additional: usize, value: bool) {
287        MutableBooleanArray::extend_constant(self, additional, Some(value))
288    }
289
290    #[inline]
291    fn extend_null_constant(&mut self, _additional: usize) {
292        unimplemented!()
293    }
294    fn freeze(self) -> Self::Freeze {
295        self.into()
296    }
297}
298
299impl Pushable<Option<bool>> for MutableBooleanArray {
300    type Freeze = BooleanArray;
301    #[inline]
302    fn reserve(&mut self, additional: usize) {
303        MutableBooleanArray::reserve(self, additional)
304    }
305
306    #[inline]
307    fn push(&mut self, value: Option<bool>) {
308        MutableBooleanArray::push(self, value)
309    }
310
311    #[inline]
312    fn len(&self) -> usize {
313        self.values().len()
314    }
315
316    #[inline]
317    fn push_null(&mut self) {
318        MutableBooleanArray::push_null(self)
319    }
320
321    #[inline]
322    fn extend_constant(&mut self, additional: usize, value: Option<bool>) {
323        MutableBooleanArray::extend_constant(self, additional, value)
324    }
325
326    #[inline]
327    fn extend_null_constant(&mut self, additional: usize) {
328        MutableBooleanArray::extend_constant(self, additional, None)
329    }
330    fn freeze(self) -> Self::Freeze {
331        self.into()
332    }
333}