Generalized ZipIterator. (#296)

src/array/binary/iterator.rs

@@ -52,7 +52,10 @@ impl<'a, O: Offset> IntoIterator for &'a BinaryArray<O> {
 impl<'a, O: Offset> BinaryArray<O> {
     /// 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]`

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`

src/array/dictionary/iterator.rs

@@ -70,7 +70,10 @@ impl<'a, K: DictionaryKey> IntoIterator for &'a DictionaryArray<K> {
 impl<'a, K: DictionaryKey> DictionaryArray<K> {
     /// Returns an iterator of `Option<Box<dyn Array>>`
     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<dyn Array>`

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()),
+        )
     }
 }

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<Box<dyn Array>>`
     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<dyn Array>`

src/array/list/iterator.rs

@@ -76,7 +76,10 @@ impl<'a, O: Offset> IntoIterator for &'a ListArray<O> {
 impl<'a, O: Offset> ListArray<O> {
     /// Returns an iterator of `Option<Box<dyn Array>>`
     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<dyn Array>`

src/array/primitive/iterator.rs

@@ -19,6 +19,9 @@ impl<'a, T: NativeType> PrimitiveArray<T> {
     /// 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()),
+        )
     }
 }

src/array/utf8/iterator.rs

@@ -68,7 +68,10 @@ impl<'a, O: Offset> IntoIterator for &'a Utf8Array<O> {
 impl<'a, O: Offset> Utf8Array<O> {
     /// 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`

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<Item = T> + Clone> Clone for ZipValidity<'a, T, I> {
 }
 
 impl<'a, T, I: Iterator<Item = T>> ZipValidity<'a, T, I> {
-    #[inline]
-    pub fn new(values: I, validity: &'a Option<Bitmap>) -> 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<BitmapIter<'a>>) -> 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<Item = T>> std::iter::ExactSizeIterator for ZipValidity<
 unsafe impl<T, I: TrustedLen<Item = T>> TrustedLen for ZipValidity<'_, T, I> {}
 
 /// Returns an iterator adapter that returns Option<T> according to an optional validity.
-#[inline]
 pub fn zip_validity<T, I: Iterator<Item = T>>(
     values: I,
-    validity: &Option<Bitmap>,
+    validity: Option<BitmapIter>,
 ) -> ZipValidity<T, I> {
-    ZipValidity::<T, I>::new(values, validity)
+    ZipValidity::new(values, validity)
 }

tests/it/bitmap/utils/zip_validity.rs

@@ -2,21 +2,21 @@ 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::<Vec<_>>();
     assert_eq!(a, vec![Some(0), None]);
 }
 
 #[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::<Vec<_>>();
     assert_eq!(
@@ -27,27 +27,27 @@ 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| {
         let start = x[0];
         let end = x[1];
         &values[start..end]
     });
-    let zip = zip_validity(iter, &a);
+    let zip = zip_validity(iter, a);
 
     let a = zip.collect::<Vec<_>>();
     assert_eq!(a, vec![Some([1, 2].as_ref()), None]);
 }
 
 #[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::<Vec<_>>();
     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::<Vec<_>>();
     assert_eq!(
@@ -83,18 +83,18 @@ 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::<Vec<_>>();
     assert_eq!(a, vec![Some(0), Some(1), Some(2)]);
 }
 
 #[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::<Vec<_>>();
     let expected = vec![None, Some(1), None, None, Some(4), Some(5), None, Some(7)]