polars_arrow/array/struct_/
builder.rs

1use polars_utils::IdxSize;
2
3use super::StructArray;
4use crate::array::builder::{ArrayBuilder, ShareStrategy, StaticArrayBuilder};
5use crate::bitmap::OptBitmapBuilder;
6use crate::datatypes::ArrowDataType;
7
8pub struct StructArrayBuilder {
9    dtype: ArrowDataType,
10    length: usize,
11    inner_builders: Vec<Box<dyn ArrayBuilder>>,
12    validity: OptBitmapBuilder,
13}
14
15impl StructArrayBuilder {
16    pub fn new(dtype: ArrowDataType, inner_builders: Vec<Box<dyn ArrayBuilder>>) -> Self {
17        Self {
18            dtype,
19            length: 0,
20            inner_builders,
21            validity: OptBitmapBuilder::default(),
22        }
23    }
24}
25
26impl StaticArrayBuilder for StructArrayBuilder {
27    type Array = StructArray;
28
29    fn dtype(&self) -> &ArrowDataType {
30        &self.dtype
31    }
32
33    fn reserve(&mut self, additional: usize) {
34        for builder in &mut self.inner_builders {
35            builder.reserve(additional);
36        }
37        self.validity.reserve(additional);
38    }
39
40    fn freeze(self) -> StructArray {
41        let values = self
42            .inner_builders
43            .into_iter()
44            .map(|b| b.freeze())
45            .collect();
46        let validity = self.validity.into_opt_validity();
47        StructArray::new(self.dtype, self.length, values, validity)
48    }
49
50    fn freeze_reset(&mut self) -> Self::Array {
51        let values = self
52            .inner_builders
53            .iter_mut()
54            .map(|b| b.freeze_reset())
55            .collect();
56        let validity = core::mem::take(&mut self.validity).into_opt_validity();
57        let out = StructArray::new(self.dtype.clone(), self.length, values, validity);
58        self.length = 0;
59        out
60    }
61
62    fn len(&self) -> usize {
63        self.length
64    }
65
66    fn extend_nulls(&mut self, length: usize) {
67        for builder in &mut self.inner_builders {
68            builder.extend_nulls(length);
69        }
70        self.validity.extend_constant(length, false);
71        self.length += length;
72    }
73
74    fn subslice_extend(
75        &mut self,
76        other: &StructArray,
77        start: usize,
78        length: usize,
79        share: ShareStrategy,
80    ) {
81        for (builder, other_values) in self.inner_builders.iter_mut().zip(other.values()) {
82            builder.subslice_extend(&**other_values, start, length, share);
83        }
84        self.validity
85            .subslice_extend_from_opt_validity(other.validity(), start, length);
86        self.length += length.min(other.len().saturating_sub(start));
87    }
88
89    fn subslice_extend_each_repeated(
90        &mut self,
91        other: &StructArray,
92        start: usize,
93        length: usize,
94        repeats: usize,
95        share: ShareStrategy,
96    ) {
97        for (builder, other_values) in self.inner_builders.iter_mut().zip(other.values()) {
98            builder.subslice_extend_each_repeated(&**other_values, start, length, repeats, share);
99        }
100        self.validity
101            .subslice_extend_from_opt_validity(other.validity(), start, length);
102        self.length += length.min(other.len().saturating_sub(start)) * repeats;
103    }
104
105    unsafe fn gather_extend(
106        &mut self,
107        other: &StructArray,
108        idxs: &[IdxSize],
109        share: ShareStrategy,
110    ) {
111        for (builder, other_values) in self.inner_builders.iter_mut().zip(other.values()) {
112            builder.gather_extend(&**other_values, idxs, share);
113        }
114        self.validity
115            .gather_extend_from_opt_validity(other.validity(), idxs);
116        self.length += idxs.len();
117    }
118
119    fn opt_gather_extend(&mut self, other: &StructArray, idxs: &[IdxSize], share: ShareStrategy) {
120        for (builder, other_values) in self.inner_builders.iter_mut().zip(other.values()) {
121            builder.opt_gather_extend(&**other_values, idxs, share);
122        }
123        self.validity
124            .opt_gather_extend_from_opt_validity(other.validity(), idxs, other.len());
125        self.length += idxs.len();
126    }
127}
polars_arrow/array/struct_/builder.rs

polars_arrow/array/struct_/
builder.rs
