Repository: jorgecarleitao/arrow2
Pull request #1081 (2022-06-18 03:20:44)

Fixed error in chunked_mut bitmap

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1081](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (2d84f29) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/d4873fcb9b201d591049055c0e554dbdd2c45ad3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (d4873fc) will **increase** coverage by `0.01%`.
> The diff coverage is `83.33%`.

```diff
@@            Coverage Diff             @@
##             main    #1081      +/-   ##
==========================================
+ Coverage   80.96%   80.97%   +0.01%     
==========================================
  Files         366      366              
  Lines       35136    35126      -10     
==========================================
- Hits        28447    28444       -3     
+ Misses       6689     6682       -7     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/compute/arity\_assign.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2NvbXB1dGUvYXJpdHlfYXNzaWduLnJz) | `74.07% <75.00%> (+2.01%)` | :arrow_up: |
| [src/bitmap/utils/chunk\_iterator/chunks\_exact.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9jaHVua19pdGVyYXRvci9jaHVua3NfZXhhY3QucnM=) | `85.45% <100.00%> (+0.26%)` | :arrow_up: |
| [src/bitmap/utils/chunks\_exact\_mut.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9jaHVua3NfZXhhY3RfbXV0LnJz) | `77.41% <100.00%> (+2.41%)` | :arrow_up: |
| [src/bitmap/immutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC9pbW11dGFibGUucnM=) | `83.33% <0.00%> (+0.64%)` | :arrow_up: |
| [src/bitmap/assign\_ops.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC9hc3NpZ25fb3BzLnJz) | `86.29% <0.00%> (+0.80%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081?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/1081?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [d4873fc...2d84f29](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1081?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/src/bitmap/immutable.rs b/src/bitmap/immutable.rs
index b63310b07c..c9514e30b8 100644
--- a/src/bitmap/immutable.rs
+++ b/src/bitmap/immutable.rs
@@ -5,6 +5,7 @@ use std::sync::Arc;
 use crate::{buffer::bytes::Bytes, error::Error, trusted_len::TrustedLen};
 
 use super::{
+    chunk_iter_to_vec,
     utils::{count_zeros, fmt, get_bit, get_bit_unchecked, BitChunk, BitChunks, BitmapIter},
     MutableBitmap,
 };
@@ -240,8 +241,11 @@ impl Bitmap {
         match self.into_mut() {
             Either::Left(data) => {
                 if data.offset > 0 {
-                    // we have to recreate the bytes because a MutableBitmap does not have an `offset`.
-                    data.iter().collect()
+                    // re-align the bits (remove the offset)
+                    let chunks = data.chunks::<u64>();
+                    let remainder = chunks.remainder();
+                    let vec = chunk_iter_to_vec(chunks.chain(std::iter::once(remainder)));
+                    MutableBitmap::from_vec(vec, data.length)
                 } else {
                     MutableBitmap::from_vec(data.bytes.as_ref().to_vec(), data.length)
                 }
diff --git a/src/bitmap/utils/chunk_iterator/chunks_exact.rs b/src/bitmap/utils/chunk_iterator/chunks_exact.rs
index 5b2eaf4781..df4520bd20 100644
--- a/src/bitmap/utils/chunk_iterator/chunks_exact.rs
+++ b/src/bitmap/utils/chunk_iterator/chunks_exact.rs
@@ -16,14 +16,15 @@ pub struct BitChunksExact<'a, T: BitChunk> {
 impl<'a, T: BitChunk> BitChunksExact<'a, T> {
     /// Creates a new [`BitChunksExact`].
     #[inline]
-    pub fn new(slice: &'a [u8], len: usize) -> Self {
-        assert!(len <= slice.len() * 8);
+    pub fn new(bitmap: &'a [u8], length: usize) -> Self {
+        assert!(length <= bitmap.len() * 8);
         let size_of = std::mem::size_of::<T>();
 
-        let split = (len / 8 / size_of) * size_of;
-        let chunks = &slice[..split];
-        let remainder = &slice[split..];
-        let remainder_len = len - chunks.len() * 8;
+        let bitmap = &bitmap[..length.saturating_add(7) / 8];
+
+        let split = (length / 8 / size_of) * size_of;
+        let (chunks, remainder) = bitmap.split_at(split);
+        let remainder_len = length - chunks.len() * 8;
         let iter = chunks.chunks_exact(size_of);
 
         Self {
diff --git a/src/bitmap/utils/chunks_exact_mut.rs b/src/bitmap/utils/chunks_exact_mut.rs
index afc55ee093..7a5a91a128 100644
--- a/src/bitmap/utils/chunks_exact_mut.rs
+++ b/src/bitmap/utils/chunks_exact_mut.rs
@@ -17,8 +17,11 @@ impl<'a, T: BitChunk> BitChunksExactMut<'a, T> {
     /// Returns a new [`BitChunksExactMut`]
     #[inline]
     pub fn new(bitmap: &'a mut [u8], length: usize) -> Self {
+        assert!(length <= bitmap.len() * 8);
         let size_of = std::mem::size_of::<T>();
 
+        let bitmap = &mut bitmap[..length.saturating_add(7) / 8];
+
         let split = (length / 8 / size_of) * size_of;
         let (chunks, remainder) = bitmap.split_at_mut(split);
         let remainder_len = length - chunks.len() * 8;
diff --git a/src/compute/arity_assign.rs b/src/compute/arity_assign.rs
index b8746cdd81..ff504c9690 100644
--- a/src/compute/arity_assign.rs
+++ b/src/compute/arity_assign.rs
@@ -46,42 +46,37 @@ where
     // 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(|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);
+    if let Some(rhs) = rhs.validity() {
+        if lhs.validity().is_none() {
+            *lhs = lhs.with_validity(Some(rhs.clone()))
+        } else {
+            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()
-                        }
+                *bitmap = match owned_lhs.into_mut() {
+                    // we take alloc and write to new buffer
+                    Either::Left(immutable) => {
+                        // we allocate a new bitmap
+                        &immutable & rhs
                     }
-                });
-            }
+                    // we can mutate in place, happy days.
+                    Either::Right(mut mutable) => {
+                        let mut mutable_ref = &mut mutable;
+                        mutable_ref &= rhs;
+                        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);
 
-    match owned_lhs.into_mut() {
+    *lhs = match owned_lhs.into_mut() {
         // we take alloc and write to new buffer
         Either::Left(mut immutable) => {
             let values = immutable
@@ -91,7 +86,7 @@ where
                 .map(|(l, r)| op(*l, *r))
                 .collect::<Vec<_>>();
             immutable.set_values(values.into());
-            *lhs = immutable;
+            immutable
         }
         // we can mutate in place
         Either::Right(mut mutable) => {
@@ -100,7 +95,7 @@ where
                     .zip(rhs.values().iter())
                     .for_each(|(l, r)| *l = op(*l, *r))
             });
-            *lhs = mutable.into()
+            mutable.into()
         }
     }
 }


Test patch:
diff --git a/tests/it/bitmap/assign_ops.rs b/tests/it/bitmap/assign_ops.rs
index ef01fdd208..caf5586cf9 100644
--- a/tests/it/bitmap/assign_ops.rs
+++ b/tests/it/bitmap/assign_ops.rs
@@ -1,5 +1,9 @@
+use proptest::prelude::*;
+
 use arrow2::bitmap::{binary_assign, unary_assign, Bitmap, MutableBitmap};
 
+use crate::bitmap::bitmap_strategy;
+
 #[test]
 fn basics() {
     let mut b = MutableBitmap::from_iter(std::iter::repeat(true).take(10));
@@ -58,3 +62,18 @@ fn fast_paths() {
     let b = b | &c;
     assert_eq!(b, MutableBitmap::from_iter([true, false]));
 }
+
+proptest! {
+    /// Asserts that !bitmap equals all bits flipped
+    #[test]
+    #[cfg_attr(miri, ignore)] // miri and proptest do not work well :(
+    fn not(b in bitmap_strategy()) {
+        let not_b: MutableBitmap = b.iter().map(|x| !x).collect();
+
+        let mut b = b.make_mut();
+
+        unary_assign(&mut b, |x: u8| !x);
+
+        assert_eq!(b, not_b);
+    }
+}
diff --git a/tests/it/bitmap/utils/bit_chunks_exact.rs b/tests/it/bitmap/utils/bit_chunks_exact.rs
index 38280624dd..56ac806f16 100644
--- a/tests/it/bitmap/utils/bit_chunks_exact.rs
+++ b/tests/it/bitmap/utils/bit_chunks_exact.rs
@@ -29,5 +29,5 @@ fn remainder_u16() {
     );
     assert_eq!(iter.next(), Some(511));
     assert_eq!(iter.next(), None);
-    assert_eq!(iter.remainder(), 0b1101_1011_0000_0001u16);
+    assert_eq!(iter.remainder(), 1u16);
 }