polars_arrow/array/growable/
structure.rs

1use std::sync::Arc;
2
3use super::{make_growable, Growable};
4use crate::array::growable::utils::{extend_validity, prepare_validity};
5use crate::array::{Array, StructArray};
6use crate::bitmap::BitmapBuilder;
7
8/// Concrete [`Growable`] for the [`StructArray`].
9pub struct GrowableStruct<'a> {
10    arrays: Vec<&'a StructArray>,
11    length: usize,
12    validity: Option<BitmapBuilder>,
13    values: Vec<Box<dyn Growable<'a> + 'a>>,
14}
15
16impl<'a> GrowableStruct<'a> {
17    /// Creates a new [`GrowableStruct`] bound to `arrays` with a pre-allocated `capacity`.
18    /// # Panics
19    /// If `arrays` is empty.
20    pub fn new(arrays: Vec<&'a StructArray>, mut use_validity: bool, capacity: usize) -> Self {
21        assert!(!arrays.is_empty());
22
23        // if any of the arrays has nulls, insertions from any array requires setting bits
24        // as there is at least one array with nulls.
25        if arrays.iter().any(|array| array.null_count() > 0) {
26            use_validity = true;
27        };
28
29        let arrays = arrays
30            .iter()
31            .map(|array| array.as_any().downcast_ref::<StructArray>().unwrap())
32            .collect::<Vec<_>>();
33
34        // ([field1, field2], [field3, field4]) -> ([field1, field3], [field2, field3])
35        let values = (0..arrays[0].values().len())
36            .map(|i| {
37                make_growable(
38                    &arrays
39                        .iter()
40                        .map(|x| x.values()[i].as_ref())
41                        .collect::<Vec<_>>(),
42                    use_validity,
43                    capacity,
44                )
45            })
46            .collect::<Vec<Box<dyn Growable>>>();
47
48        Self {
49            arrays,
50            length: 0,
51            values,
52            validity: prepare_validity(use_validity, capacity),
53        }
54    }
55
56    fn to(&mut self) -> StructArray {
57        let validity = std::mem::take(&mut self.validity);
58        let values = std::mem::take(&mut self.values);
59        let values = values.into_iter().map(|mut x| x.as_box()).collect();
60
61        StructArray::new(
62            self.arrays[0].dtype().clone(),
63            self.length,
64            values,
65            validity.map(|v| v.freeze()),
66        )
67    }
68}
69
70impl<'a> Growable<'a> for GrowableStruct<'a> {
71    unsafe fn extend(&mut self, index: usize, start: usize, len: usize) {
72        let array = *self.arrays.get_unchecked(index);
73        extend_validity(&mut self.validity, array, start, len);
74
75        self.length += len;
76
77        if array.null_count() == 0 {
78            self.values
79                .iter_mut()
80                .for_each(|child| child.extend(index, start, len))
81        } else {
82            (start..start + len).for_each(|i| {
83                if array.is_valid(i) {
84                    self.values
85                        .iter_mut()
86                        .for_each(|child| child.extend(index, i, 1))
87                } else {
88                    self.values
89                        .iter_mut()
90                        .for_each(|child| child.extend_validity(1))
91                }
92            })
93        }
94    }
95
96    fn extend_validity(&mut self, additional: usize) {
97        self.values
98            .iter_mut()
99            .for_each(|child| child.extend_validity(additional));
100        if let Some(validity) = &mut self.validity {
101            validity.extend_constant(additional, false);
102        }
103        self.length += additional;
104    }
105
106    #[inline]
107    fn len(&self) -> usize {
108        if let Some(child) = self.values.first() {
109            child.len()
110        } else {
111            unreachable!()
112        }
113    }
114
115    fn as_arc(&mut self) -> Arc<dyn Array> {
116        Arc::new(self.to())
117    }
118
119    fn as_box(&mut self) -> Box<dyn Array> {
120        Box::new(self.to())
121    }
122}
123
124impl<'a> From<GrowableStruct<'a>> for StructArray {
125    fn from(val: GrowableStruct<'a>) -> Self {
126        let values = val.values.into_iter().map(|mut x| x.as_box()).collect();
127
128        StructArray::new(
129            val.arrays[0].dtype().clone(),
130            val.length,
131            values,
132            val.validity.map(|v| v.freeze()),
133        )
134    }
135}
polars_arrow/array/growable/structure.rs

polars_arrow/array/growable/
structure.rs
