Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Improved MutableBooleanArray.
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao committed Aug 18, 2021
1 parent 2c7f6f9 commit d2411bd
Show file tree
Hide file tree
Showing 5 changed files with 84 additions and 24 deletions.
30 changes: 29 additions & 1 deletion src/array/boolean/iterator.rs
Original file line number Diff line number Diff line change
@@ -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>;
Expand Down Expand Up @@ -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()
}
}
3 changes: 1 addition & 2 deletions src/array/boolean/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}
22 changes: 17 additions & 5 deletions src/array/boolean/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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`].
Expand Down Expand Up @@ -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())
}
}
21 changes: 16 additions & 5 deletions src/bitmap/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<MutableBitmap> for Bitmap {
Expand Down Expand Up @@ -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)
}
}
32 changes: 21 additions & 11 deletions tests/it/array/boolean/mutable.rs
Original file line number Diff line number Diff line change
@@ -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)]));
}

0 comments on commit d2411bd

Please sign in to comment.