Improved MutableBooleanArray

src/array/boolean/iterator.rs

@@ -1,6 +1,7 @@
 use crate::bitmap::utils::{zip_validity, BitmapIter, ZipValidity};
 
-use super::BooleanArray;
+use super::super::MutableArray;
+use super::{BooleanArray, MutableBooleanArray};
 
 impl<'a> IntoIterator for &'a BooleanArray {
     type Item = Option<bool>;
@@ -28,3 +29,30 @@ impl<'a> BooleanArray {
         self.values().iter()
     }
 }
+
+impl<'a> IntoIterator for &'a MutableBooleanArray {
+    type Item = Option<bool>;
+    type IntoIter = ZipValidity<'a, bool, BitmapIter<'a>>;
+
+    #[inline]
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+impl<'a> MutableBooleanArray {
+    /// Returns an iterator over `Option<bool>`
+    #[inline]
+    pub fn iter(&'a self) -> ZipValidity<'a, bool, BitmapIter<'a>> {
+        zip_validity(
+            self.values().iter(),
+            self.validity().as_ref().map(|x| x.iter()),
+        )
+    }
+
+    /// Returns an iterator of `bool`
+    #[inline]
+    pub fn values_iter(&'a self) -> BitmapIter<'a> {
+        self.values().iter()
+    }
+}

src/array/boolean/mod.rs

@@ -125,8 +125,7 @@ impl std::fmt::Display for BooleanArray {
 
 impl<P: AsRef<[Option<bool>]>> From<P> for BooleanArray {
     /// Creates a new [`BooleanArray`] out of a slice of Optional `bool`.
-    #[inline]
     fn from(slice: P) -> Self {
-        Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref()))
+        MutableBooleanArray::from(slice).into()
     }
 }

src/array/boolean/mutable.rs

@@ -25,6 +25,13 @@ impl From<MutableBooleanArray> for BooleanArray {
     }
 }
 
+impl<P: AsRef<[Option<bool>]>> From<P> for MutableBooleanArray {
+    /// Creates a new [`MutableBooleanArray`] out of a slice of Optional `bool`.
+    fn from(slice: P) -> Self {
+        Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref()))
+    }
+}
+
 impl Default for MutableBooleanArray {
     fn default() -> Self {
         Self::new()
@@ -74,11 +81,10 @@ impl MutableBooleanArray {
     }
 
     fn init_validity(&mut self) {
-        self.validity = Some(MutableBitmap::from_trusted_len_iter(
-            std::iter::repeat(true)
-                .take(self.len() - 1)
-                .chain(std::iter::once(false)),
-        ))
+        let mut validity = MutableBitmap::new();
+        validity.extend_constant(self.len(), true);
+        validity.set(self.len() - 1, false);
+        self.validity = Some(validity)
     }
 
     /// Converts itself into an [`Array`].
@@ -360,3 +366,9 @@ impl TryExtend<Option<bool>> for MutableBooleanArray {
         Ok(())
     }
 }
+
+impl PartialEq for MutableBooleanArray {
+    fn eq(&self, other: &Self) -> bool {
+        self.iter().eq(other.iter())
+    }
+}

src/bitmap/mutable.rs

@@ -78,10 +78,8 @@ impl MutableBitmap {
         if self.length % 8 == 0 {
             self.buffer.push(0);
         }
-        if value {
-            let byte = self.buffer.as_mut_slice().last_mut().unwrap();
-            *byte = set(*byte, self.length % 8, true);
-        };
+        let byte = self.buffer.as_mut_slice().last_mut().unwrap();
+        *byte = set(*byte, self.length % 8, value);
         self.length += 1;
     }
 
@@ -217,11 +215,6 @@ impl MutableBitmap {
         assert!(length <= buffer.len() * 8);
         Self { buffer, length }
     }
-
-    /// Returns an iterator over the values of the [`MutableBitmap`].
-    fn iter(&self) -> BitmapIter {
-        BitmapIter::new(&self.buffer, 0, self.len())
-    }
 }
 
 impl From<MutableBitmap> for Bitmap {
@@ -546,3 +539,19 @@ impl Default for MutableBitmap {
         Self::new()
     }
 }
+
+impl<'a> IntoIterator for &'a MutableBitmap {
+    type Item = bool;
+    type IntoIter = BitmapIter<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        BitmapIter::<'a>::new(&self.buffer, 0, self.length)
+    }
+}
+
+impl<'a> MutableBitmap {
+    /// constructs a new iterator over the values of [`MutableBitmap`].
+    pub fn iter(&'a self) -> BitmapIter<'a> {
+        BitmapIter::<'a>::new(&self.buffer, 0, self.length)
+    }
+}

tests/it/array/boolean/mutable.rs

@@ -1,30 +1,40 @@
-use arrow2::{
-    array::{MutableArray, MutableBooleanArray},
-    bitmap::MutableBitmap,
-};
+use arrow2::array::{MutableArray, MutableBooleanArray};
+
+#[test]
+fn set() {
+    let mut a = MutableBooleanArray::from(&[Some(false), Some(true), Some(false)]);
+
+    a.set(1, None);
+    a.set(0, Some(true));
+    assert_eq!(
+        a,
+        MutableBooleanArray::from([Some(true), None, Some(false)])
+    )
+}
 
 #[test]
 fn push() {
     let mut a = MutableBooleanArray::new();
     a.push(Some(true));
+    a.push(Some(false));
     a.push(None);
-    assert_eq!(a.len(), 2);
-    assert!(a.is_valid(0));
-    assert!(!a.is_valid(1));
-
-    assert_eq!(a.values(), &MutableBitmap::from([true, false]));
+    a.push_null();
+    assert_eq!(
+        a,
+        MutableBooleanArray::from([Some(true), Some(false), None, None])
+    );
 }
 
 #[test]
 fn from_trusted_len_iter() {
     let iter = std::iter::repeat(true).take(2).map(Some);
     let a = MutableBooleanArray::from_trusted_len_iter(iter);
-    assert_eq!(a.len(), 2);
+    assert_eq!(a, MutableBooleanArray::from([Some(true), Some(true)]));
 }
 
 #[test]
 fn from_iter() {
     let iter = std::iter::repeat(true).take(2).map(Some);
     let a: MutableBooleanArray = iter.collect();
-    assert_eq!(a.len(), 2);
+    assert_eq!(a, MutableBooleanArray::from([Some(true), Some(true)]));
 }