From 14ed0b5c3aeaae7472a69a20afb36c163500327b Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Wed, 18 Aug 2021 08:04:40 +0000 Subject: [PATCH] Generalized ZipIterator. --- src/array/binary/iterator.rs | 5 +++- src/array/boolean/iterator.rs | 5 +++- src/array/dictionary/iterator.rs | 5 +++- src/array/fixed_size_binary/iterator.rs | 5 +++- src/array/fixed_size_list/iterator.rs | 5 +++- src/array/list/iterator.rs | 5 +++- src/array/primitive/iterator.rs | 5 +++- src/array/utf8/iterator.rs | 5 +++- src/bitmap/utils/zip_validity.rs | 18 +++++------- tests/it/bitmap/utils/zip_validity.rs | 38 ++++++++++++------------- 10 files changed, 58 insertions(+), 38 deletions(-) diff --git a/src/array/binary/iterator.rs b/src/array/binary/iterator.rs index d0b116f17f5..7d2679798a6 100644 --- a/src/array/binary/iterator.rs +++ b/src/array/binary/iterator.rs @@ -52,7 +52,10 @@ impl<'a, O: Offset> IntoIterator for &'a BinaryArray { impl<'a, O: Offset> BinaryArray { /// Returns an iterator of `Option<&[u8]>` pub fn iter(&'a self) -> ZipValidity<'a, &'a [u8], BinaryValueIter<'a, O>> { - zip_validity(BinaryValueIter::new(self), &self.validity) + zip_validity( + BinaryValueIter::new(self), + self.validity.as_ref().map(|x| x.iter()), + ) } /// Returns an iterator of `&[u8]` diff --git a/src/array/boolean/iterator.rs b/src/array/boolean/iterator.rs index 8cc334b6d47..9b2d3821da8 100644 --- a/src/array/boolean/iterator.rs +++ b/src/array/boolean/iterator.rs @@ -16,7 +16,10 @@ impl<'a> BooleanArray { /// constructs a new iterator #[inline] pub fn iter(&'a self) -> ZipValidity<'a, bool, BitmapIter<'a>> { - zip_validity(self.values().iter(), &self.validity) + zip_validity( + self.values().iter(), + self.validity.as_ref().map(|x| x.iter()), + ) } /// Returns an iterator of `bool` diff --git a/src/array/dictionary/iterator.rs b/src/array/dictionary/iterator.rs index 12530cf96fd..ae7d6951dd5 100644 --- a/src/array/dictionary/iterator.rs +++ b/src/array/dictionary/iterator.rs @@ -70,7 +70,10 @@ impl<'a, K: DictionaryKey> IntoIterator for &'a DictionaryArray { impl<'a, K: DictionaryKey> DictionaryArray { /// Returns an iterator of `Option>` pub fn iter(&'a self) -> ZipIter<'a, K> { - zip_validity(DictionaryValuesIter::new(self), self.keys.validity()) + zip_validity( + DictionaryValuesIter::new(self), + self.keys.validity().as_ref().map(|x| x.iter()), + ) } /// Returns an iterator of `Box` diff --git a/src/array/fixed_size_binary/iterator.rs b/src/array/fixed_size_binary/iterator.rs index 91781f0011d..b24117fea75 100644 --- a/src/array/fixed_size_binary/iterator.rs +++ b/src/array/fixed_size_binary/iterator.rs @@ -50,6 +50,9 @@ impl<'a> IntoIterator for &'a FixedSizeBinaryArray { impl<'a> FixedSizeBinaryArray { /// constructs a new iterator pub fn iter(&'a self) -> ZipValidity<'a, &'a [u8], FixedSizeBinaryValuesIter<'a>> { - zip_validity(FixedSizeBinaryValuesIter::new(self), &self.validity) + zip_validity( + FixedSizeBinaryValuesIter::new(self), + self.validity.as_ref().map(|x| x.iter()), + ) } } diff --git a/src/array/fixed_size_list/iterator.rs b/src/array/fixed_size_list/iterator.rs index 883144dbd01..031f3f15f48 100644 --- a/src/array/fixed_size_list/iterator.rs +++ b/src/array/fixed_size_list/iterator.rs @@ -26,7 +26,10 @@ impl<'a> IntoIterator for &'a FixedSizeListArray { impl<'a> FixedSizeListArray { /// Returns an iterator of `Option>` pub fn iter(&'a self) -> ZipIter<'a> { - zip_validity(ListValuesIter::new(self), &self.validity) + zip_validity( + ListValuesIter::new(self), + self.validity.as_ref().map(|x| x.iter()), + ) } /// Returns an iterator of `Box` diff --git a/src/array/list/iterator.rs b/src/array/list/iterator.rs index b6c184fb973..7b55ff03135 100644 --- a/src/array/list/iterator.rs +++ b/src/array/list/iterator.rs @@ -76,7 +76,10 @@ impl<'a, O: Offset> IntoIterator for &'a ListArray { impl<'a, O: Offset> ListArray { /// Returns an iterator of `Option>` pub fn iter(&'a self) -> ZipIter<'a, O> { - zip_validity(ListValuesIter::new(self), &self.validity) + zip_validity( + ListValuesIter::new(self), + self.validity.as_ref().map(|x| x.iter()), + ) } /// Returns an iterator of `Box` diff --git a/src/array/primitive/iterator.rs b/src/array/primitive/iterator.rs index 281fd99ec59..2d878baa595 100644 --- a/src/array/primitive/iterator.rs +++ b/src/array/primitive/iterator.rs @@ -19,6 +19,9 @@ impl<'a, T: NativeType> PrimitiveArray { /// constructs a new iterator #[inline] pub fn iter(&'a self) -> ZipValidity<'a, &'a T, std::slice::Iter<'a, T>> { - zip_validity(self.values().iter(), &self.validity) + zip_validity( + self.values().iter(), + self.validity.as_ref().map(|x| x.iter()), + ) } } diff --git a/src/array/utf8/iterator.rs b/src/array/utf8/iterator.rs index 5fe1b285f31..fd6f5d01dfa 100644 --- a/src/array/utf8/iterator.rs +++ b/src/array/utf8/iterator.rs @@ -68,7 +68,10 @@ impl<'a, O: Offset> IntoIterator for &'a Utf8Array { impl<'a, O: Offset> Utf8Array { /// Returns an iterator of `Option<&str>` pub fn iter(&'a self) -> ZipValidity<'a, &'a str, Utf8ValuesIter<'a, O>> { - zip_validity(Utf8ValuesIter::new(self), &self.validity) + zip_validity( + Utf8ValuesIter::new(self), + self.validity.as_ref().map(|x| x.iter()), + ) } /// Returns an iterator of `&str` diff --git a/src/bitmap/utils/zip_validity.rs b/src/bitmap/utils/zip_validity.rs index 615d749ce28..b23983aff17 100644 --- a/src/bitmap/utils/zip_validity.rs +++ b/src/bitmap/utils/zip_validity.rs @@ -1,4 +1,4 @@ -use crate::{bitmap::Bitmap, trusted_len::TrustedLen}; +use crate::trusted_len::TrustedLen; use super::BitmapIter; @@ -21,17 +21,14 @@ impl<'a, T, I: Iterator + Clone> Clone for ZipValidity<'a, T, I> { } impl<'a, T, I: Iterator> ZipValidity<'a, T, I> { - #[inline] - pub fn new(values: I, validity: &'a Option) -> Self { - let validity_iter = validity - .as_ref() - .map(|x| x.iter()) - .unwrap_or_else(|| BitmapIter::new(&[], 0, 0)); + pub fn new(values: I, validity: Option>) -> Self { + let has_validity = validity.as_ref().is_some(); + let validity_iter = validity.unwrap_or_else(|| BitmapIter::new(&[], 0, 0)); Self { values, validity_iter, - has_validity: validity.is_some(), + has_validity, } } } @@ -75,10 +72,9 @@ impl<'a, T, I: Iterator> std::iter::ExactSizeIterator for ZipValidity< unsafe impl> TrustedLen for ZipValidity<'_, T, I> {} /// Returns an iterator adapter that returns Option according to an optional validity. -#[inline] pub fn zip_validity>( values: I, - validity: &Option, + validity: Option, ) -> ZipValidity { - ZipValidity::::new(values, validity) + ZipValidity::new(values, validity) } diff --git a/tests/it/bitmap/utils/zip_validity.rs b/tests/it/bitmap/utils/zip_validity.rs index b49ec623810..d778fb001b0 100644 --- a/tests/it/bitmap/utils/zip_validity.rs +++ b/tests/it/bitmap/utils/zip_validity.rs @@ -2,9 +2,10 @@ use arrow2::bitmap::{utils::zip_validity, Bitmap}; #[test] fn basic() { - let a = Some(Bitmap::from([true, false])); + let a = Bitmap::from([true, false]); + let a = Some(a.iter()); let values = vec![0, 1]; - let zip = zip_validity(values.into_iter(), &a); + let zip = zip_validity(values.into_iter(), a); let a = zip.collect::>(); assert_eq!(a, vec![Some(0), None]); @@ -12,11 +13,10 @@ fn basic() { #[test] fn complete() { - let a = Some(Bitmap::from([ - true, false, true, false, true, false, true, false, - ])); + let a = Bitmap::from([true, false, true, false, true, false, true, false]); + let a = Some(a.iter()); let values = vec![0, 1, 2, 3, 4, 5, 6, 7]; - let zip = zip_validity(values.into_iter(), &a); + let zip = zip_validity(values.into_iter(), a); let a = zip.collect::>(); assert_eq!( @@ -27,7 +27,8 @@ fn complete() { #[test] fn slices() { - let a = Some(Bitmap::from([true, false])); + let a = Bitmap::from([true, false]); + let a = Some(a.iter()); let offsets = vec![0, 2, 3]; let values = vec![1, 2, 3]; let iter = offsets.windows(2).map(|x| { @@ -35,7 +36,7 @@ fn slices() { let end = x[1]; &values[start..end] }); - let zip = zip_validity(iter, &a); + let zip = zip_validity(iter, a); let a = zip.collect::>(); assert_eq!(a, vec![Some([1, 2].as_ref()), None]); @@ -43,11 +44,10 @@ fn slices() { #[test] fn byte() { - let a = Some(Bitmap::from([ - true, false, true, false, false, true, true, false, true, - ])); + let a = Bitmap::from([true, false, true, false, false, true, true, false, true]); + let a = Some(a.iter()); let values = vec![0, 1, 2, 3, 4, 5, 6, 7, 8]; - let zip = zip_validity(values.into_iter(), &a); + let zip = zip_validity(values.into_iter(), a); let a = zip.collect::>(); assert_eq!( @@ -68,10 +68,10 @@ fn byte() { #[test] fn offset() { - let a = Bitmap::from([true, false, true, false, false, true, true, false, true]); - let a = Some(a.slice(1, 8)); + let a = Bitmap::from([true, false, true, false, false, true, true, false, true]).slice(1, 8); + let a = Some(a.iter()); let values = vec![0, 1, 2, 3, 4, 5, 6, 7]; - let zip = zip_validity(values.into_iter(), &a); + let zip = zip_validity(values.into_iter(), a); let a = zip.collect::>(); assert_eq!( @@ -83,7 +83,7 @@ fn offset() { #[test] fn none() { let values = vec![0, 1, 2]; - let zip = zip_validity(values.into_iter(), &None); + let zip = zip_validity(values.into_iter(), None); let a = zip.collect::>(); assert_eq!(a, vec![Some(0), Some(1), Some(2)]); @@ -91,10 +91,10 @@ fn none() { #[test] fn rev() { - let a = Bitmap::from([true, false, true, false, false, true, true, false, true]); - let a = Some(a.slice(1, 8)); + let a = Bitmap::from([true, false, true, false, false, true, true, false, true]).slice(1, 8); + let a = Some(a.iter()); let values = vec![0, 1, 2, 3, 4, 5, 6, 7]; - let zip = zip_validity(values.into_iter(), &a); + let zip = zip_validity(values.into_iter(), a); let result = zip.rev().collect::>(); let expected = vec![None, Some(1), None, None, Some(4), Some(5), None, Some(7)]