Improved ZipValidity iterators

jorgecarleitao · Oct 29, 2022 · 03e981e · 03e981e
1 parent e106cff
commit 03e981e
Show file tree

Hide file tree

Showing 17 changed files with 85 additions and 19 deletions.
diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs
@@ -117,7 +117,12 @@ impl<O: Offset> BinaryArray<O> {
 
     /// Returns an iterator of `Option<&[u8]>` over every element of this array.
     pub fn iter(&self) -> ZipValidity<&[u8], BinaryValueIter<O>, BitmapIter> {
-        ZipValidity::new(self.values_iter(), self.validity.as_ref().map(|x| x.iter()))
+        let null_count = self.validity.as_ref().map(|validity| validity.unset_bits());
+        ZipValidity::new(
+            self.values_iter(),
+            self.validity.as_ref().map(|x| x.iter()),
+            null_count,
+        )
     }
 
     /// Returns an iterator of `&[u8]` over every element of this array, ignoring the validity

diff --git a/src/array/boolean/iterator.rs b/src/array/boolean/iterator.rs
@@ -22,8 +22,9 @@ impl IntoIterator for BooleanArray {
     fn into_iter(self) -> Self::IntoIter {
         let (_, values, validity) = self.into_inner();
         let values = values.into_iter();
+        let null_count = validity.as_ref().map(|validity| validity.unset_bits());
         let validity = validity.map(|x| x.into_iter());
-        ZipValidity::new(values, validity)
+        ZipValidity::new(values, validity, null_count)
     }
 }
 
@@ -41,9 +42,14 @@ impl<'a> MutableBooleanArray {
     /// Returns an iterator over the optional values of this [`MutableBooleanArray`].
     #[inline]
     pub fn iter(&'a self) -> ZipValidity<bool, BitmapIter<'a>, BitmapIter<'a>> {
+        let null_count = self
+            .validity()
+            .as_ref()
+            .map(|validity| validity.unset_bits());
         ZipValidity::new(
             self.values().iter(),
             self.validity().as_ref().map(|x| x.iter()),
+            null_count,
         )
     }
 

diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
@@ -91,6 +91,9 @@ impl BooleanArray {
         ZipValidity::new(
             self.values().iter(),
             self.validity.as_ref().map(|x| x.iter()),
+            self.validity()
+                .as_ref()
+                .map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs
@@ -194,6 +194,10 @@ impl<K: DictionaryKey> DictionaryArray<K> {
         ZipValidity::new(
             DictionaryValuesIter::new(self),
             self.keys.validity().as_ref().map(|x| x.iter()),
+            self.keys
+                .validity()
+                .as_ref()
+                .map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/fixed_size_binary/iterator.rs b/src/array/fixed_size_binary/iterator.rs
@@ -19,7 +19,11 @@ impl<'a> FixedSizeBinaryArray {
     pub fn iter(
         &'a self,
     ) -> ZipValidity<&'a [u8], std::slice::ChunksExact<'a, u8>, BitmapIter<'a>> {
-        ZipValidity::new(self.values_iter(), self.validity.as_ref().map(|x| x.iter()))
+        ZipValidity::new(
+            self.values_iter(),
+            self.validity.as_ref().map(|x| x.iter()),
+            self.validity.as_ref().map(|validity| validity.unset_bits()),
+        )
     }
 
     /// Returns iterator over the values of [`FixedSizeBinaryArray`]
@@ -45,6 +49,9 @@ impl<'a> MutableFixedSizeBinaryArray {
         ZipValidity::new(
             self.iter_values(),
             self.validity().as_ref().map(|x| x.iter()),
+            self.validity()
+                .as_ref()
+                .map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/fixed_size_list/iterator.rs b/src/array/fixed_size_list/iterator.rs
@@ -39,6 +39,9 @@ impl<'a> FixedSizeListArray {
         ZipValidity::new(
             FixedSizeListValuesIter::new(self),
             self.validity.as_ref().map(|x| x.iter()),
+            self.validity()
+                .as_ref()
+                .map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/list/iterator.rs b/src/array/list/iterator.rs
@@ -38,6 +38,9 @@ impl<'a, O: Offset> ListArray<O> {
         ZipValidity::new(
             ListValuesIter::new(self),
             self.validity.as_ref().map(|x| x.iter()),
+            self.validity()
+                .as_ref()
+                .map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/map/iterator.rs b/src/array/map/iterator.rs
@@ -75,6 +75,7 @@ impl<'a> MapArray {
         ZipValidity::new(
             MapValuesIter::new(self),
             self.validity.as_ref().map(|x| x.iter()),
+            self.validity.as_ref().map(|x| x.unset_bits()),
         )
     }
 

diff --git a/src/array/primitive/iterator.rs b/src/array/primitive/iterator.rs
@@ -16,8 +16,9 @@ impl<T: NativeType> IntoIterator for PrimitiveArray<T> {
     fn into_iter(self) -> Self::IntoIter {
         let (_, values, validity) = self.into_inner();
         let values = values.into_iter();
+        let null_count = validity.as_ref().map(|x| x.unset_bits());
         let validity = validity.map(|x| x.into_iter());
-        ZipValidity::new(values, validity)
+        ZipValidity::new(values, validity, null_count)
     }
 }
 
@@ -38,6 +39,9 @@ impl<'a, T: NativeType> MutablePrimitiveArray<T> {
         ZipValidity::new(
             self.values().iter(),
             self.validity().as_ref().map(|x| x.iter()),
+            self.validity()
+                .as_ref()
+                .map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
@@ -144,6 +144,7 @@ impl<T: NativeType> PrimitiveArray<T> {
         ZipValidity::new(
             self.values().iter(),
             self.validity().as_ref().map(|x| x.iter()),
+            self.validity.as_ref().map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/struct_/iterator.rs b/src/array/struct_/iterator.rs
@@ -92,6 +92,9 @@ impl<'a> StructArray {
         ZipValidity::new(
             StructValueIter::new(self),
             self.validity.as_ref().map(|x| x.iter()),
+            self.validity()
+                .as_ref()
+                .map(|validity| validity.unset_bits()),
         )
     }
 

diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs
@@ -133,7 +133,11 @@ impl<O: Offset> Utf8Array<O> {
 
     /// Returns an iterator of `Option<&str>`
     pub fn iter(&self) -> ZipValidity<&str, Utf8ValuesIter<O>, BitmapIter> {
-        ZipValidity::new(self.values_iter(), self.validity.as_ref().map(|x| x.iter()))
+        ZipValidity::new(
+            self.values_iter(),
+            self.validity.as_ref().map(|x| x.iter()),
+            self.validity.as_ref().map(|validity| validity.unset_bits()),
+        )
     }
 
     /// Returns an iterator of `&str`

diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
@@ -206,7 +206,11 @@ impl<O: Offset> MutableUtf8Array<O> {
 
     /// Returns an iterator of `Option<&str>`
     pub fn iter(&self) -> ZipValidity<&str, MutableUtf8ValuesIter<O>, BitmapIter> {
-        ZipValidity::new(self.values_iter(), self.validity.as_ref().map(|x| x.iter()))
+        ZipValidity::new(
+            self.values_iter(),
+            self.validity.as_ref().map(|x| x.iter()),
+            self.validity.as_ref().map(|validity| validity.unset_bits()),
+        )
     }
 
     /// Converts itself into an [`Array`].

diff --git a/src/bitmap/utils/zip_validity.rs b/src/bitmap/utils/zip_validity.rs
@@ -101,10 +101,13 @@ where
     V: Iterator<Item = bool>,
 {
     /// Returns a new [`ZipValidity`]
-    pub fn new(values: I, validity: Option<V>) -> Self {
+    pub fn new(values: I, validity: Option<V>, null_count: Option<usize>) -> Self {
         match validity {
-            Some(validity) => Self::Optional(ZipValidityIter::new(values, validity)),
-            None => Self::Required(values),
+            // only if we have a validity and there are nulls we will iterate them
+            Some(validity) if null_count != Some(0) => {
+                Self::Optional(ZipValidityIter::new(values, validity))
+            }
+            _ => Self::Required(values),
         }
     }
 }

diff --git a/src/io/avro/write/serialize.rs b/src/io/avro/write/serialize.rs
@@ -126,7 +126,11 @@ fn list_optional<'a, O: Offset>(array: &'a ListArray<O>, schema: &AvroSchema) ->
         .offsets()
         .windows(2)
         .map(|w| (w[1] - w[0]).to_usize() as i64);
-    let lengths = ZipValidity::new(lengths, array.validity().as_ref().map(|x| x.iter()));
+    let lengths = ZipValidity::new(
+        lengths,
+        array.validity().as_ref().map(|x| x.iter()),
+        array.validity().as_ref().map(|x| x.unset_bits()),
+    );
 
     Box::new(BufStreamingIterator::new(
         lengths,
@@ -180,7 +184,11 @@ fn struct_optional<'a>(array: &'a StructArray, schema: &Record) -> BoxSerializer
         .map(|(x, schema)| new_serializer(x.as_ref(), schema))
         .collect::<Vec<_>>();
 
-    let iterator = ZipValidity::new(0..array.len(), array.validity().as_ref().map(|x| x.iter()));
+    let iterator = ZipValidity::new(
+        0..array.len(),
+        array.validity().as_ref().map(|x| x.iter()),
+        array.validity().as_ref().map(|x| x.unset_bits()),
+    );
 
     Box::new(BufStreamingIterator::new(
         iterator,

diff --git a/src/io/json/write/serialize.rs b/src/io/json/write/serialize.rs
@@ -103,7 +103,11 @@ fn struct_serializer<'a>(
     let names = array.fields().iter().map(|f| f.name.as_str());
 
     Box::new(BufStreamingIterator::new(
-        ZipValidity::new(0..array.len(), array.validity().map(|x| x.iter())),
+        ZipValidity::new(
+            0..array.len(),
+            array.validity().map(|x| x.iter()),
+            array.validity().map(|x| x.unset_bits()),
+        ),
         move |maybe, buf| {
             if maybe.is_some() {
                 let names = names.clone();
@@ -143,6 +147,7 @@ fn list_serializer<'a, O: Offset>(
         ZipValidity::new(
             array.offsets().windows(2),
             array.validity().map(|x| x.iter()),
+            array.validity().map(|x| x.unset_bits()),
         ),
         move |offset, buf| {
             if let Some(offset) = offset {

diff --git a/tests/it/bitmap/utils/zip_validity.rs b/tests/it/bitmap/utils/zip_validity.rs
@@ -8,7 +8,7 @@ fn basic() {
     let a = Bitmap::from([true, false]);
     let a = Some(a.iter());
     let values = vec![0, 1];
-    let zip = ZipValidity::new(values.into_iter(), a);
+    let zip = ZipValidity::new(values.into_iter(), a, None);
 
     let a = zip.collect::<Vec<_>>();
     assert_eq!(a, vec![Some(0), None]);
@@ -19,7 +19,7 @@ fn complete() {
     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 = ZipValidity::new(values.into_iter(), a);
+    let zip = ZipValidity::new(values.into_iter(), a, None);
 
     let a = zip.collect::<Vec<_>>();
     assert_eq!(
@@ -31,6 +31,7 @@ fn complete() {
 #[test]
 fn slices() {
     let a = Bitmap::from([true, false]);
+    let null_count = Some(a.unset_bits());
     let a = Some(a.iter());
     let offsets = vec![0, 2, 3];
     let values = vec![1, 2, 3];
@@ -39,7 +40,7 @@ fn slices() {
         let end = x[1];
         &values[start..end]
     });
-    let zip = ZipValidity::new(iter, a);
+    let zip = ZipValidity::new(iter, a, null_count);
 
     let a = zip.collect::<Vec<_>>();
     assert_eq!(a, vec![Some([1, 2].as_ref()), None]);
@@ -50,7 +51,7 @@ fn byte() {
     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 = ZipValidity::new(values.into_iter(), a);
+    let zip = ZipValidity::new(values.into_iter(), a, None);
 
     let a = zip.collect::<Vec<_>>();
     assert_eq!(
@@ -74,7 +75,7 @@ fn offset() {
     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 = ZipValidity::new(values.into_iter(), a);
+    let zip = ZipValidity::new(values.into_iter(), a, None);
 
     let a = zip.collect::<Vec<_>>();
     assert_eq!(
@@ -86,7 +87,7 @@ fn offset() {
 #[test]
 fn none() {
     let values = vec![0, 1, 2];
-    let zip = ZipValidity::new(values.into_iter(), None::<BitmapIter>);
+    let zip = ZipValidity::new(values.into_iter(), None::<BitmapIter>, None);
 
     let a = zip.collect::<Vec<_>>();
     assert_eq!(a, vec![Some(0), Some(1), Some(2)]);
@@ -95,9 +96,10 @@ fn none() {
 #[test]
 fn rev() {
     let a = Bitmap::from([true, false, true, false, false, true, true, false, true]).slice(1, 8);
+    let null_count = Some(a.unset_bits());
     let a = Some(a.iter());
     let values = vec![0, 1, 2, 3, 4, 5, 6, 7];
-    let zip = ZipValidity::new(values.into_iter(), a);
+    let zip = ZipValidity::new(values.into_iter(), a, null_count);
 
     let result = zip.rev().collect::<Vec<_>>();
     let expected = vec![None, Some(1), None, None, Some(4), Some(5), None, Some(7)]