polars_arrow/array/fixed_size_binary/
mutable.rs

1use std::sync::Arc;
2
3use polars_error::{PolarsResult, polars_bail};
4
5use super::FixedSizeBinaryArray;
6use crate::array::physical_binary::extend_validity;
7use crate::array::{Array, MutableArray, TryExtendFromSelf};
8use crate::bitmap::MutableBitmap;
9use crate::datatypes::ArrowDataType;
10
11/// The Arrow's equivalent to a mutable `Vec<Option<[u8; size]>>`.
12/// Converting a [`MutableFixedSizeBinaryArray`] into a [`FixedSizeBinaryArray`] is `O(1)`.
13/// # Implementation
14/// This struct does not allocate a validity until one is required (i.e. push a null to it).
15#[derive(Debug, Clone)]
16pub struct MutableFixedSizeBinaryArray {
17    dtype: ArrowDataType,
18    size: usize,
19    values: Vec<u8>,
20    validity: Option<MutableBitmap>,
21}
22
23impl From<MutableFixedSizeBinaryArray> for FixedSizeBinaryArray {
24    fn from(other: MutableFixedSizeBinaryArray) -> Self {
25        FixedSizeBinaryArray::new(
26            other.dtype,
27            other.values.into(),
28            other.validity.map(|x| x.into()),
29        )
30    }
31}
32
33impl MutableFixedSizeBinaryArray {
34    /// Creates a new [`MutableFixedSizeBinaryArray`].
35    ///
36    /// # Errors
37    /// This function returns an error iff:
38    /// * The `dtype`'s physical type is not [`crate::datatypes::PhysicalType::FixedSizeBinary`]
39    /// * The length of `values` is not a multiple of `size` in `dtype`
40    /// * the validity's length is not equal to `values.len() / size`.
41    pub fn try_new(
42        dtype: ArrowDataType,
43        values: Vec<u8>,
44        validity: Option<MutableBitmap>,
45    ) -> PolarsResult<Self> {
46        let size = FixedSizeBinaryArray::maybe_get_size(&dtype)?;
47
48        if !values.len().is_multiple_of(size) {
49            polars_bail!(ComputeError:
50                "values (of len {}) must be a multiple of size ({}) in FixedSizeBinaryArray.",
51                values.len(),
52                size
53            )
54        }
55        let len = values.len() / size;
56
57        if validity
58            .as_ref()
59            .is_some_and(|validity| validity.len() != len)
60        {
61            polars_bail!(ComputeError: "validity mask length must be equal to the number of values divided by size")
62        }
63
64        Ok(Self {
65            size,
66            dtype,
67            values,
68            validity,
69        })
70    }
71
72    /// Creates a new empty [`MutableFixedSizeBinaryArray`].
73    pub fn new(size: usize) -> Self {
74        Self::with_capacity(size, 0)
75    }
76
77    /// Creates a new [`MutableFixedSizeBinaryArray`] with capacity for `capacity` entries.
78    pub fn with_capacity(size: usize, capacity: usize) -> Self {
79        Self::try_new(
80            ArrowDataType::FixedSizeBinary(size),
81            Vec::<u8>::with_capacity(capacity * size),
82            None,
83        )
84        .unwrap()
85    }
86
87    /// Creates a new [`MutableFixedSizeBinaryArray`] from a slice of optional `[u8]`.
88    // Note: this can't be `impl From` because Rust does not allow double `AsRef` on it.
89    pub fn from<const N: usize, P: AsRef<[Option<[u8; N]>]>>(slice: P) -> Self {
90        let values = slice
91            .as_ref()
92            .iter()
93            .copied()
94            .flat_map(|x| x.unwrap_or([0; N]))
95            .collect::<Vec<_>>();
96        let validity = slice
97            .as_ref()
98            .iter()
99            .map(|x| x.is_some())
100            .collect::<MutableBitmap>();
101        Self::try_new(ArrowDataType::FixedSizeBinary(N), values, validity.into()).unwrap()
102    }
103
104    /// tries to push a new entry to [`MutableFixedSizeBinaryArray`].
105    /// # Error
106    /// Errors iff the size of `value` is not equal to its own size.
107    #[inline]
108    pub fn try_push<P: AsRef<[u8]>>(&mut self, value: Option<P>) -> PolarsResult<()> {
109        match value {
110            Some(bytes) => {
111                let bytes = bytes.as_ref();
112                if self.size != bytes.len() {
113                    polars_bail!(
114                        ComputeError:
115                        "MutableFixedSizeBinaryArray(row_width: {}): attempted to push \
116                        row with {} bytes",
117                        self.size, bytes.len(),
118                    )
119                }
120                self.values.extend_from_slice(bytes);
121
122                if let Some(validity) = &mut self.validity {
123                    validity.push(true)
124                }
125            },
126            None => {
127                self.values.resize(self.values.len() + self.size, 0);
128                match &mut self.validity {
129                    Some(validity) => validity.push(false),
130                    None => self.init_validity(),
131                }
132            },
133        }
134        Ok(())
135    }
136
137    /// pushes a new entry to [`MutableFixedSizeBinaryArray`].
138    /// # Panics
139    /// Panics iff the size of `value` is not equal to its own size.
140    #[inline]
141    pub fn push<P: AsRef<[u8]>>(&mut self, value: Option<P>) {
142        self.try_push(value).unwrap()
143    }
144
145    /// Returns the length of this array
146    #[inline]
147    pub fn len(&self) -> usize {
148        self.values.len() / self.size
149    }
150
151    /// Pop the last entry from [`MutableFixedSizeBinaryArray`].
152    /// This function returns `None` iff this array is empty
153    pub fn pop(&mut self) -> Option<Vec<u8>> {
154        if self.values.len() < self.size {
155            return None;
156        }
157        let value_start = self.values.len() - self.size;
158        let value = self.values.split_off(value_start);
159        self.validity
160            .as_mut()
161            .map(|x| x.pop()?.then(|| ()))
162            .unwrap_or_else(|| Some(()))
163            .map(|_| value)
164    }
165
166    /// Creates a new [`MutableFixedSizeBinaryArray`] from an iterator of values.
167    /// # Errors
168    /// Errors iff the size of any of the `value` is not equal to its own size.
169    pub fn try_from_iter<P: AsRef<[u8]>, I: IntoIterator<Item = Option<P>>>(
170        iter: I,
171        size: usize,
172    ) -> PolarsResult<Self> {
173        let iterator = iter.into_iter();
174        let (lower, _) = iterator.size_hint();
175        let mut primitive = Self::with_capacity(size, lower);
176        for item in iterator {
177            primitive.try_push(item)?
178        }
179        Ok(primitive)
180    }
181
182    /// returns the (fixed) size of the [`MutableFixedSizeBinaryArray`].
183    #[inline]
184    pub fn size(&self) -> usize {
185        self.size
186    }
187
188    /// Returns the capacity of this array
189    pub fn capacity(&self) -> usize {
190        self.values.capacity() / self.size
191    }
192
193    fn init_validity(&mut self) {
194        let mut validity = MutableBitmap::new();
195        validity.extend_constant(self.len(), true);
196        validity.set(self.len() - 1, false);
197        self.validity = Some(validity)
198    }
199
200    /// Returns the element at index `i` as `&[u8]`
201    #[inline]
202    pub fn value(&self, i: usize) -> &[u8] {
203        &self.values[i * self.size..(i + 1) * self.size]
204    }
205
206    /// Returns the element at index `i` as `&[u8]`
207    ///
208    /// # Safety
209    /// Assumes that the `i < self.len`.
210    #[inline]
211    pub unsafe fn value_unchecked(&self, i: usize) -> &[u8] {
212        std::slice::from_raw_parts(self.values.as_ptr().add(i * self.size), self.size)
213    }
214
215    /// Reserves `additional` slots.
216    pub fn reserve(&mut self, additional: usize) {
217        self.values.reserve(additional * self.size);
218        if let Some(x) = self.validity.as_mut() {
219            x.reserve(additional)
220        }
221    }
222
223    /// Shrinks the capacity of the [`MutableFixedSizeBinaryArray`] to fit its current length.
224    pub fn shrink_to_fit(&mut self) {
225        self.values.shrink_to_fit();
226        if let Some(validity) = &mut self.validity {
227            validity.shrink_to_fit()
228        }
229    }
230
231    pub fn freeze(self) -> FixedSizeBinaryArray {
232        FixedSizeBinaryArray::new(
233            ArrowDataType::FixedSizeBinary(self.size),
234            self.values.into(),
235            self.validity.map(|x| x.into()),
236        )
237    }
238}
239
240/// Accessors
241impl MutableFixedSizeBinaryArray {
242    /// Returns its values.
243    pub fn values(&self) -> &Vec<u8> {
244        &self.values
245    }
246
247    /// Returns a mutable slice of values.
248    pub fn values_mut_slice(&mut self) -> &mut [u8] {
249        self.values.as_mut_slice()
250    }
251}
252
253impl MutableArray for MutableFixedSizeBinaryArray {
254    fn len(&self) -> usize {
255        self.values.len() / self.size
256    }
257
258    fn validity(&self) -> Option<&MutableBitmap> {
259        self.validity.as_ref()
260    }
261
262    fn as_box(&mut self) -> Box<dyn Array> {
263        FixedSizeBinaryArray::new(
264            ArrowDataType::FixedSizeBinary(self.size),
265            std::mem::take(&mut self.values).into(),
266            std::mem::take(&mut self.validity).map(|x| x.into()),
267        )
268        .boxed()
269    }
270
271    fn as_arc(&mut self) -> Arc<dyn Array> {
272        FixedSizeBinaryArray::new(
273            ArrowDataType::FixedSizeBinary(self.size),
274            std::mem::take(&mut self.values).into(),
275            std::mem::take(&mut self.validity).map(|x| x.into()),
276        )
277        .arced()
278    }
279
280    fn dtype(&self) -> &ArrowDataType {
281        &self.dtype
282    }
283
284    fn as_any(&self) -> &dyn std::any::Any {
285        self
286    }
287
288    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
289        self
290    }
291
292    fn push_null(&mut self) {
293        self.push::<&[u8]>(None);
294    }
295
296    fn reserve(&mut self, additional: usize) {
297        self.reserve(additional)
298    }
299
300    fn shrink_to_fit(&mut self) {
301        self.shrink_to_fit()
302    }
303}
304
305impl PartialEq for MutableFixedSizeBinaryArray {
306    fn eq(&self, other: &Self) -> bool {
307        self.iter().eq(other.iter())
308    }
309}
310
311impl TryExtendFromSelf for MutableFixedSizeBinaryArray {
312    fn try_extend_from_self(&mut self, other: &Self) -> PolarsResult<()> {
313        extend_validity(self.len(), &mut self.validity, &other.validity);
314
315        let slice = other.values.as_slice();
316        self.values.extend_from_slice(slice);
317        Ok(())
318    }
319}
polars_arrow/array/fixed_size_binary/mutable.rs

polars_arrow/array/fixed_size_binary/
mutable.rs
