polars_arrow/array/
static_array.rs

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

polars_arrow/array/
static_array.rs
