arrow2/array/growable/
primitive.rs

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

arrow2/array/growable/
primitive.rs
