CoCalc -- static

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-arrow/src/array/static_array.rs
⁶⁹³⁹ views
1
use bytemuck::Zeroable;
2
use polars_utils::no_call_const;
3

4
use crate::array::binview::BinaryViewValueIter;
5
use crate::array::builder::{ShareStrategy, StaticArrayBuilder, make_builder};
6
use crate::array::fixed_size_list::FixedSizeListArrayBuilder;
7
use crate::array::static_array_collect::ArrayFromIterDtype;
8
use crate::array::{
9
    Array, ArrayValuesIter, BinaryArray, BinaryValueIter, BinaryViewArray, BooleanArray,
10
    FixedSizeListArray, ListArray, ListValuesIter, MutableBinaryViewArray, PrimitiveArray,
11
    StructArray, Utf8Array, Utf8ValuesIter, Utf8ViewArray,
12
};
13
use crate::bitmap::Bitmap;
14
use crate::bitmap::utils::{BitmapIter, ZipValidity};
15
use crate::datatypes::ArrowDataType;
16
use crate::trusted_len::TrustedLen;
17
use crate::types::NativeType;
18

19
pub 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

102
pub trait ParameterFreeDtypeStaticArray: StaticArray {
103
    fn get_dtype() -> ArrowDataType;
104
}
105

106
impl<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

150
impl<T: NativeType> ParameterFreeDtypeStaticArray for PrimitiveArray<T> {
151
    fn get_dtype() -> ArrowDataType {
152
        T::PRIMITIVE.into()
153
    }
154
}
155

156
impl 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

195
impl ParameterFreeDtypeStaticArray for BooleanArray {
196
    fn get_dtype() -> ArrowDataType {
197
        ArrowDataType::Boolean
198
    }
199
}
200

201
impl 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

228
impl ParameterFreeDtypeStaticArray for Utf8Array<i64> {
229
    fn get_dtype() -> ArrowDataType {
230
        ArrowDataType::LargeUtf8
231
    }
232
}
233

234
impl 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

261
impl ParameterFreeDtypeStaticArray for BinaryArray<i64> {
262
    fn get_dtype() -> ArrowDataType {
263
        ArrowDataType::LargeBinary
264
    }
265
}
266

267
impl 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

299
impl ParameterFreeDtypeStaticArray for BinaryViewArray {
300
    fn get_dtype() -> ArrowDataType {
301
        ArrowDataType::BinaryView
302
    }
303
}
304

305
impl 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

338
impl ParameterFreeDtypeStaticArray for Utf8ViewArray {
339
    fn get_dtype() -> ArrowDataType {
340
        ArrowDataType::Utf8View
341
    }
342
}
343

344
impl 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

371
impl 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

406
impl 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
}
419

420
Product

Resources

Company
