don't count nulls when iterating validities

pola-rs · Nov 4, 2021 · 7eb17c0 · 7eb17c0
1 parent e0dfa2e
commit 7eb17c0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/polars/polars-core/src/chunked_array/mod.rs b/polars/polars-core/src/chunked_array/mod.rs
@@ -180,14 +180,14 @@ impl<T> ChunkedArray<T> {
     /// Get the index of the first non null value in this ChunkedArray.
     pub fn first_non_null(&self) -> Option<usize> {
         let mut offset = 0;
-        for (_, null_bitmap) in self.null_bits() {
-            if let Some(null_bitmap) = null_bitmap {
-                for (idx, is_valid) in null_bitmap.iter().enumerate() {
+        for validity in self.iter_validities() {
+            if let Some(validity) = validity {
+                for (idx, is_valid) in validity.iter().enumerate() {
                     if is_valid {
                         return Some(offset + idx);
                     }
                 }
-                offset += null_bitmap.len()
+                offset += validity.len()
             } else {
                 return Some(offset);
             }
@@ -196,10 +196,8 @@ impl<T> ChunkedArray<T> {
     }
 
     /// Get the buffer of bits representing null values
-    pub fn null_bits(&self) -> impl Iterator<Item = (usize, Option<&Bitmap>)> + '_ {
-        self.chunks
-            .iter()
-            .map(|arr| (arr.null_count(), arr.validity()))
+    pub fn iter_validities(&self) -> impl Iterator<Item = Option<&Bitmap>> + '_ {
+        self.chunks.iter().map(|arr| arr.validity())
     }
 
     /// Shrink the capacity of this array to fit it's length.

diff --git a/polars/polars-core/src/chunked_array/ops/apply.rs b/polars/polars-core/src/chunked_array/ops/apply.rs
@@ -60,8 +60,8 @@ where
     {
         let chunks = self
             .data_views()
-            .zip(self.null_bits())
-            .map(|(slice, (_, validity))| {
+            .zip(self.iter_validities())
+            .map(|(slice, validity)| {
                 let values = AlignedVec::<_>::from_trusted_len_iter(slice.iter().map(|&v| f(v)));
                 to_array::<S>(values, validity.cloned())
             })
@@ -97,8 +97,8 @@ where
         let chunks = self
             .data_views()
             .into_iter()
-            .zip(self.null_bits())
-            .map(|(slice, (_, validity))| {
+            .zip(self.iter_validities())
+            .map(|(slice, validity)| {
                 let values = slice.iter().copied().map(f);
                 let values = AlignedVec::<_>::from_trusted_len_iter(values);
                 to_array::<T>(values, validity.cloned())
@@ -114,8 +114,8 @@ where
         let mut ca: ChunkedArray<T> = self
             .data_views()
             .into_iter()
-            .zip(self.null_bits())
-            .map(|(slice, (_null_count, validity))| {
+            .zip(self.iter_validities())
+            .map(|(slice, validity)| {
                 let vec: Result<AlignedVec<_>> = slice.iter().copied().map(f).collect();
                 Ok((vec?, validity.cloned()))
             })