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    }
108
109    unsafe fn gather_extend(
110        &mut self,
111        other: &FixedSizeListArray,
112        idxs: &[IdxSize],
113        share: ShareStrategy,
114    ) {
115        let other_values = &**other.values();
116        self.inner_builder.reserve(idxs.len() * self.size);
117
118        // Group consecutive indices into larger copies.
119        let mut group_start = 0;
120        while group_start < idxs.len() {
121            let start_idx = idxs[group_start] as usize;
122            let mut group_len = 1;
123            while group_start + group_len < idxs.len()
124                && idxs[group_start + group_len] as usize == start_idx + group_len
125            {
126                group_len += 1;
127            }
128            self.inner_builder.subslice_extend(
129                other_values,
130                start_idx * self.size,
131                group_len * self.size,
132                share,
133            );
134            group_start += group_len;
135        }
136
137        self.validity
138            .gather_extend_from_opt_validity(other.validity(), idxs);
139        self.length += idxs.len();
140    }
141
142    fn opt_gather_extend(
143        &mut self,
144        other: &FixedSizeListArray,
145        idxs: &[IdxSize],
146        share: ShareStrategy,
147    ) {
148        let other_values = &**other.values();
149        self.inner_builder.reserve(idxs.len() * self.size);
150
151        // Group consecutive indices into larger copies.
152        let mut group_start = 0;
153        while group_start < idxs.len() {
154            let start_idx = idxs[group_start] as usize;
155            let mut group_len = 1;
156            let in_bounds = start_idx < other.len();
157
158            if in_bounds {
159                while group_start + group_len < idxs.len()
160                    && idxs[group_start + group_len] as usize == start_idx + group_len
161                    && start_idx + group_len < other.len()
162                {
163                    group_len += 1;
164                }
165
166                self.inner_builder.subslice_extend(
167                    other_values,
168                    start_idx * self.size,
169                    group_len * self.size,
170                    share,
171                );
172            } else {
173                while group_start + group_len < idxs.len()
174                    && idxs[group_start + group_len] as usize >= other.len()
175                {
176                    group_len += 1;
177                }
178
179                self.inner_builder.extend_nulls(group_len * self.size);
180            }
181            group_start += group_len;
182        }
183
184        self.validity
185            .opt_gather_extend_from_opt_validity(other.validity(), idxs, other.len());
186        self.length += idxs.len();
187    }
188}
polars_arrow/array/fixed_size_list/builder.rs

polars_arrow/array/fixed_size_list/
builder.rs
