Repository: jorgecarleitao/arrow2
Pull request #994 (2022-05-18 06:12:04)

delay null_counts
On several `extend` (which may be called in hot loops) there were still `null_count` check on mutable arrays.

Upon freezing a `MutableArray` to `Array` there was also this pattern

```
// first null_count
if mutable_bitmap::null_count > 0 {
  // second null_count on conversion to bitmap
  mutable_bitmap::into_bitmap
}
```

I rewrote it by first doing the `bitmap` coversions (which computes the nulls). And then if the `null_count == 0` then we set `validity = None`.

So basically this PR delays all `null_counts` until we freeze the array as we will compute it upon conversion and can then decide if we can drop the validity or not.


codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#994](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (85c1cbe) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/aafba7b4eb4991e016638cbc1d4df676912e8236?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (aafba7b) will **increase** coverage by `0.00%`.
> The diff coverage is `89.47%`.

```diff
@@           Coverage Diff           @@
##             main     #994   +/-   ##
=======================================
  Coverage   71.41%   71.41%           
=======================================
  Files         357      357           
  Lines       19801    19803    +2     
=======================================
+ Hits        14140    14142    +2     
  Misses       5661     5661           
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/array/utf8/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3V0ZjgvbXV0YWJsZS5ycw==) | `86.84% <87.50%> (-0.19%)` | :arrow_down: |
| [src/array/primitive/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3ByaW1pdGl2ZS9tdXRhYmxlLnJz) | `89.07% <90.90%> (+0.27%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [aafba7b...85c1cbe](https://codecov.io/gh/jorgecarleitao/arrow2/pull/994?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).

ritchie46: @jorgecarleitao I added an extra test and that required met to clone the `MutableArray`. I derived `Clone` for `MutablePrimitiveArray` and `MutableBitmap` as that seemed a reasonable trait to have to me.

Was there any reason it was not deriving clone other than just not yet added?

Patch:
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index 43f734f74b..24768f8361 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -1,5 +1,6 @@
 use std::{iter::FromIterator, sync::Arc};
 
+use crate::bitmap::Bitmap;
 use crate::{
     array::{Array, MutableArray, TryExtend, TryPush},
     bitmap::MutableBitmap,
@@ -13,7 +14,7 @@ use super::PrimitiveArray;
 
 /// The Arrow's equivalent to `Vec<Option<T>>` where `T` is byte-size (e.g. `i32`).
 /// Converting a [`MutablePrimitiveArray`] into a [`PrimitiveArray`] is `O(1)`.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct MutablePrimitiveArray<T: NativeType> {
     data_type: DataType,
     values: Vec<T>,
@@ -22,11 +23,14 @@ pub struct MutablePrimitiveArray<T: NativeType> {
 
 impl<T: NativeType> From<MutablePrimitiveArray<T>> for PrimitiveArray<T> {
     fn from(other: MutablePrimitiveArray<T>) -> Self {
-        let validity = if other.validity.as_ref().map(|x| x.null_count()).unwrap_or(0) > 0 {
-            other.validity.map(|x| x.into())
-        } else {
-            None
-        };
+        let validity = other.validity.and_then(|x| {
+            let bitmap: Bitmap = x.into();
+            if bitmap.null_count() == 0 {
+                None
+            } else {
+                Some(bitmap)
+            }
+        });
 
         PrimitiveArray::<T>::new(other.data_type, other.values.into(), validity)
     }
@@ -194,9 +198,7 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
             let mut validity = MutableBitmap::new();
             validity.extend_constant(self.len(), true);
             extend_trusted_len_unzip(iterator, &mut validity, &mut self.values);
-            if validity.null_count() > 0 {
-                self.validity = Some(validity);
-            }
+            self.validity = Some(validity);
         }
     }
     /// Extends the [`MutablePrimitiveArray`] from an iterator of values of trusted len.
@@ -526,11 +528,7 @@ impl<T: NativeType, Ptr: std::borrow::Borrow<Option<T>>> FromIterator<Ptr>
             })
             .collect();
 
-        let validity = if validity.null_count() > 0 {
-            Some(validity)
-        } else {
-            None
-        };
+        let validity = Some(validity);
 
         Self {
             data_type: T::PRIMITIVE.into(),
@@ -588,11 +586,7 @@ where
 
     extend_trusted_len_unzip(iterator, &mut validity, &mut buffer);
 
-    let validity = if validity.null_count() > 0 {
-        Some(validity)
-    } else {
-        None
-    };
+    let validity = Some(validity);
 
     (validity, buffer)
 }
@@ -634,11 +628,7 @@ where
     buffer.set_len(len);
     null.set_len(len);
 
-    let validity = if null.null_count() > 0 {
-        Some(null)
-    } else {
-        None
-    };
+    let validity = Some(null);
 
     Ok((validity, buffer))
 }
diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
index f9f53a6846..f19e331cfa 100644
--- a/src/array/utf8/mutable.rs
+++ b/src/array/utf8/mutable.rs
@@ -13,6 +13,7 @@ use crate::{
 
 use super::Utf8Array;
 use crate::array::physical_binary::*;
+use crate::bitmap::Bitmap;
 
 struct StrAsBytes<P>(P);
 impl<T: AsRef<str>> AsRef<[u8]> for StrAsBytes<T> {
@@ -36,12 +37,21 @@ impl<O: Offset> From<MutableUtf8Array<O>> for Utf8Array<O> {
         // Safety:
         // `MutableUtf8Array` has the same invariants as `Utf8Array` and thus
         // `Utf8Array` can be safely created from `MutableUtf8Array` without checks.
+        let validity = other.validity.and_then(|x| {
+            let bitmap: Bitmap = x.into();
+            if bitmap.null_count() == 0 {
+                None
+            } else {
+                Some(bitmap)
+            }
+        });
+
         unsafe {
             Utf8Array::<O>::from_data_unchecked(
                 other.data_type,
                 other.offsets.into(),
                 other.values.into(),
-                other.validity.map(|x| x.into()),
+                validity,
             )
         }
     }
@@ -368,10 +378,6 @@ impl<O: Offset> MutableUtf8Array<O> {
             self.validity.as_mut().unwrap(),
             iterator,
         );
-
-        if self.validity.as_mut().unwrap().null_count() == 0 {
-            self.validity = None;
-        }
     }
 
     /// Creates a [`MutableUtf8Array`] from an iterator of trusted length.
diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs
index 94e0b00d7c..c8e5e04649 100644
--- a/src/bitmap/mutable.rs
+++ b/src/bitmap/mutable.rs
@@ -14,6 +14,7 @@ use super::Bitmap;
 /// The main difference against [`Vec<bool>`] is that a bitmap cannot be represented as `&[bool]`.
 /// # Implementation
 /// This container is backed by [`Vec<u8>`].
+#[derive(Clone)]
 pub struct MutableBitmap {
     buffer: Vec<u8>,
     // invariant: length.saturating_add(7) / 8 == buffer.len();


Test patch:
diff --git a/tests/it/array/primitive/mutable.rs b/tests/it/array/primitive/mutable.rs
index 6ff96a9e86..99f897076d 100644
--- a/tests/it/array/primitive/mutable.rs
+++ b/tests/it/array/primitive/mutable.rs
@@ -126,7 +126,8 @@ fn set() {
 fn from_iter() {
     let a = MutablePrimitiveArray::<i32>::from_iter((0..2).map(Some));
     assert_eq!(a.len(), 2);
-    assert_eq!(a.validity(), None);
+    let validity = a.validity().unwrap();
+    assert_eq!(validity.null_count(), 0);
 }
 
 #[test]
@@ -176,7 +177,8 @@ fn from_trusted_len() {
 fn extend_trusted_len() {
     let mut a = MutablePrimitiveArray::<i32>::new();
     a.extend_trusted_len(vec![Some(1), Some(2)].into_iter());
-    assert_eq!(a.validity(), None);
+    let validity = a.validity().unwrap();
+    assert_eq!(validity.null_count(), 0);
     a.extend_trusted_len(vec![None, Some(4)].into_iter());
     assert_eq!(
         a.validity(),
@@ -266,7 +268,14 @@ fn extend_from_slice() {
 fn set_validity() {
     let mut a = MutablePrimitiveArray::<i32>::new();
     a.extend_trusted_len(vec![Some(1), Some(2)].into_iter());
-    assert_eq!(a.validity(), None);
+    let validity = a.validity().unwrap();
+    assert_eq!(validity.null_count(), 0);
+
+    // test that upon conversion to array the bitmap is set to None
+    let arr: PrimitiveArray<_> = a.clone().into();
+    assert_eq!(arr.validity(), None);
+
+    // test set_validity
     a.set_validity(Some(MutableBitmap::from([false, true])));
     assert_eq!(a.validity(), Some(&MutableBitmap::from([false, true])));
 }