polars_arrow/array/fixed_size_list/
builder.rs

1use polars_utils::IdxSize;
2
3use super::FixedSizeListArray;
4use crate::array::builder::{ArrayBuilder, ShareStrategy, StaticArrayBuilder};
5use crate::bitmap::OptBitmapBuilder;
6use crate::datatypes::ArrowDataType;
7
8pub struct FixedSizeListArrayBuilder<B: ArrayBuilder> {
9    dtype: ArrowDataType,
10    size: usize,
11    length: usize,
12    inner_builder: B,
13    validity: OptBitmapBuilder,
14}
15impl<B: ArrayBuilder> FixedSizeListArrayBuilder<B> {
16    pub fn new(dtype: ArrowDataType, inner_builder: B) -> Self {
17        Self {
18            size: FixedSizeListArray::get_child_and_size(&dtype).1,
19            dtype,
20            length: 0,
21            inner_builder,
22            validity: OptBitmapBuilder::default(),
23        }
24    }
25}
26
27impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
28    type Array = FixedSizeListArray;
29
30    fn dtype(&self) -> &ArrowDataType {
31        &self.dtype
32    }
33
34    fn reserve(&mut self, additional: usize) {
35        self.inner_builder.reserve(additional);
36        self.validity.reserve(additional);
37    }
38
39    fn freeze(self) -> FixedSizeListArray {
40        let values = self.inner_builder.freeze();
41        let validity = self.validity.into_opt_validity();
42        FixedSizeListArray::new(self.dtype, self.length, values, validity)
43    }
44
45    fn freeze_reset(&mut self) -> Self::Array {
46        let values = self.inner_builder.freeze_reset();
47        let validity = core::mem::take(&mut self.validity).into_opt_validity();
48        let out = FixedSizeListArray::new(self.dtype.clone(), self.length, values, validity);
49        self.length = 0;
50        out
51    }
52
53    fn len(&self) -> usize {
54        self.length
55    }
56
57    fn extend_nulls(&mut self, length: usize) {
58        self.inner_builder.extend_nulls(length * self.size);
59        self.validity.extend_constant(length, false);
60        self.length += length;
61    }
62
63    fn subslice_extend(
64        &mut self,
65        other: &FixedSizeListArray,
66        start: usize,
67        length: usize,
68        share: ShareStrategy,
69    ) {
70        self.inner_builder.subslice_extend(
71            &**other.values(),
72            start * self.size,
73            length * self.size,
74            share,
75        );
76        self.validity
77            .subslice_extend_from_opt_validity(other.validity(), start, length);
78        self.length += length.min(other.len().saturating_sub(start));
79    }
80
81    fn subslice_extend_each_repeated(
82        &mut self,
83        other: &FixedSizeListArray,
84        start: usize,
85        length: usize,
86        repeats: usize,
87        share: ShareStrategy,
88    ) {
89        let other_values = &**other.values();
90        self.inner_builder.reserve(repeats * length * self.size);
91        for outer_idx in start..start + length {
92            self.inner_builder.subslice_extend_repeated(
93                other_values,
94                outer_idx * self.size,
95                self.size,
96                repeats,
97                share,
98            );
99        }
100        self.validity
101            .subslice_extend_each_repeated_from_opt_validity(
102                other.validity(),
103                start,
104                length,
105                repeats,
106            );
107        self.length += repeats * length.min(other.len().saturating_sub(start));
108    }
109
110    unsafe fn gather_extend(
111        &mut self,
112        other: &FixedSizeListArray,
113        idxs: &[IdxSize],
114        share: ShareStrategy,
115    ) {
116        let other_values = &**other.values();
117        self.inner_builder.reserve(idxs.len() * self.size);
118
119        // Group consecutive indices into larger copies.
120        let mut group_start = 0;
121        while group_start < idxs.len() {
122            let start_idx = idxs[group_start] as usize;
123            let mut group_len = 1;
124            while group_start + group_len < idxs.len()
125                && idxs[group_start + group_len] as usize == start_idx + group_len
126            {
127                group_len += 1;
128            }
129            self.inner_builder.subslice_extend(
130                other_values,
131                start_idx * self.size,
132                group_len * self.size,
133                share,
134            );
135            group_start += group_len;
136        }
137
138        self.validity
139            .gather_extend_from_opt_validity(other.validity(), idxs);
140        self.length += idxs.len();
141    }
142
143    fn opt_gather_extend(
144        &mut self,
145        other: &FixedSizeListArray,
146        idxs: &[IdxSize],
147        share: ShareStrategy,
148    ) {
149        let other_values = &**other.values();
150        self.inner_builder.reserve(idxs.len() * self.size);
151
152        // Group consecutive indices into larger copies.
153        let mut group_start = 0;
154        while group_start < idxs.len() {
155            let start_idx = idxs[group_start] as usize;
156            let mut group_len = 1;
157            let in_bounds = start_idx < other.len();
158
159            if in_bounds {
160                while group_start + group_len < idxs.len()
161                    && idxs[group_start + group_len] as usize == start_idx + group_len
162                    && start_idx + group_len < other.len()
163                {
164                    group_len += 1;
165                }
166
167                self.inner_builder.subslice_extend(
168                    other_values,
169                    start_idx * self.size,
170                    group_len * self.size,
171                    share,
172                );
173            } else {
174                while group_start + group_len < idxs.len()
175                    && idxs[group_start + group_len] as usize >= other.len()
176                {
177                    group_len += 1;
178                }
179
180                self.inner_builder.extend_nulls(group_len * self.size);
181            }
182            group_start += group_len;
183        }
184
185        self.validity
186            .opt_gather_extend_from_opt_validity(other.validity(), idxs, other.len());
187        self.length += idxs.len();
188    }
189}
polars_arrow/array/fixed_size_list/builder.rs

polars_arrow/array/fixed_size_list/
builder.rs
