polars_arrow/array/growable/
primitive.rs

1use std::sync::Arc;
2
3use super::Growable;
4use crate::array::growable::utils::{extend_validity, extend_validity_copies, prepare_validity};
5use crate::array::{Array, PrimitiveArray};
6use crate::bitmap::BitmapBuilder;
7use crate::datatypes::ArrowDataType;
8use crate::types::NativeType;
9
10/// Concrete [`Growable`] for the [`PrimitiveArray`].
11pub struct GrowablePrimitive<'a, T: NativeType> {
12    dtype: ArrowDataType,
13    arrays: Vec<&'a PrimitiveArray<T>>,
14    validity: Option<BitmapBuilder>,
15    values: Vec<T>,
16}
17
18impl<'a, T: NativeType> GrowablePrimitive<'a, T> {
19    /// Creates a new [`GrowablePrimitive`] bound to `arrays` with a pre-allocated `capacity`.
20    /// # Panics
21    /// If `arrays` is empty.
22    pub fn new(
23        arrays: Vec<&'a PrimitiveArray<T>>,
24        mut use_validity: bool,
25        capacity: usize,
26    ) -> Self {
27        // if any of the arrays has nulls, insertions from any array requires setting bits
28        // as there is at least one array with nulls.
29        if !use_validity & arrays.iter().any(|array| array.null_count() > 0) {
30            use_validity = true;
31        };
32
33        let dtype = arrays[0].dtype().clone();
34
35        Self {
36            dtype,
37            arrays,
38            values: Vec::with_capacity(capacity),
39            validity: prepare_validity(use_validity, capacity),
40        }
41    }
42
43    #[inline]
44    fn to(&mut self) -> PrimitiveArray<T> {
45        let validity = std::mem::take(&mut self.validity);
46        let values = std::mem::take(&mut self.values);
47
48        PrimitiveArray::<T>::new(
49            self.dtype.clone(),
50            values.into(),
51            validity.map(|v| v.freeze()),
52        )
53    }
54}
55
56impl<'a, T: NativeType> Growable<'a> for GrowablePrimitive<'a, T> {
57    #[inline]
58    unsafe fn extend(&mut self, index: usize, start: usize, len: usize) {
59        let array = *self.arrays.get_unchecked(index);
60        extend_validity(&mut self.validity, array, start, len);
61
62        let values = array.values().as_slice();
63        self.values
64            .extend_from_slice(values.get_unchecked(start..start + len));
65    }
66
67    #[inline]
68    unsafe fn extend_copies(&mut self, index: usize, start: usize, len: usize, copies: usize) {
69        let array = *self.arrays.get_unchecked(index);
70        extend_validity_copies(&mut self.validity, array, start, len, copies);
71
72        let values = array.values().as_slice();
73        self.values.reserve(len * copies);
74        for _ in 0..copies {
75            self.values
76                .extend_from_slice(values.get_unchecked(start..start + len));
77        }
78    }
79
80    #[inline]
81    fn extend_validity(&mut self, additional: usize) {
82        self.values
83            .resize(self.values.len() + additional, T::default());
84        if let Some(validity) = &mut self.validity {
85            validity.extend_constant(additional, false);
86        }
87    }
88
89    #[inline]
90    fn len(&self) -> usize {
91        self.values.len()
92    }
93
94    #[inline]
95    fn as_arc(&mut self) -> Arc<dyn Array> {
96        Arc::new(self.to())
97    }
98
99    #[inline]
100    fn as_box(&mut self) -> Box<dyn Array> {
101        Box::new(self.to())
102    }
103}
104
105impl<'a, T: NativeType> From<GrowablePrimitive<'a, T>> for PrimitiveArray<T> {
106    #[inline]
107    fn from(val: GrowablePrimitive<'a, T>) -> Self {
108        PrimitiveArray::<T>::new(
109            val.dtype,
110            val.values.into(),
111            val.validity.map(|v| v.freeze()),
112        )
113    }
114}
polars_arrow/array/growable/primitive.rs

polars_arrow/array/growable/
primitive.rs
