arrow2/array/boolean/
mutable.rs

1use std::iter::FromIterator;
2use std::sync::Arc;
3
4use crate::{
5    array::{
6        physical_binary::extend_validity, Array, MutableArray, TryExtend, TryExtendFromSelf,
7        TryPush,
8    },
9    bitmap::MutableBitmap,
10    datatypes::{DataType, PhysicalType},
11    error::Error,
12    trusted_len::TrustedLen,
13};
14
15use super::BooleanArray;
16
17/// The Arrow's equivalent to `Vec<Option<bool>>`, but with `1/16` of its size.
18/// Converting a [`MutableBooleanArray`] into a [`BooleanArray`] is `O(1)`.
19/// # Implementation
20/// This struct does not allocate a validity until one is required (i.e. push a null to it).
21#[derive(Debug, Clone)]
22pub struct MutableBooleanArray {
23    data_type: DataType,
24    values: MutableBitmap,
25    validity: Option<MutableBitmap>,
26}
27
28impl From<MutableBooleanArray> for BooleanArray {
29    fn from(other: MutableBooleanArray) -> Self {
30        BooleanArray::new(
31            other.data_type,
32            other.values.into(),
33            other.validity.map(|x| x.into()),
34        )
35    }
36}
37
38impl<P: AsRef<[Option<bool>]>> From<P> for MutableBooleanArray {
39    /// Creates a new [`MutableBooleanArray`] out of a slice of Optional `bool`.
40    fn from(slice: P) -> Self {
41        Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref()))
42    }
43}
44
45impl Default for MutableBooleanArray {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl MutableBooleanArray {
52    /// Creates an new empty [`MutableBooleanArray`].
53    pub fn new() -> Self {
54        Self::with_capacity(0)
55    }
56
57    /// The canonical method to create a [`MutableBooleanArray`] out of low-end APIs.
58    /// # Errors
59    /// This function errors iff:
60    /// * The validity is not `None` and its length is different from `values`'s length
61    /// * The `data_type`'s [`PhysicalType`] is not equal to [`PhysicalType::Boolean`].
62    pub fn try_new(
63        data_type: DataType,
64        values: MutableBitmap,
65        validity: Option<MutableBitmap>,
66    ) -> Result<Self, Error> {
67        if validity
68            .as_ref()
69            .map_or(false, |validity| validity.len() != values.len())
70        {
71            return Err(Error::oos(
72                "validity mask length must match the number of values",
73            ));
74        }
75
76        if data_type.to_physical_type() != PhysicalType::Boolean {
77            return Err(Error::oos(
78                "MutableBooleanArray can only be initialized with a DataType whose physical type is Boolean",
79            ));
80        }
81
82        Ok(Self {
83            data_type,
84            values,
85            validity,
86        })
87    }
88
89    /// Creates an new [`MutableBooleanArray`] with a capacity of values.
90    pub fn with_capacity(capacity: usize) -> Self {
91        Self {
92            data_type: DataType::Boolean,
93            values: MutableBitmap::with_capacity(capacity),
94            validity: None,
95        }
96    }
97
98    /// Reserves `additional` slots.
99    pub fn reserve(&mut self, additional: usize) {
100        self.values.reserve(additional);
101        if let Some(x) = self.validity.as_mut() {
102            x.reserve(additional)
103        }
104    }
105
106    /// Pushes a new entry to [`MutableBooleanArray`].
107    pub fn push(&mut self, value: Option<bool>) {
108        match value {
109            Some(value) => {
110                self.values.push(value);
111                match &mut self.validity {
112                    Some(validity) => validity.push(true),
113                    None => {}
114                }
115            }
116            None => {
117                self.values.push(false);
118                match &mut self.validity {
119                    Some(validity) => validity.push(false),
120                    None => self.init_validity(),
121                }
122            }
123        }
124    }
125
126    /// Pop an entry from [`MutableBooleanArray`].
127    /// Note If the values is empty, this method will return None.
128    pub fn pop(&mut self) -> Option<bool> {
129        let value = self.values.pop()?;
130        self.validity
131            .as_mut()
132            .map(|x| x.pop()?.then(|| value))
133            .unwrap_or_else(|| Some(value))
134    }
135
136    /// Extends the [`MutableBooleanArray`] from an iterator of values of trusted len.
137    /// This differs from `extend_trusted_len` which accepts in iterator of optional values.
138    #[inline]
139    pub fn extend_trusted_len_values<I>(&mut self, iterator: I)
140    where
141        I: TrustedLen<Item = bool>,
142    {
143        // Safety: `I` is `TrustedLen`
144        unsafe { self.extend_trusted_len_values_unchecked(iterator) }
145    }
146
147    /// Extends the [`MutableBooleanArray`] from an iterator of values of trusted len.
148    /// This differs from `extend_trusted_len_unchecked`, which accepts in iterator of optional values.
149    /// # Safety
150    /// The iterator must be trusted len.
151    #[inline]
152    pub unsafe fn extend_trusted_len_values_unchecked<I>(&mut self, iterator: I)
153    where
154        I: Iterator<Item = bool>,
155    {
156        let (_, upper) = iterator.size_hint();
157        let additional =
158            upper.expect("extend_trusted_len_values_unchecked requires an upper limit");
159
160        if let Some(validity) = self.validity.as_mut() {
161            validity.extend_constant(additional, true);
162        }
163
164        self.values.extend_from_trusted_len_iter_unchecked(iterator)
165    }
166
167    /// Extends the [`MutableBooleanArray`] from an iterator of trusted len.
168    #[inline]
169    pub fn extend_trusted_len<I, P>(&mut self, iterator: I)
170    where
171        P: std::borrow::Borrow<bool>,
172        I: TrustedLen<Item = Option<P>>,
173    {
174        // Safety: `I` is `TrustedLen`
175        unsafe { self.extend_trusted_len_unchecked(iterator) }
176    }
177
178    /// Extends the [`MutableBooleanArray`] from an iterator of trusted len.
179    /// # Safety
180    /// The iterator must be trusted len.
181    #[inline]
182    pub unsafe fn extend_trusted_len_unchecked<I, P>(&mut self, iterator: I)
183    where
184        P: std::borrow::Borrow<bool>,
185        I: Iterator<Item = Option<P>>,
186    {
187        if let Some(validity) = self.validity.as_mut() {
188            extend_trusted_len_unzip(iterator, validity, &mut self.values);
189        } else {
190            let mut validity = MutableBitmap::new();
191            validity.extend_constant(self.len(), true);
192
193            extend_trusted_len_unzip(iterator, &mut validity, &mut self.values);
194
195            if validity.unset_bits() > 0 {
196                self.validity = Some(validity);
197            }
198        }
199    }
200
201    fn init_validity(&mut self) {
202        let mut validity = MutableBitmap::with_capacity(self.values.capacity());
203        validity.extend_constant(self.len(), true);
204        validity.set(self.len() - 1, false);
205        self.validity = Some(validity)
206    }
207
208    /// Converts itself into an [`Array`].
209    pub fn into_arc(self) -> Arc<dyn Array> {
210        let a: BooleanArray = self.into();
211        Arc::new(a)
212    }
213}
214
215/// Getters
216impl MutableBooleanArray {
217    /// Returns its values.
218    pub fn values(&self) -> &MutableBitmap {
219        &self.values
220    }
221}
222
223/// Setters
224impl MutableBooleanArray {
225    /// Sets position `index` to `value`.
226    /// Note that if it is the first time a null appears in this array,
227    /// this initializes the validity bitmap (`O(N)`).
228    /// # Panic
229    /// Panics iff index is larger than `self.len()`.
230    pub fn set(&mut self, index: usize, value: Option<bool>) {
231        self.values.set(index, value.unwrap_or_default());
232
233        if value.is_none() && self.validity.is_none() {
234            // When the validity is None, all elements so far are valid. When one of the elements is set fo null,
235            // the validity must be initialized.
236            self.validity = Some(MutableBitmap::from_trusted_len_iter(
237                std::iter::repeat(true).take(self.len()),
238            ));
239        }
240        if let Some(x) = self.validity.as_mut() {
241            x.set(index, value.is_some())
242        }
243    }
244}
245
246/// From implementations
247impl MutableBooleanArray {
248    /// Creates a new [`MutableBooleanArray`] from an [`TrustedLen`] of `bool`.
249    #[inline]
250    pub fn from_trusted_len_values_iter<I: TrustedLen<Item = bool>>(iterator: I) -> Self {
251        Self::try_new(
252            DataType::Boolean,
253            MutableBitmap::from_trusted_len_iter(iterator),
254            None,
255        )
256        .unwrap()
257    }
258
259    /// Creates a new [`MutableBooleanArray`] from an [`TrustedLen`] of `bool`.
260    /// Use this over [`BooleanArray::from_trusted_len_iter`] when the iterator is trusted len
261    /// but this crate does not mark it as such.
262    /// # Safety
263    /// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
264    /// I.e. that `size_hint().1` correctly reports its length.
265    #[inline]
266    pub unsafe fn from_trusted_len_values_iter_unchecked<I: Iterator<Item = bool>>(
267        iterator: I,
268    ) -> Self {
269        let mut mutable = MutableBitmap::new();
270        mutable.extend_from_trusted_len_iter_unchecked(iterator);
271        MutableBooleanArray::try_new(DataType::Boolean, mutable, None).unwrap()
272    }
273
274    /// Creates a new [`MutableBooleanArray`] from a slice of `bool`.
275    #[inline]
276    pub fn from_slice<P: AsRef<[bool]>>(slice: P) -> Self {
277        Self::from_trusted_len_values_iter(slice.as_ref().iter().copied())
278    }
279
280    /// Creates a [`BooleanArray`] from an iterator of trusted length.
281    /// Use this over [`BooleanArray::from_trusted_len_iter`] when the iterator is trusted len
282    /// but this crate does not mark it as such.
283    /// # Safety
284    /// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
285    /// I.e. that `size_hint().1` correctly reports its length.
286    #[inline]
287    pub unsafe fn from_trusted_len_iter_unchecked<I, P>(iterator: I) -> Self
288    where
289        P: std::borrow::Borrow<bool>,
290        I: Iterator<Item = Option<P>>,
291    {
292        let (validity, values) = trusted_len_unzip(iterator);
293
294        Self::try_new(DataType::Boolean, values, validity).unwrap()
295    }
296
297    /// Creates a [`BooleanArray`] from a [`TrustedLen`].
298    #[inline]
299    pub fn from_trusted_len_iter<I, P>(iterator: I) -> Self
300    where
301        P: std::borrow::Borrow<bool>,
302        I: TrustedLen<Item = Option<P>>,
303    {
304        // Safety: `I` is `TrustedLen`
305        unsafe { Self::from_trusted_len_iter_unchecked(iterator) }
306    }
307
308    /// Creates a [`BooleanArray`] from an falible iterator of trusted length.
309    /// # Safety
310    /// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
311    /// I.e. that `size_hint().1` correctly reports its length.
312    #[inline]
313    pub unsafe fn try_from_trusted_len_iter_unchecked<E, I, P>(
314        iterator: I,
315    ) -> std::result::Result<Self, E>
316    where
317        P: std::borrow::Borrow<bool>,
318        I: Iterator<Item = std::result::Result<Option<P>, E>>,
319    {
320        let (validity, values) = try_trusted_len_unzip(iterator)?;
321
322        let validity = if validity.unset_bits() > 0 {
323            Some(validity)
324        } else {
325            None
326        };
327
328        Ok(Self::try_new(DataType::Boolean, values, validity).unwrap())
329    }
330
331    /// Creates a [`BooleanArray`] from a [`TrustedLen`].
332    #[inline]
333    pub fn try_from_trusted_len_iter<E, I, P>(iterator: I) -> std::result::Result<Self, E>
334    where
335        P: std::borrow::Borrow<bool>,
336        I: TrustedLen<Item = std::result::Result<Option<P>, E>>,
337    {
338        // Safety: `I` is `TrustedLen`
339        unsafe { Self::try_from_trusted_len_iter_unchecked(iterator) }
340    }
341
342    /// Shrinks the capacity of the [`MutableBooleanArray`] to fit its current length.
343    pub fn shrink_to_fit(&mut self) {
344        self.values.shrink_to_fit();
345        if let Some(validity) = &mut self.validity {
346            validity.shrink_to_fit()
347        }
348    }
349}
350
351/// Creates a Bitmap and an optional [`MutableBitmap`] from an iterator of `Option<bool>`.
352/// The first buffer corresponds to a bitmap buffer, the second one
353/// corresponds to a values buffer.
354/// # Safety
355/// The caller must ensure that `iterator` is `TrustedLen`.
356#[inline]
357pub(crate) unsafe fn trusted_len_unzip<I, P>(iterator: I) -> (Option<MutableBitmap>, MutableBitmap)
358where
359    P: std::borrow::Borrow<bool>,
360    I: Iterator<Item = Option<P>>,
361{
362    let mut validity = MutableBitmap::new();
363    let mut values = MutableBitmap::new();
364
365    extend_trusted_len_unzip(iterator, &mut validity, &mut values);
366
367    let validity = if validity.unset_bits() > 0 {
368        Some(validity)
369    } else {
370        None
371    };
372
373    (validity, values)
374}
375
376/// Extends validity [`MutableBitmap`] and values [`MutableBitmap`] from an iterator of `Option`.
377/// # Safety
378/// The caller must ensure that `iterator` is `TrustedLen`.
379#[inline]
380pub(crate) unsafe fn extend_trusted_len_unzip<I, P>(
381    iterator: I,
382    validity: &mut MutableBitmap,
383    values: &mut MutableBitmap,
384) where
385    P: std::borrow::Borrow<bool>,
386    I: Iterator<Item = Option<P>>,
387{
388    let (_, upper) = iterator.size_hint();
389    let additional = upper.expect("extend_trusted_len_unzip requires an upper limit");
390
391    // Length of the array before new values are pushed,
392    // variable created for assertion post operation
393    let pre_length = values.len();
394
395    validity.reserve(additional);
396    values.reserve(additional);
397
398    for item in iterator {
399        let item = if let Some(item) = item {
400            validity.push_unchecked(true);
401            *item.borrow()
402        } else {
403            validity.push_unchecked(false);
404            bool::default()
405        };
406        values.push_unchecked(item);
407    }
408
409    debug_assert_eq!(
410        values.len(),
411        pre_length + additional,
412        "Trusted iterator length was not accurately reported"
413    );
414}
415
416/// # Safety
417/// The caller must ensure that `iterator` is `TrustedLen`.
418#[inline]
419pub(crate) unsafe fn try_trusted_len_unzip<E, I, P>(
420    iterator: I,
421) -> std::result::Result<(MutableBitmap, MutableBitmap), E>
422where
423    P: std::borrow::Borrow<bool>,
424    I: Iterator<Item = std::result::Result<Option<P>, E>>,
425{
426    let (_, upper) = iterator.size_hint();
427    let len = upper.expect("trusted_len_unzip requires an upper limit");
428
429    let mut null = MutableBitmap::with_capacity(len);
430    let mut values = MutableBitmap::with_capacity(len);
431
432    for item in iterator {
433        let item = if let Some(item) = item? {
434            null.push(true);
435            *item.borrow()
436        } else {
437            null.push(false);
438            false
439        };
440        values.push(item);
441    }
442    assert_eq!(
443        values.len(),
444        len,
445        "Trusted iterator length was not accurately reported"
446    );
447    values.set_len(len);
448    null.set_len(len);
449
450    Ok((null, values))
451}
452
453impl<Ptr: std::borrow::Borrow<Option<bool>>> FromIterator<Ptr> for MutableBooleanArray {
454    fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
455        let iter = iter.into_iter();
456        let (lower, _) = iter.size_hint();
457
458        let mut validity = MutableBitmap::with_capacity(lower);
459
460        let values: MutableBitmap = iter
461            .map(|item| {
462                if let Some(a) = item.borrow() {
463                    validity.push(true);
464                    *a
465                } else {
466                    validity.push(false);
467                    false
468                }
469            })
470            .collect();
471
472        let validity = if validity.unset_bits() > 0 {
473            Some(validity)
474        } else {
475            None
476        };
477
478        MutableBooleanArray::try_new(DataType::Boolean, values, validity).unwrap()
479    }
480}
481
482impl MutableArray for MutableBooleanArray {
483    fn len(&self) -> usize {
484        self.values.len()
485    }
486
487    fn validity(&self) -> Option<&MutableBitmap> {
488        self.validity.as_ref()
489    }
490
491    fn as_box(&mut self) -> Box<dyn Array> {
492        let array: BooleanArray = std::mem::take(self).into();
493        array.boxed()
494    }
495
496    fn as_arc(&mut self) -> Arc<dyn Array> {
497        let array: BooleanArray = std::mem::take(self).into();
498        array.arced()
499    }
500
501    fn data_type(&self) -> &DataType {
502        &self.data_type
503    }
504
505    fn as_any(&self) -> &dyn std::any::Any {
506        self
507    }
508
509    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
510        self
511    }
512
513    #[inline]
514    fn push_null(&mut self) {
515        self.push(None)
516    }
517
518    fn reserve(&mut self, additional: usize) {
519        self.reserve(additional)
520    }
521
522    fn shrink_to_fit(&mut self) {
523        self.shrink_to_fit()
524    }
525}
526
527impl Extend<Option<bool>> for MutableBooleanArray {
528    fn extend<I: IntoIterator<Item = Option<bool>>>(&mut self, iter: I) {
529        let iter = iter.into_iter();
530        self.reserve(iter.size_hint().0);
531        iter.for_each(|x| self.push(x))
532    }
533}
534
535impl TryExtend<Option<bool>> for MutableBooleanArray {
536    /// This is infalible and is implemented for consistency with all other types
537    fn try_extend<I: IntoIterator<Item = Option<bool>>>(&mut self, iter: I) -> Result<(), Error> {
538        self.extend(iter);
539        Ok(())
540    }
541}
542
543impl TryPush<Option<bool>> for MutableBooleanArray {
544    /// This is infalible and is implemented for consistency with all other types
545    fn try_push(&mut self, item: Option<bool>) -> Result<(), Error> {
546        self.push(item);
547        Ok(())
548    }
549}
550
551impl PartialEq for MutableBooleanArray {
552    fn eq(&self, other: &Self) -> bool {
553        self.iter().eq(other.iter())
554    }
555}
556
557impl TryExtendFromSelf for MutableBooleanArray {
558    fn try_extend_from_self(&mut self, other: &Self) -> Result<(), Error> {
559        extend_validity(self.len(), &mut self.validity, &other.validity);
560
561        let slice = other.values.as_slice();
562        // safety: invariant offset + length <= slice.len()
563        unsafe {
564            self.values
565                .extend_from_slice_unchecked(slice, 0, other.values.len());
566        }
567        Ok(())
568    }
569}
arrow2/array/boolean/mutable.rs

arrow2/array/boolean/
mutable.rs
