Repository: jorgecarleitao/arrow2
Pull request #1076 (2022-06-15 18:31:14)

Improve `arity_assign`: `~2x` improvement if we share data.
This continues upon #1073 (you may close that one if you prefer in one PR, otherwise I rebase).

I did a local benchmark to see what the cost of the new `arity_assign` would be in the case we share data. In polars expressions that is every first operation:

```
Gnuplot not found, using plotters backend
bench_1                 time:   [608.34 us 609.30 us 610.73 us]                    
                        change: [+44.741% +45.228% +45.698%] (p = 0.00 < 0.05)
                        Performance has regressed.
Found 7 outliers among 100 measurements (7.00%)
  2 (2.00%) high mild
  5 (5.00%) high severe
```

This showed, that in case of shared data, we first do a `memcpy` and this is `~2x` slower than the kernels where we allocate a new memory buffer.

Luckily we can have best of both worlds. :)

This PR branches depending on the shared state:

* in case we share state -> we alloc (and thus choose the fastest path possible)
* in case we own -> we mutate in place (and thus choose the fastest path possible)

I also added a distinction between `set_validity` and `with_validity` as we now have mutable data, it is not always logical return a new `PrimitiveArray`.

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1076](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (37e366a) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/39db6fb7514364bfea08d594793b23e1ed5a7def?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (39db6fb) will **decrease** coverage by `0.04%`.
> The diff coverage is `46.80%`.

```diff
@@            Coverage Diff             @@
##             main    #1076      +/-   ##
==========================================
- Coverage   81.01%   80.96%   -0.05%     
==========================================
  Files         366      366              
  Lines       35056    35096      +40     
==========================================
+ Hits        28399    28414      +15     
- Misses       6657     6682      +25     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/array/primitive/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3ByaW1pdGl2ZS9tdXRhYmxlLnJz) | `95.37% <42.85%> (-0.92%)` | :arrow_down: |
| [src/compute/arity\_assign.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2NvbXB1dGUvYXJpdHlfYXNzaWduLnJz) | `71.42% <47.36%> (-21.17%)` | :arrow_down: |
| [src/array/primitive/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3ByaW1pdGl2ZS9tb2QucnM=) | `81.44% <47.61%> (-4.27%)` | :arrow_down: |
| [src/bitmap/utils/slice\_iterator.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9zbGljZV9pdGVyYXRvci5ycw==) | `98.78% <0.00%> (+1.21%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076?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/1076?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [39db6fb...37e366a](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1076?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).


Patch:
diff --git a/benches/assign_ops.rs b/benches/assign_ops.rs
index 35f29b643c..dc4c9b182f 100644
--- a/benches/assign_ops.rs
+++ b/benches/assign_ops.rs
@@ -14,7 +14,7 @@ fn add_benchmark(c: &mut Criterion) {
         c.bench_function(&format!("apply_mul 2^{}", log2_size), |b| {
             b.iter(|| {
                 criterion::black_box(&mut arr_a)
-                    .apply_values(|x| x.iter_mut().for_each(|x| *x *= 1.01));
+                    .apply_values_mut(|x| x.iter_mut().for_each(|x| *x *= 1.01));
                 assert!(!arr_a.value(10).is_nan());
             })
         });
@@ -30,7 +30,7 @@ fn add_benchmark(c: &mut Criterion) {
         let mut arr_a = create_primitive_array::<f32>(size, 0.2);
         let mut arr_b = create_primitive_array_with_seed::<f32>(size, 0.2, 10);
         // convert to be close to 1.01
-        arr_b.apply_values(|x| x.iter_mut().for_each(|x| *x = 1.01 + *x / 20.0));
+        arr_b.apply_values_mut(|x| x.iter_mut().for_each(|x| *x = 1.01 + *x / 20.0));
 
         c.bench_function(&format!("apply_mul null 2^{}", log2_size), |b| {
             b.iter(|| {
diff --git a/examples/cow.rs b/examples/cow.rs
index 4b76972a79..3255d365f4 100644
--- a/examples/cow.rs
+++ b/examples/cow.rs
@@ -13,7 +13,7 @@ fn main() {
         .unwrap();
 
     // 2. call `apply_values` with the function to apply over the values
-    array.apply_values(|x| x.iter_mut().for_each(|x| *x *= 10));
+    array.apply_values_mut(|x| x.iter_mut().for_each(|x| *x *= 10));
 
     // confirm that it gives the right result :)
     assert_eq!(array, &PrimitiveArray::from_vec(vec![10i32, 20]));
diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
index 6faef21186..815ec6b134 100644
--- a/src/array/boolean/mod.rs
+++ b/src/array/boolean/mod.rs
@@ -102,7 +102,7 @@ impl BooleanArray {
     /// if it is being shared (since it results in a `O(N)` memcopy).
     /// # Panics
     /// This function panics if the function modifies the length of the [`MutableBitmap`].
-    pub fn apply_values<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+    pub fn apply_values_mut<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
         let values = std::mem::take(&mut self.values);
         let mut values = values.make_mut();
         f(&mut values);
@@ -121,7 +121,7 @@ impl BooleanArray {
     /// if it is being shared (since it results in a `O(N)` memcopy).
     /// # Panics
     /// This function panics if the function modifies the length of the [`MutableBitmap`].
-    pub fn apply_validity<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+    pub fn apply_validity_mut<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
         if let Some(validity) = self.validity.as_mut() {
             let values = std::mem::take(validity);
             let mut bitmap = values.make_mut();
diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
index de55b6bb28..8443fca7d5 100644
--- a/src/array/primitive/mod.rs
+++ b/src/array/primitive/mod.rs
@@ -245,12 +245,41 @@ impl<T: NativeType> PrimitiveArray<T> {
     /// This function panics iff `validity.len() != self.len()`.
     #[must_use]
     pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
+        let mut out = self.clone();
+        out.set_validity(validity);
+        out
+    }
+
+    /// Update the validity buffer of this [`PrimitiveArray`].
+    /// # Panics
+    /// This function panics iff `values.len() != self.len()`.
+    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
         if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
             panic!("validity should be as least as large as the array")
         }
-        let mut arr = self.clone();
-        arr.validity = validity;
-        arr
+        self.validity = validity;
+    }
+
+    /// Returns a clone of this [`PrimitiveArray`] with a new values.
+    /// # Panics
+    /// This function panics iff `values.len() != self.len()`.
+    #[must_use]
+    pub fn with_values(&self, values: Buffer<T>) -> Self {
+        let mut out = self.clone();
+        out.set_values(values);
+        out
+    }
+
+    /// Update the values buffer of this [`PrimitiveArray`].
+    /// # Panics
+    /// This function panics iff `values.len() != self.len()`.
+    pub fn set_values(&mut self, values: Buffer<T>) {
+        assert_eq!(
+            values.len(),
+            self.len(),
+            "values length should be equal to this arrays length"
+        );
+        self.values = values;
     }
 
     /// Applies a function `f` to the values of this array, cloning the values
@@ -260,10 +289,14 @@ impl<T: NativeType> PrimitiveArray<T> {
     /// # Implementation
     /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
     /// if it is being shared (since it results in a `O(N)` memcopy).
-    pub fn apply_values<F: Fn(&mut [T])>(&mut self, f: F) {
+    /// # Panics
+    /// This function panics, if `f` modifies the length of `&mut [T]`
+    pub fn apply_values_mut<F: Fn(&mut [T])>(&mut self, f: F) {
         let values = std::mem::take(&mut self.values);
         let mut values = values.make_mut();
+        let len = values.len();
         f(&mut values);
+        assert_eq!(values.len(), len, "values length must remain the same");
         self.values = values.into();
     }
 
@@ -276,13 +309,29 @@ impl<T: NativeType> PrimitiveArray<T> {
     /// if it is being shared (since it results in a `O(N)` memcopy).
     /// # Panics
     /// This function panics if the function modifies the length of the [`MutableBitmap`].
-    pub fn apply_validity<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+    pub fn apply_validity_mut<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+        if let Some(validity) = self.validity.as_mut() {
+            let owned_validity = std::mem::take(validity);
+            let mut mut_bitmap = owned_validity.make_mut();
+            f(&mut mut_bitmap);
+            assert_eq!(mut_bitmap.len(), self.values.len());
+            *validity = mut_bitmap.into();
+        }
+    }
+
+    /// Applies a function `f` to the validity of this array, the caller can decide to make
+    /// it mutable or not.
+    ///
+    /// This is an API to leverage clone-on-write
+    /// # Implementation
+    /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
+    /// if it is being shared (since it results in a `O(N)` memcopy).
+    /// # Panics
+    /// This function panics if the function modifies the length of the [`MutableBitmap`].
+    pub fn apply_validity<F: Fn(&mut Bitmap)>(&mut self, f: F) {
         if let Some(validity) = self.validity.as_mut() {
-            let values = std::mem::take(validity);
-            let mut bitmap = values.make_mut();
-            f(&mut bitmap);
-            assert_eq!(bitmap.len(), self.values.len());
-            *validity = bitmap.into();
+            f(validity);
+            assert_eq!(validity.len(), self.values.len());
         }
     }
 
@@ -511,3 +560,9 @@ pub type UInt16Vec = MutablePrimitiveArray<u16>;
 pub type UInt32Vec = MutablePrimitiveArray<u32>;
 /// A type definition [`MutablePrimitiveArray`] for `u64`
 pub type UInt64Vec = MutablePrimitiveArray<u64>;
+
+impl<T: NativeType> Default for PrimitiveArray<T> {
+    fn default() -> Self {
+        PrimitiveArray::new(T::PRIMITIVE.into(), Default::default(), None)
+    }
+}
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index d59df1b845..ae983c10af 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -81,6 +81,37 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
     pub fn into_data(self) -> (DataType, Vec<T>, Option<MutableBitmap>) {
         (self.data_type, self.values, self.validity)
     }
+
+    /// Applies a function `f` to the values of this array, cloning the values
+    /// iff they are being shared with others
+    ///
+    /// This is an API to use clone-on-write
+    /// # Implementation
+    /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
+    /// if it is being shared (since it results in a `O(N)` memcopy).
+    /// # Panics
+    /// This function panics, if `f` modifies the length of `&mut [T]`
+    pub fn apply_values<F: Fn(&mut [T])>(&mut self, f: F) {
+        let len = self.values.len();
+        f(&mut self.values);
+        assert_eq!(len, self.values.len(), "values length must remain the same")
+    }
+
+    /// Applies a function `f` to the validity of this array, cloning it
+    /// iff it is being shared.
+    ///
+    /// This is an API to leverage clone-on-write
+    /// # Implementation
+    /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
+    /// if it is being shared (since it results in a `O(N)` memcopy).
+    /// # Panics
+    /// This function panics if the function modifies the length of the [`MutableBitmap`].
+    pub fn apply_validity<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+        if let Some(validity) = &mut self.validity {
+            f(validity);
+            assert_eq!(validity.len(), self.values.len());
+        }
+    }
 }
 
 impl<T: NativeType> Default for MutablePrimitiveArray<T> {
diff --git a/src/bitmap/immutable.rs b/src/bitmap/immutable.rs
index 45d8e89e4e..b63310b07c 100644
--- a/src/bitmap/immutable.rs
+++ b/src/bitmap/immutable.rs
@@ -239,7 +239,12 @@ impl Bitmap {
     pub fn make_mut(self) -> MutableBitmap {
         match self.into_mut() {
             Either::Left(data) => {
-                MutableBitmap::from_vec(data.bytes.as_ref().to_vec(), data.length)
+                if data.offset > 0 {
+                    // we have to recreate the bytes because a MutableBitmap does not have an `offset`.
+                    data.iter().collect()
+                } else {
+                    MutableBitmap::from_vec(data.bytes.as_ref().to_vec(), data.length)
+                }
             }
             Either::Right(data) => data,
         }
diff --git a/src/compute/arity_assign.rs b/src/compute/arity_assign.rs
index d10b100117..b8746cdd81 100644
--- a/src/compute/arity_assign.rs
+++ b/src/compute/arity_assign.rs
@@ -2,6 +2,7 @@
 
 use super::utils::check_same_len;
 use crate::{array::PrimitiveArray, types::NativeType};
+use either::Either;
 
 /// Applies an unary function to a [`PrimitiveArray`] in-place via cow semantics.
 ///
@@ -17,7 +18,7 @@ where
     I: NativeType,
     F: Fn(I) -> I,
 {
-    array.apply_values(|values| values.iter_mut().for_each(|v| *v = op(*v)));
+    array.apply_values_mut(|values| values.iter_mut().for_each(|v| *v = op(*v)));
 }
 
 /// Applies a binary operations to two [`PrimitiveArray`], applying the operation
@@ -38,20 +39,68 @@ where
 {
     check_same_len(lhs, rhs).unwrap();
 
+    // both for the validity and for the values
+    // we branch to check if we can mutate in place
+    // if we can, great that is fastest.
+    // if we cannot, we allocate a new buffer and assign values to that
+    // new buffer, that is benchmarked to be ~2x faster than first memcpy and assign in place
+    // for the validity bits it can be much faster as we might need to iterate all bits if the
+    // bitmap has an offset.
     match rhs.validity() {
         None => {}
         Some(rhs) => {
             if lhs.validity().is_none() {
                 *lhs = lhs.with_validity(Some(rhs.clone()))
             } else {
-                lhs.apply_validity(|mut lhs| lhs &= rhs)
+                lhs.apply_validity(|bitmap| {
+                    // we need to take ownership for the `into_mut` call, but leave the `&mut` lhs intact
+                    // so that we can later assign the result to out `&mut bitmap`
+                    let owned_lhs = std::mem::take(bitmap);
+
+                    match owned_lhs.into_mut() {
+                        // we take alloc and write to new buffer
+                        Either::Left(immutable) => {
+                            // we allocate a new bitmap because that is a lot faster
+                            // than doing the memcpy or the potential iteration of bits if
+                            // we are dealing with an offset
+                            let new = &immutable & rhs;
+                            *bitmap = new;
+                        }
+                        // we can mutate in place, happy days.
+                        Either::Right(mut mutable) => {
+                            let mut mutable_ref = &mut mutable;
+                            mutable_ref &= rhs;
+                            *bitmap = mutable.into()
+                        }
+                    }
+                });
             }
         }
     }
+    // we need to take ownership for the `into_mut` call, but leave the `&mut` lhs intact
+    // so that we can later assign the result to out `&mut lhs`
+    let owned_lhs = std::mem::take(lhs);
 
-    lhs.apply_values(|x| {
-        x.iter_mut()
-            .zip(rhs.values().iter())
-            .for_each(|(l, r)| *l = op(*l, *r))
-    });
+    match owned_lhs.into_mut() {
+        // we take alloc and write to new buffer
+        Either::Left(mut immutable) => {
+            let values = immutable
+                .values()
+                .iter()
+                .zip(rhs.values().iter())
+                .map(|(l, r)| op(*l, *r))
+                .collect::<Vec<_>>();
+            immutable.set_values(values.into());
+            *lhs = immutable;
+        }
+        // we can mutate in place
+        Either::Right(mut mutable) => {
+            mutable.apply_values(|x| {
+                x.iter_mut()
+                    .zip(rhs.values().iter())
+                    .for_each(|(l, r)| *l = op(*l, *r))
+            });
+            *lhs = mutable.into()
+        }
+    }
 }


Test patch:
diff --git a/tests/it/array/boolean/mod.rs b/tests/it/array/boolean/mod.rs
index 786fdb1dbc..ffd5867bbb 100644
--- a/tests/it/array/boolean/mod.rs
+++ b/tests/it/array/boolean/mod.rs
@@ -135,7 +135,7 @@ fn from_iter() {
 #[test]
 fn apply_values() {
     let mut a = BooleanArray::from([Some(true), Some(false), None]);
-    a.apply_values(|x| {
+    a.apply_values_mut(|x| {
         let mut a = std::mem::take(x);
         a = !a;
         *x = a;
@@ -147,7 +147,7 @@ fn apply_values() {
 #[test]
 fn apply_validity() {
     let mut a = BooleanArray::from([Some(true), Some(false), None]);
-    a.apply_validity(|x| {
+    a.apply_validity_mut(|x| {
         let mut a = std::mem::take(x);
         a = !a;
         *x = a;
diff --git a/tests/it/array/primitive/mod.rs b/tests/it/array/primitive/mod.rs
index 2af71d504a..bb0402f400 100644
--- a/tests/it/array/primitive/mod.rs
+++ b/tests/it/array/primitive/mod.rs
@@ -128,7 +128,7 @@ fn into_mut_3() {
 #[test]
 fn apply_values() {
     let mut a = PrimitiveArray::from([Some(1), Some(2), None]);
-    a.apply_values(|x| {
+    a.apply_values_mut(|x| {
         x[0] = 10;
     });
     let expected = PrimitiveArray::from([Some(10), Some(2), None]);
@@ -138,7 +138,7 @@ fn apply_values() {
 #[test]
 fn apply_validity() {
     let mut a = PrimitiveArray::from([Some(1), Some(2), None]);
-    a.apply_validity(|x| {
+    a.apply_validity_mut(|x| {
         let mut a = std::mem::take(x);
         a = !a;
         *x = a;
diff --git a/tests/it/bitmap/assign_ops.rs b/tests/it/bitmap/assign_ops.rs
index bea2b0dcbf..ef01fdd208 100644
--- a/tests/it/bitmap/assign_ops.rs
+++ b/tests/it/bitmap/assign_ops.rs
@@ -18,6 +18,24 @@ fn basics() {
     );
 }
 
+#[test]
+fn binary_assign_oob() {
+    // this check we don't have an oob access if the bitmaps are size T + 1
+    // and we do some slicing.
+    let a = MutableBitmap::from_iter(std::iter::repeat(true).take(65));
+    let b = MutableBitmap::from_iter(std::iter::repeat(true).take(65));
+
+    let a: Bitmap = a.into();
+    let a = a.slice(10, 20);
+
+    let b: Bitmap = b.into();
+    let b = b.slice(10, 20);
+
+    let mut a = a.make_mut();
+
+    binary_assign(&mut a, &b, |x: u64, y| x & y);
+}
+
 #[test]
 fn fast_paths() {
     let b = MutableBitmap::from([true, false]);