polars_arrow/array/fixed_size_list/
mutable.rs

1use std::sync::Arc;
2
3use polars_error::{PolarsResult, polars_bail};
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().is_multiple_of(self.size) {
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    pub fn freeze(mut self) -> FixedSizeListArray {
154        FixedSizeListArray::new(
155            self.dtype,
156            self.length,
157            self.values.as_box(),
158            self.validity.map(|b| b.freeze()),
159        )
160    }
161}
162
163impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {
164    fn len(&self) -> usize {
165        debug_assert!(self.has_valid_invariants());
166        self.length
167    }
168
169    fn validity(&self) -> Option<&MutableBitmap> {
170        self.validity.as_ref()
171    }
172
173    fn as_box(&mut self) -> Box<dyn Array> {
174        FixedSizeListArray::new(
175            self.dtype.clone(),
176            self.length,
177            self.values.as_box(),
178            std::mem::take(&mut self.validity).map(|x| x.into()),
179        )
180        .boxed()
181    }
182
183    fn as_arc(&mut self) -> Arc<dyn Array> {
184        FixedSizeListArray::new(
185            self.dtype.clone(),
186            self.length,
187            self.values.as_box(),
188            std::mem::take(&mut self.validity).map(|x| x.into()),
189        )
190        .arced()
191    }
192
193    fn dtype(&self) -> &ArrowDataType {
194        &self.dtype
195    }
196
197    fn as_any(&self) -> &dyn std::any::Any {
198        self
199    }
200
201    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
202        self
203    }
204
205    #[inline]
206    fn push_null(&mut self) {
207        (0..self.size).for_each(|_| {
208            self.values.push_null();
209        });
210        if let Some(validity) = &mut self.validity {
211            validity.push(false)
212        } else {
213            self.init_validity()
214        }
215        self.length += 1;
216
217        debug_assert!(self.has_valid_invariants());
218    }
219
220    fn reserve(&mut self, additional: usize) {
221        self.reserve(additional)
222    }
223
224    fn shrink_to_fit(&mut self) {
225        self.shrink_to_fit()
226    }
227}
228
229impl<M, I, T> TryExtend<Option<I>> for MutableFixedSizeListArray<M>
230where
231    M: MutableArray + TryExtend<Option<T>>,
232    I: IntoIterator<Item = Option<T>>,
233{
234    #[inline]
235    fn try_extend<II: IntoIterator<Item = Option<I>>>(&mut self, iter: II) -> PolarsResult<()> {
236        for items in iter {
237            self.try_push(items)?;
238        }
239
240        debug_assert!(self.has_valid_invariants());
241
242        Ok(())
243    }
244}
245
246impl<M, I, T> TryPush<Option<I>> for MutableFixedSizeListArray<M>
247where
248    M: MutableArray + TryExtend<Option<T>>,
249    I: IntoIterator<Item = Option<T>>,
250{
251    #[inline]
252    fn try_push(&mut self, item: Option<I>) -> PolarsResult<()> {
253        if let Some(items) = item {
254            self.values.try_extend(items)?;
255            self.try_push_valid()?;
256        } else {
257            self.push_null();
258        }
259
260        debug_assert!(self.has_valid_invariants());
261
262        Ok(())
263    }
264}
265
266impl<M, I, T> PushUnchecked<Option<I>> for MutableFixedSizeListArray<M>
267where
268    M: MutableArray + Extend<Option<T>>,
269    I: IntoIterator<Item = Option<T>>,
270{
271    /// # Safety
272    /// The caller must ensure that the `I` iterates exactly over `size`
273    /// items, where `size` is the fixed size width.
274    #[inline]
275    unsafe fn push_unchecked(&mut self, item: Option<I>) {
276        if let Some(items) = item {
277            self.values.extend(items);
278            self.push_valid();
279        } else {
280            self.push_null();
281        }
282
283        debug_assert!(self.has_valid_invariants());
284    }
285}
286
287impl<M> TryExtendFromSelf for MutableFixedSizeListArray<M>
288where
289    M: MutableArray + TryExtendFromSelf,
290{
291    fn try_extend_from_self(&mut self, other: &Self) -> PolarsResult<()> {
292        extend_validity(self.len(), &mut self.validity, &other.validity);
293
294        self.values.try_extend_from_self(&other.values)?;
295        self.length += other.len();
296
297        debug_assert!(self.has_valid_invariants());
298
299        Ok(())
300    }
301}
polars_arrow/array/fixed_size_list/mutable.rs

polars_arrow/array/fixed_size_list/
mutable.rs
