diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs index 17c29f1b30d..5a84241359b 100644 --- a/src/array/binary/mod.rs +++ b/src/array/binary/mod.rs @@ -117,7 +117,12 @@ impl BinaryArray { /// Returns an iterator of `Option<&[u8]>` over every element of this array. pub fn iter(&self) -> ZipValidity<&[u8], BinaryValueIter, 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 index 019ba5b388f..9beb53b1dbb 100644 --- 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, 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 index 2c33ef155ab..816e901144e 100644 --- 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 index 02d72ccd260..81d2eec354d 100644 --- a/src/array/dictionary/mod.rs +++ b/src/array/dictionary/mod.rs @@ -194,6 +194,10 @@ impl DictionaryArray { 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 index ad3c36968fa..2e4acaf60f2 100644 --- 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 index 4ee178da5a7..046050d5671 100644 --- 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 index 90ec9b4cb58..8fff5fa8204 100644 --- a/src/array/list/iterator.rs +++ b/src/array/list/iterator.rs @@ -38,6 +38,9 @@ impl<'a, O: Offset> ListArray { 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 index 8bd8ce820ed..c20324a5d1e 100644 --- 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 index 89a8842e8fd..7b4a40a0355 100644 --- a/src/array/primitive/iterator.rs +++ b/src/array/primitive/iterator.rs @@ -16,8 +16,9 @@ impl IntoIterator for PrimitiveArray { 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 { 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 index 2bbdbd05239..c27e98395fb 100644 --- a/src/array/primitive/mod.rs +++ b/src/array/primitive/mod.rs @@ -144,6 +144,7 @@ impl PrimitiveArray { 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 index c5c75a6b925..884cc300126 100644 --- 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 index dc453a5c37b..84ab1b941f0 100644 --- a/src/array/utf8/mod.rs +++ b/src/array/utf8/mod.rs @@ -133,7 +133,11 @@ impl Utf8Array { /// Returns an iterator of `Option<&str>` pub fn iter(&self) -> ZipValidity<&str, Utf8ValuesIter, 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 index 4cc83da9011..c993ad18f8b 100644 --- a/src/array/utf8/mutable.rs +++ b/src/array/utf8/mutable.rs @@ -206,7 +206,11 @@ impl MutableUtf8Array { /// Returns an iterator of `Option<&str>` pub fn iter(&self) -> ZipValidity<&str, MutableUtf8ValuesIter, 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 index 88ad018db33..201cb071ae2 100644 --- a/src/bitmap/utils/zip_validity.rs +++ b/src/bitmap/utils/zip_validity.rs @@ -101,10 +101,13 @@ where V: Iterator, { /// Returns a new [`ZipValidity`] - pub fn new(values: I, validity: Option) -> Self { + pub fn new(values: I, validity: Option, null_count: Option) -> 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/tests/it/bitmap/utils/zip_validity.rs b/tests/it/bitmap/utils/zip_validity.rs index 9c6187f60a6..c86240eb0c2 100644 --- 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::>(); 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::>(); 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::>(); 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::>(); 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::>(); 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::); + let zip = ZipValidity::new(values.into_iter(), None::, None); let a = zip.collect::>(); 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::>(); let expected = vec![None, Some(1), None, None, Some(4), Some(5), None, Some(7)]