polars_arrow/array/
static_array.rs

1use bytemuck::Zeroable;
2use polars_utils::no_call_const;
3
4use super::growable::{Growable, GrowableFixedSizeList};
5use crate::array::binview::BinaryViewValueIter;
6use crate::array::static_array_collect::ArrayFromIterDtype;
7use crate::array::{
8    Array, ArrayValuesIter, BinaryArray, BinaryValueIter, BinaryViewArray, BooleanArray,
9    FixedSizeListArray, ListArray, ListValuesIter, MutableBinaryViewArray, PrimitiveArray,
10    StructArray, Utf8Array, Utf8ValuesIter, Utf8ViewArray,
11};
12use crate::bitmap::utils::{BitmapIter, ZipValidity};
13use crate::bitmap::Bitmap;
14use crate::datatypes::ArrowDataType;
15use crate::trusted_len::TrustedLen;
16use crate::types::NativeType;
17
18pub trait StaticArray:
19    Array
20    + for<'a> ArrayFromIterDtype<Self::ValueT<'a>>
21    + for<'a> ArrayFromIterDtype<Self::ZeroableValueT<'a>>
22    + for<'a> ArrayFromIterDtype<Option<Self::ValueT<'a>>>
23    + Clone
24{
25    type ValueT<'a>: Clone
26    where
27        Self: 'a;
28    type ZeroableValueT<'a>: Zeroable + From<Self::ValueT<'a>>
29    where
30        Self: 'a;
31    type ValueIterT<'a>: DoubleEndedIterator<Item = Self::ValueT<'a>> + TrustedLen + Send + Sync
32    where
33        Self: 'a;
34
35    #[inline]
36    fn get(&self, idx: usize) -> Option<Self::ValueT<'_>> {
37        if idx >= self.len() {
38            None
39        } else {
40            unsafe { self.get_unchecked(idx) }
41        }
42    }
43
44    /// # Safety
45    /// It is the callers responsibility that the `idx < self.len()`.
46    #[inline]
47    unsafe fn get_unchecked(&self, idx: usize) -> Option<Self::ValueT<'_>> {
48        if self.is_null_unchecked(idx) {
49            None
50        } else {
51            Some(self.value_unchecked(idx))
52        }
53    }
54
55    #[inline]
56    fn last(&self) -> Option<Self::ValueT<'_>> {
57        unsafe { self.get_unchecked(self.len().checked_sub(1)?) }
58    }
59
60    #[inline]
61    fn value(&self, idx: usize) -> Self::ValueT<'_> {
62        assert!(idx < self.len());
63        unsafe { self.value_unchecked(idx) }
64    }
65
66    /// # Safety
67    /// It is the callers responsibility that the `idx < self.len()`.
68    #[allow(unused_variables)]
69    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
70        no_call_const!()
71    }
72
73    #[inline(always)]
74    fn as_slice(&self) -> Option<&[Self::ValueT<'_>]> {
75        None
76    }
77
78    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
79        no_call_const!()
80    }
81    fn values_iter(&self) -> Self::ValueIterT<'_> {
82        no_call_const!()
83    }
84    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self;
85
86    fn from_vec(v: Vec<Self::ValueT<'_>>, dtype: ArrowDataType) -> Self {
87        Self::arr_from_iter_with_dtype(dtype, v)
88    }
89
90    fn from_zeroable_vec(v: Vec<Self::ZeroableValueT<'_>>, dtype: ArrowDataType) -> Self {
91        Self::arr_from_iter_with_dtype(dtype, v)
92    }
93
94    fn full_null(length: usize, dtype: ArrowDataType) -> Self;
95
96    fn full(length: usize, value: Self::ValueT<'_>, dtype: ArrowDataType) -> Self {
97        Self::arr_from_iter_with_dtype(dtype, std::iter::repeat(value).take(length))
98    }
99}
100
101pub trait ParameterFreeDtypeStaticArray: StaticArray {
102    fn get_dtype() -> ArrowDataType;
103}
104
105impl<T: NativeType> StaticArray for PrimitiveArray<T> {
106    type ValueT<'a> = T;
107    type ZeroableValueT<'a> = T;
108    type ValueIterT<'a> = std::iter::Copied<std::slice::Iter<'a, T>>;
109
110    #[inline]
111    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
112        self.value_unchecked(idx)
113    }
114
115    fn values_iter(&self) -> Self::ValueIterT<'_> {
116        self.values_iter().copied()
117    }
118
119    #[inline(always)]
120    fn as_slice(&self) -> Option<&[Self::ValueT<'_>]> {
121        Some(self.values().as_slice())
122    }
123
124    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
125        ZipValidity::new_with_validity(self.values().iter().copied(), self.validity())
126    }
127
128    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
129        self.with_validity(validity)
130    }
131
132    fn from_vec(v: Vec<Self::ValueT<'_>>, _dtype: ArrowDataType) -> Self {
133        PrimitiveArray::from_vec(v)
134    }
135
136    fn from_zeroable_vec(v: Vec<Self::ZeroableValueT<'_>>, _dtype: ArrowDataType) -> Self {
137        PrimitiveArray::from_vec(v)
138    }
139
140    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
141        Self::new_null(dtype, length)
142    }
143
144    fn full(length: usize, value: Self::ValueT<'_>, _dtype: ArrowDataType) -> Self {
145        PrimitiveArray::from_vec(vec![value; length])
146    }
147}
148
149impl<T: NativeType> ParameterFreeDtypeStaticArray for PrimitiveArray<T> {
150    fn get_dtype() -> ArrowDataType {
151        T::PRIMITIVE.into()
152    }
153}
154
155impl StaticArray for BooleanArray {
156    type ValueT<'a> = bool;
157    type ZeroableValueT<'a> = bool;
158    type ValueIterT<'a> = BitmapIter<'a>;
159
160    #[inline]
161    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
162        self.value_unchecked(idx)
163    }
164
165    fn values_iter(&self) -> Self::ValueIterT<'_> {
166        self.values_iter()
167    }
168
169    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
170        self.iter()
171    }
172
173    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
174        self.with_validity(validity)
175    }
176
177    fn from_vec(v: Vec<Self::ValueT<'_>>, _dtype: ArrowDataType) -> Self {
178        BooleanArray::from_slice(v)
179    }
180
181    fn from_zeroable_vec(v: Vec<Self::ValueT<'_>>, _dtype: ArrowDataType) -> Self {
182        BooleanArray::from_slice(v)
183    }
184
185    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
186        Self::new_null(dtype, length)
187    }
188
189    fn full(length: usize, value: Self::ValueT<'_>, _dtype: ArrowDataType) -> Self {
190        Bitmap::new_with_value(value, length).into()
191    }
192}
193
194impl ParameterFreeDtypeStaticArray for BooleanArray {
195    fn get_dtype() -> ArrowDataType {
196        ArrowDataType::Boolean
197    }
198}
199
200impl StaticArray for Utf8Array<i64> {
201    type ValueT<'a> = &'a str;
202    type ZeroableValueT<'a> = Option<&'a str>;
203    type ValueIterT<'a> = Utf8ValuesIter<'a, i64>;
204
205    #[inline]
206    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
207        self.value_unchecked(idx)
208    }
209
210    fn values_iter(&self) -> Self::ValueIterT<'_> {
211        self.values_iter()
212    }
213
214    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
215        self.iter()
216    }
217
218    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
219        self.with_validity(validity)
220    }
221
222    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
223        Self::new_null(dtype, length)
224    }
225}
226
227impl ParameterFreeDtypeStaticArray for Utf8Array<i64> {
228    fn get_dtype() -> ArrowDataType {
229        ArrowDataType::LargeUtf8
230    }
231}
232
233impl StaticArray for BinaryArray<i64> {
234    type ValueT<'a> = &'a [u8];
235    type ZeroableValueT<'a> = Option<&'a [u8]>;
236    type ValueIterT<'a> = BinaryValueIter<'a, i64>;
237
238    #[inline]
239    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
240        self.value_unchecked(idx)
241    }
242
243    fn values_iter(&self) -> Self::ValueIterT<'_> {
244        self.values_iter()
245    }
246
247    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
248        self.iter()
249    }
250
251    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
252        self.with_validity(validity)
253    }
254
255    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
256        Self::new_null(dtype, length)
257    }
258}
259
260impl ParameterFreeDtypeStaticArray for BinaryArray<i64> {
261    fn get_dtype() -> ArrowDataType {
262        ArrowDataType::LargeBinary
263    }
264}
265
266impl StaticArray for BinaryViewArray {
267    type ValueT<'a> = &'a [u8];
268    type ZeroableValueT<'a> = Option<&'a [u8]>;
269    type ValueIterT<'a> = BinaryViewValueIter<'a, [u8]>;
270
271    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
272        self.value_unchecked(idx)
273    }
274
275    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
276        self.iter()
277    }
278
279    fn values_iter(&self) -> Self::ValueIterT<'_> {
280        self.values_iter()
281    }
282
283    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
284        self.with_validity(validity)
285    }
286
287    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
288        Self::new_null(dtype, length)
289    }
290
291    fn full(length: usize, value: Self::ValueT<'_>, _dtype: ArrowDataType) -> Self {
292        let mut builder = MutableBinaryViewArray::with_capacity(length);
293        builder.extend_constant(length, Some(value));
294        builder.into()
295    }
296}
297
298impl ParameterFreeDtypeStaticArray for BinaryViewArray {
299    fn get_dtype() -> ArrowDataType {
300        ArrowDataType::BinaryView
301    }
302}
303
304impl StaticArray for Utf8ViewArray {
305    type ValueT<'a> = &'a str;
306    type ZeroableValueT<'a> = Option<&'a str>;
307    type ValueIterT<'a> = BinaryViewValueIter<'a, str>;
308
309    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
310        self.value_unchecked(idx)
311    }
312
313    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
314        self.iter()
315    }
316
317    fn values_iter(&self) -> Self::ValueIterT<'_> {
318        self.values_iter()
319    }
320
321    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
322        self.with_validity(validity)
323    }
324
325    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
326        Self::new_null(dtype, length)
327    }
328
329    fn full(length: usize, value: Self::ValueT<'_>, _dtype: ArrowDataType) -> Self {
330        unsafe {
331            BinaryViewArray::full(length, value.as_bytes(), ArrowDataType::BinaryView)
332                .to_utf8view_unchecked()
333        }
334    }
335}
336
337impl ParameterFreeDtypeStaticArray for Utf8ViewArray {
338    fn get_dtype() -> ArrowDataType {
339        ArrowDataType::Utf8View
340    }
341}
342
343impl StaticArray for ListArray<i64> {
344    type ValueT<'a> = Box<dyn Array>;
345    type ZeroableValueT<'a> = Option<Box<dyn Array>>;
346    type ValueIterT<'a> = ListValuesIter<'a, i64>;
347
348    #[inline]
349    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
350        self.value_unchecked(idx)
351    }
352
353    fn values_iter(&self) -> Self::ValueIterT<'_> {
354        self.values_iter()
355    }
356
357    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
358        self.iter()
359    }
360
361    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
362        self.with_validity(validity)
363    }
364
365    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
366        Self::new_null(dtype, length)
367    }
368}
369
370impl StaticArray for FixedSizeListArray {
371    type ValueT<'a> = Box<dyn Array>;
372    type ZeroableValueT<'a> = Option<Box<dyn Array>>;
373    type ValueIterT<'a> = ArrayValuesIter<'a, FixedSizeListArray>;
374
375    #[inline]
376    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
377        self.value_unchecked(idx)
378    }
379
380    fn values_iter(&self) -> Self::ValueIterT<'_> {
381        self.values_iter()
382    }
383
384    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
385        self.iter()
386    }
387
388    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
389        self.with_validity(validity)
390    }
391
392    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
393        Self::new_null(dtype, length)
394    }
395
396    fn full(length: usize, value: Self::ValueT<'_>, dtype: ArrowDataType) -> Self {
397        let singular_arr = FixedSizeListArray::new(dtype, 1, value, None);
398        let mut arr = GrowableFixedSizeList::new(vec![&singular_arr], false, length);
399        unsafe { arr.extend_copies(0, 0, 1, length) }
400        arr.into()
401    }
402}
403
404impl StaticArray for StructArray {
405    type ValueT<'a> = ();
406    type ZeroableValueT<'a> = ();
407    type ValueIterT<'a> = std::iter::Repeat<()>;
408
409    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
410        self.with_validity(validity)
411    }
412
413    fn full_null(length: usize, dtype: ArrowDataType) -> Self {
414        Self::new_null(dtype, length)
415    }
416}
polars_arrow/array/static_array.rs

polars_arrow/array/
static_array.rs
