polars_arrow/array/fixed_size_list/
mutable.rs

1use std::sync::Arc;
2
3use polars_error::{polars_bail, PolarsResult};
4use polars_utils::pl_str::PlSmallStr;
5
6use super::FixedSizeListArray;
7use crate::array::physical_binary::extend_validity;
8use crate::array::{Array, MutableArray, PushUnchecked, TryExtend, TryExtendFromSelf, TryPush};
9use crate::bitmap::MutableBitmap;
10use crate::datatypes::{ArrowDataType, Field};
11
12/// The mutable version of [`FixedSizeListArray`].
13#[derive(Debug, Clone)]
14pub struct MutableFixedSizeListArray<M: MutableArray> {
15    dtype: ArrowDataType,
16    size: usize,
17    length: usize,
18    values: M,
19    validity: Option<MutableBitmap>,
20}
21
22impl<M: MutableArray> From<MutableFixedSizeListArray<M>> for FixedSizeListArray {
23    fn from(mut other: MutableFixedSizeListArray<M>) -> Self {
24        FixedSizeListArray::new(
25            other.dtype,
26            other.length,
27            other.values.as_box(),
28            other.validity.map(|x| x.into()),
29        )
30    }
31}
32
33impl<M: MutableArray> MutableFixedSizeListArray<M> {
34    /// Creates a new [`MutableFixedSizeListArray`] from a [`MutableArray`] and size.
35    pub fn new(values: M, size: usize) -> Self {
36        let dtype = FixedSizeListArray::default_datatype(values.dtype().clone(), size);
37        Self::new_from(values, dtype, size)
38    }
39
40    /// Creates a new [`MutableFixedSizeListArray`] from a [`MutableArray`] and size.
41    pub fn new_with_field(values: M, name: PlSmallStr, nullable: bool, size: usize) -> Self {
42        let dtype = ArrowDataType::FixedSizeList(
43            Box::new(Field::new(name, values.dtype().clone(), nullable)),
44            size,
45        );
46        Self::new_from(values, dtype, size)
47    }
48
49    /// Creates a new [`MutableFixedSizeListArray`] from a [`MutableArray`], [`ArrowDataType`] and size.
50    pub fn new_from(values: M, dtype: ArrowDataType, size: usize) -> Self {
51        assert_eq!(values.len(), 0);
52        match dtype {
53            ArrowDataType::FixedSizeList(..) => (),
54            _ => panic!("data type must be FixedSizeList (got {dtype:?})"),
55        };
56        Self {
57            size,
58            length: 0,
59            dtype,
60            values,
61            validity: None,
62        }
63    }
64
65    #[inline]
66    fn has_valid_invariants(&self) -> bool {
67        (self.size == 0 && self.values().len() == 0)
68            || (self.size > 0 && self.values.len() / self.size == self.length)
69    }
70
71    /// Returns the size (number of elements per slot) of this [`FixedSizeListArray`].
72    pub const fn size(&self) -> usize {
73        self.size
74    }
75
76    /// The length of this array
77    pub fn len(&self) -> usize {
78        debug_assert!(self.has_valid_invariants());
79        self.length
80    }
81
82    /// The inner values
83    pub fn values(&self) -> &M {
84        &self.values
85    }
86
87    fn init_validity(&mut self) {
88        let len = self.values.len() / self.size;
89
90        let mut validity = MutableBitmap::new();
91        validity.extend_constant(len, true);
92        validity.set(len - 1, false);
93        self.validity = Some(validity)
94    }
95
96    #[inline]
97    /// Needs to be called when a valid value was extended to this array.
98    /// This is a relatively low level function, prefer `try_push` when you can.
99    pub fn try_push_valid(&mut self) -> PolarsResult<()> {
100        if self.values.len() % self.size != 0 {
101            polars_bail!(ComputeError: "overflow")
102        };
103        if let Some(validity) = &mut self.validity {
104            validity.push(true)
105        }
106        self.length += 1;
107
108        debug_assert!(self.has_valid_invariants());
109
110        Ok(())
111    }
112
113    #[inline]
114    /// Needs to be called when a valid value was extended to this array.
115    /// This is a relatively low level function, prefer `try_push` when you can.
116    pub fn push_valid(&mut self) {
117        if let Some(validity) = &mut self.validity {
118            validity.push(true)
119        }
120        self.length += 1;
121
122        debug_assert!(self.has_valid_invariants());
123    }
124
125    #[inline]
126    fn push_null(&mut self) {
127        (0..self.size).for_each(|_| self.values.push_null());
128        match &mut self.validity {
129            Some(validity) => validity.push(false),
130            None => self.init_validity(),
131        }
132        self.length += 1;
133
134        debug_assert!(self.has_valid_invariants());
135    }
136
137    /// Reserves `additional` slots.
138    pub fn reserve(&mut self, additional: usize) {
139        self.values.reserve(additional);
140        if let Some(x) = self.validity.as_mut() {
141            x.reserve(additional)
142        }
143    }
144
145    /// Shrinks the capacity of the [`MutableFixedSizeListArray`] to fit its current length.
146    pub fn shrink_to_fit(&mut self) {
147        self.values.shrink_to_fit();
148        if let Some(validity) = &mut self.validity {
149            validity.shrink_to_fit()
150        }
151    }
152}
153
154impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {
155    fn len(&self) -> usize {
156        debug_assert!(self.has_valid_invariants());
157        self.length
158    }
159
160    fn validity(&self) -> Option<&MutableBitmap> {
161        self.validity.as_ref()
162    }
163
164    fn as_box(&mut self) -> Box<dyn Array> {
165        FixedSizeListArray::new(
166            self.dtype.clone(),
167            self.length,
168            self.values.as_box(),
169            std::mem::take(&mut self.validity).map(|x| x.into()),
170        )
171        .boxed()
172    }
173
174    fn as_arc(&mut self) -> Arc<dyn Array> {
175        FixedSizeListArray::new(
176            self.dtype.clone(),
177            self.length,
178            self.values.as_box(),
179            std::mem::take(&mut self.validity).map(|x| x.into()),
180        )
181        .arced()
182    }
183
184    fn dtype(&self) -> &ArrowDataType {
185        &self.dtype
186    }
187
188    fn as_any(&self) -> &dyn std::any::Any {
189        self
190    }
191
192    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
193        self
194    }
195
196    #[inline]
197    fn push_null(&mut self) {
198        (0..self.size).for_each(|_| {
199            self.values.push_null();
200        });
201        if let Some(validity) = &mut self.validity {
202            validity.push(false)
203        } else {
204            self.init_validity()
205        }
206        self.length += 1;
207
208        debug_assert!(self.has_valid_invariants());
209    }
210
211    fn reserve(&mut self, additional: usize) {
212        self.reserve(additional)
213    }
214
215    fn shrink_to_fit(&mut self) {
216        self.shrink_to_fit()
217    }
218}
219
220impl<M, I, T> TryExtend<Option<I>> for MutableFixedSizeListArray<M>
221where
222    M: MutableArray + TryExtend<Option<T>>,
223    I: IntoIterator<Item = Option<T>>,
224{
225    #[inline]
226    fn try_extend<II: IntoIterator<Item = Option<I>>>(&mut self, iter: II) -> PolarsResult<()> {
227        for items in iter {
228            self.try_push(items)?;
229        }
230
231        debug_assert!(self.has_valid_invariants());
232
233        Ok(())
234    }
235}
236
237impl<M, I, T> TryPush<Option<I>> for MutableFixedSizeListArray<M>
238where
239    M: MutableArray + TryExtend<Option<T>>,
240    I: IntoIterator<Item = Option<T>>,
241{
242    #[inline]
243    fn try_push(&mut self, item: Option<I>) -> PolarsResult<()> {
244        if let Some(items) = item {
245            self.values.try_extend(items)?;
246            self.try_push_valid()?;
247        } else {
248            self.push_null();
249        }
250
251        debug_assert!(self.has_valid_invariants());
252
253        Ok(())
254    }
255}
256
257impl<M, I, T> PushUnchecked<Option<I>> for MutableFixedSizeListArray<M>
258where
259    M: MutableArray + Extend<Option<T>>,
260    I: IntoIterator<Item = Option<T>>,
261{
262    /// # Safety
263    /// The caller must ensure that the `I` iterates exactly over `size`
264    /// items, where `size` is the fixed size width.
265    #[inline]
266    unsafe fn push_unchecked(&mut self, item: Option<I>) {
267        if let Some(items) = item {
268            self.values.extend(items);
269            self.push_valid();
270        } else {
271            self.push_null();
272        }
273
274        debug_assert!(self.has_valid_invariants());
275    }
276}
277
278impl<M> TryExtendFromSelf for MutableFixedSizeListArray<M>
279where
280    M: MutableArray + TryExtendFromSelf,
281{
282    fn try_extend_from_self(&mut self, other: &Self) -> PolarsResult<()> {
283        extend_validity(self.len(), &mut self.validity, &other.validity);
284
285        self.values.try_extend_from_self(&other.values)?;
286        self.length += other.len();
287
288        debug_assert!(self.has_valid_invariants());
289
290        Ok(())
291    }
292}
polars_arrow/array/fixed_size_list/mutable.rs

polars_arrow/array/fixed_size_list/
mutable.rs
