arrow2/array/growable/
structure.rs

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

arrow2/array/growable/
structure.rs
