arrow2/array/fixed_size_list/
mutable.rs

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

arrow2/array/fixed_size_list/
mutable.rs
