From d2411bdc1bd5bcf67b0e418eca69fe87578a5c7d Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Wed, 18 Aug 2021 08:27:57 +0000 Subject: [PATCH 1/2] Improved MutableBooleanArray. --- src/array/boolean/iterator.rs | 30 ++++++++++++++++++++++++++++- src/array/boolean/mod.rs | 3 +-- src/array/boolean/mutable.rs | 22 ++++++++++++++++----- src/bitmap/mutable.rs | 21 +++++++++++++++----- tests/it/array/boolean/mutable.rs | 32 ++++++++++++++++++++----------- 5 files changed, 84 insertions(+), 24 deletions(-) diff --git a/src/array/boolean/iterator.rs b/src/array/boolean/iterator.rs index 9b2d3821da8..d5e0b0c7731 100644 --- a/src/array/boolean/iterator.rs +++ b/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; @@ -28,3 +29,30 @@ impl<'a> BooleanArray { self.values().iter() } } + +impl<'a> IntoIterator for &'a MutableBooleanArray { + type Item = Option; + type IntoIter = ZipValidity<'a, bool, BitmapIter<'a>>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl<'a> MutableBooleanArray { + /// Returns an iterator over `Option` + #[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() + } +} diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs index 212bc01163c..43c310b3004 100644 --- a/src/array/boolean/mod.rs +++ b/src/array/boolean/mod.rs @@ -125,8 +125,7 @@ impl std::fmt::Display for BooleanArray { impl]>> From

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() } } diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs index 44231cc89a2..d5099a48205 100644 --- a/src/array/boolean/mutable.rs +++ b/src/array/boolean/mutable.rs @@ -25,6 +25,13 @@ impl From for BooleanArray { } } +impl]>> From

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> for MutableBooleanArray { Ok(()) } } + +impl PartialEq for MutableBooleanArray { + fn eq(&self, other: &Self) -> bool { + self.iter().eq(other.iter()) + } +} diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs index e97934b8d1e..236e9992220 100644 --- a/src/bitmap/mutable.rs +++ b/src/bitmap/mutable.rs @@ -217,11 +217,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 for Bitmap { @@ -546,3 +541,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) + } +} diff --git a/tests/it/array/boolean/mutable.rs b/tests/it/array/boolean/mutable.rs index 9408c1e89de..4ea81b15ffd 100644 --- a/tests/it/array/boolean/mutable.rs +++ b/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)])); } From 29ee8982b4a97260a34af1e2a1b6ab9c6fb5bdd5 Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Wed, 18 Aug 2021 13:25:04 +0000 Subject: [PATCH 2/2] Fixed --- src/bitmap/mutable.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs index 236e9992220..eed57fcdfcd 100644 --- a/src/bitmap/mutable.rs +++ b/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; }