Repository: jorgecarleitao/arrow2
Pull request #291 (2021-08-16 23:57:20)

Improved performance of concatenating non-aligned validities (15x)
This PR significantly improves the performance of concatenating arrays whose lengths are not a multiple of 8 by improving the performance of concatenating bitmaps.

Before this PR, we concatenated bitmaps by iterating bit by bit and setting bit by bit. However, there is a more efficient way of doing this via byte operations. Specifically, given a mutable bitmap `[10101000, --101010]` (`length = 8+6=14`) a bitmap can be concatenated to by shifts. E.g. `[00000011, a, b, ..., c]` can be concatenated to it by something like

* `00000011 << 6` and OR it
* merge `a` with `00000011` with an offset of 2 and append
* merge `b` with `a` with an offset of 2 and append
* ...
*  append `c`

This results in a significantly less number of instructions, lookups, etc.

This improves performance of almost operations that in some way concatenate validities. It includes:

* Growable API (`concat`, `filter`, `merge-sort`)
* lower-level bitmap concatenation

```bash
git checkout afb05d2511d495075180436dcd16af2e4b6ed71a
cargo bench --no-default-features --features benchmarks,compute --bench bitmap --bench concat --bench filter_kernels -- "2\^20"
git checkout improve_perf
cargo bench --no-default-features --features benchmarks,compute --bench bitmap --bench concat --bench filter_kernels -- "2\^20"
```

```
bitmap extend aligned 2^20                                                                             
                        time:   [3.7567 us 3.7847 us 3.8217 us]
                        change: [-3.0540% +1.6941% +6.6158%] (p = 0.49 > 0.05)

bitmap extend unaligned 2^20                                                                            
                        time:   [247.46 us 248.23 us 249.13 us]
                        change: [-75.411% -75.289% -75.172%] (p = 0.00 < 0.05)

bitmap extend_constant aligned 2^20                                                                             
                        time:   [2.6766 us 2.6822 us 2.6883 us]
                        change: [-99.536% -99.534% -99.532%] (p = 0.00 < 0.05)

bitmap extend_constant unaligned 2^20                                                                             
                        time:   [2.6916 us 2.6970 us 2.7026 us]
                        change: [-99.531% -99.529% -99.527%] (p = 0.00 < 0.05)

int32 concat aligned 2^20                                                                            
                        time:   [487.53 us 488.09 us 488.75 us]
                        change: [-93.566% -93.548% -93.530%] (p = 0.00 < 0.05)

int32 concat unaligned 2^20                                                                            
                        time:   [758.76 us 759.91 us 761.29 us]
                        change: [-89.977% -89.951% -89.923%] (p = 0.00 < 0.05)

boolean concat aligned 2^20                                                                            
                        time:   [224.02 us 224.50 us 225.02 us]
                        change: [-98.193% -98.187% -98.181%] (p = 0.00 < 0.05)

boolean concat unaligned 2^20                                                                            
                        time:   [708.63 us 710.23 us 712.14 us]
                        change: [-94.305% -94.286% -94.268%] (p = 0.00 < 0.05)

filter 2^20 f32         time:   [2.5137 ms 2.5199 ms 2.5274 ms]                             
                        change: [-2.6963% -2.3576% -1.9729%] (p = 0.00 < 0.05)

filter null 2^20 f32    time:   [7.6607 ms 7.6773 ms 7.6954 ms]                                 
                        change: [-12.051% -11.757% -11.473%] (p = 0.00 < 0.05)
```


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

[![Impacted file tree graph](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/graphs/tree.svg?width=650&height=150&src=pr&token=AgyTF60R3D&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)

```diff
@@           Coverage Diff           @@
##             main     #291   +/-   ##
=======================================
  Coverage   77.25%   77.25%           
=======================================
  Files         315      315           
  Lines       20791    20823   +32     
=======================================
+ Hits        16062    16087   +25     
- Misses       4729     4736    +7     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/bitmap/utils/chunk\_iterator/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9jaHVua19pdGVyYXRvci9tb2QucnM=) | `85.91% <ø> (ø)` | |
| [src/bitmap/utils/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9tb2QucnM=) | `100.00% <ø> (ø)` | |
| [src/bitmap/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC9tdXRhYmxlLnJz) | `87.30% <85.18%> (-0.76%)` | :arrow_down: |
| [src/array/growable/boolean.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2dyb3dhYmxlL2Jvb2xlYW4ucnM=) | `80.76% <100.00%> (ø)` | |
| [src/array/growable/utils.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2dyb3dhYmxlL3V0aWxzLnJz) | `100.00% <100.00%> (ø)` | |
| [tests/it/bitmap/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-dGVzdHMvaXQvYml0bWFwL211dGFibGUucnM=) | `100.00% <100.00%> (ø)` | |
| [src/io/json\_integration/write.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2pzb25faW50ZWdyYXRpb24vd3JpdGUucnM=) | `0.00% <0.00%> (-6.25%)` | :arrow_down: |
| [src/io/csv/write/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2Nzdi93cml0ZS9tb2QucnM=) | `72.00% <0.00%> (-4.00%)` | :arrow_down: |
| [src/io/json/read/reader.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/291/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2pzb24vcmVhZC9yZWFkZXIucnM=) | `74.32% <0.00%> (-1.36%)` | :arrow_down: |

------

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

jorgecarleitao: cc @ritchie46 and @Dandandan , since you like these things :)
ritchie46: Love it! I see some interesting bit comments. Have you got a summary of what you do? A `memcpy` instead of iterators?
Dandandan: > cc @ritchie46 and @Dandandan , since you like these things :)

Amazing 😎
jorgecarleitao: > Love it! I see some interesting bit comments. Have you got a summary of what you do? A `memcpy` instead of iterators?

:) Updated the description with the idea 👍 

Patch:
diff --git a/Cargo.toml b/Cargo.toml
index b6baebb644..e779422396 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -176,3 +176,11 @@ harness = false
 [[bench]]
 name = "arithmetic_kernels"
 harness = false
+
+[[bench]]
+name = "bitmap"
+harness = false
+
+[[bench]]
+name = "concat"
+harness = false
diff --git a/benches/bitmap.rs b/benches/bitmap.rs
new file mode 100644
index 0000000000..afe1c62747
--- /dev/null
+++ b/benches/bitmap.rs
@@ -0,0 +1,60 @@
+extern crate arrow2;
+
+use std::iter::FromIterator;
+
+use arrow2::bitmap::*;
+
+use criterion::{criterion_group, criterion_main, Criterion};
+
+//
+
+fn add_benchmark(c: &mut Criterion) {
+    (10..=20).step_by(2).for_each(|log2_size| {
+        let size = 2usize.pow(log2_size);
+
+        let bitmap2 = Bitmap::from_iter((0..size).into_iter().map(|x| x % 3 == 0));
+
+        c.bench_function(&format!("bitmap extend aligned 2^{}", log2_size), |b| {
+            let mut bitmap1 = MutableBitmap::new();
+            b.iter(|| {
+                bitmap1.extend_from_bitmap(&bitmap2);
+                bitmap1.clear();
+            })
+        });
+
+        c.bench_function(&format!("bitmap extend unaligned 2^{}", log2_size), |b| {
+            let mut bitmap1 = MutableBitmap::with_capacity(1);
+            b.iter(|| {
+                bitmap1.push(true);
+                bitmap1.extend_from_bitmap(&bitmap2);
+                bitmap1.clear();
+            })
+        });
+
+        c.bench_function(
+            &format!("bitmap extend_constant aligned 2^{}", log2_size),
+            |b| {
+                let mut bitmap1 = MutableBitmap::new();
+                b.iter(|| {
+                    bitmap1.extend_constant(size, true);
+                    bitmap1.clear();
+                })
+            },
+        );
+
+        c.bench_function(
+            &format!("bitmap extend_constant unaligned 2^{}", log2_size),
+            |b| {
+                let mut bitmap1 = MutableBitmap::with_capacity(1);
+                b.iter(|| {
+                    bitmap1.push(true);
+                    bitmap1.extend_constant(size, true);
+                    bitmap1.clear();
+                })
+            },
+        );
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);
diff --git a/benches/concat.rs b/benches/concat.rs
new file mode 100644
index 0000000000..17273a2769
--- /dev/null
+++ b/benches/concat.rs
@@ -0,0 +1,52 @@
+extern crate arrow2;
+
+use arrow2::{
+    compute::concat,
+    datatypes::DataType,
+    util::bench_util::{create_boolean_array, create_primitive_array},
+};
+
+use criterion::{criterion_group, criterion_main, Criterion};
+
+fn add_benchmark(c: &mut Criterion) {
+    (20..=20).step_by(2).for_each(|log2_size| {
+        let size = 2usize.pow(log2_size);
+
+        let array1 = create_primitive_array::<i32>(8, DataType::Int32, 0.5);
+        let array2 = create_primitive_array::<i32>(size + 1, DataType::Int32, 0.5);
+
+        c.bench_function(&format!("int32 concat aligned 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let _ = concat::concatenate(&[&array1, &array2]);
+            })
+        });
+
+        let array1 = create_primitive_array::<i32>(9, DataType::Int32, 0.5);
+
+        c.bench_function(&format!("int32 concat unaligned 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let _ = concat::concatenate(&[&array1, &array2]);
+            })
+        });
+
+        let array1 = create_boolean_array(8, 0.5, 0.5);
+        let array2 = create_boolean_array(size + 1, 0.5, 0.5);
+
+        c.bench_function(&format!("boolean concat aligned 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let _ = concat::concatenate(&[&array1, &array2]);
+            })
+        });
+
+        let array1 = create_boolean_array(9, 0.5, 0.5);
+
+        c.bench_function(&format!("boolean concat unaligned 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let _ = concat::concatenate(&[&array1, &array2]);
+            })
+        });
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);
diff --git a/src/array/growable/boolean.rs b/src/array/growable/boolean.rs
index 4f02f82e67..0492465283 100644
--- a/src/array/growable/boolean.rs
+++ b/src/array/growable/boolean.rs
@@ -55,8 +55,9 @@ impl<'a> Growable<'a> for GrowableBoolean<'a> {
 
         let array = self.arrays[index];
         let values = array.values();
-        let iter = (start..start + len).map(|i| values.get_bit(i));
-        unsafe { self.values.extend_from_trusted_len_iter_unchecked(iter) };
+
+        let (slice, offset, _) = values.as_slice();
+        self.values.extend_from_slice(slice, start + offset, len);
     }
 
     fn extend_validity(&mut self, additional: usize) {
diff --git a/src/array/growable/utils.rs b/src/array/growable/utils.rs
index 1b6f5ff55e..6a67988e2d 100644
--- a/src/array/growable/utils.rs
+++ b/src/array/growable/utils.rs
@@ -24,17 +24,12 @@ pub(super) fn build_extend_null_bits(array: &dyn Array, use_validity: bool) -> E
     if let Some(bitmap) = array.validity() {
         Box::new(move |validity, start, len| {
             assert!(start + len <= bitmap.len());
-            unsafe {
-                let iter = (start..start + len).map(|i| bitmap.get_bit_unchecked(i));
-                validity.extend_from_trusted_len_iter_unchecked(iter);
-            };
+            let (slice, offset, _) = bitmap.as_slice();
+            validity.extend_from_slice(slice, start + offset, len);
         })
     } else if use_validity {
         Box::new(|validity, _, len| {
-            let iter = (0..len).map(|_| true);
-            unsafe {
-                validity.extend_from_trusted_len_iter_unchecked(iter);
-            };
+            validity.extend_constant(len, true);
         })
     } else {
         Box::new(|_, _, _| {})
@@ -51,10 +46,8 @@ pub(super) fn extend_validity(
 ) {
     if let Some(bitmap) = validity {
         assert!(start + len <= bitmap.len());
-        unsafe {
-            let iter = (start..start + len).map(|i| bitmap.get_bit_unchecked(i));
-            mutable_validity.extend_from_trusted_len_iter_unchecked(iter);
-        };
+        let (slice, offset, _) = bitmap.as_slice();
+        mutable_validity.extend_from_slice(slice, start + offset, len);
     } else if use_validity {
         mutable_validity.extend_constant(len, true);
     };
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index 69220c3cfa..2e80840249 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -222,9 +222,9 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
         if value.is_none() && self.validity.is_none() {
             // When the validity is None, all elements so far are valid. When one of the elements is set fo null,
             // the validity must be initialized.
-            self.validity = Some(MutableBitmap::from_trusted_len_iter(
-                std::iter::repeat(true).take(self.len()),
-            ));
+            let mut validity = MutableBitmap::new();
+            validity.extend_constant(self.len(), true);
+            self.validity = Some(validity);
         }
         if let Some(x) = self.validity.as_mut() {
             x.set(index, value.is_some())
diff --git a/src/bitmap/immutable.rs b/src/bitmap/immutable.rs
index a6a3a70b2e..dccc52bdd6 100644
--- a/src/bitmap/immutable.rs
+++ b/src/bitmap/immutable.rs
@@ -26,7 +26,8 @@ pub struct Bitmap {
 
 impl std::fmt::Debug for Bitmap {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        fmt(&self.bytes, self.offset, self.length, f)
+        let (bytes, offset, len) = self.as_slice();
+        fmt(bytes, offset, len, f)
     }
 }
 
@@ -111,8 +112,8 @@ impl Bitmap {
     /// exceeds the allocated capacity of `self`.
     #[inline]
     pub fn slice(mut self, offset: usize, length: usize) -> Self {
+        assert!(offset + length <= self.length);
         self.offset += offset;
-        assert!(self.offset + length <= self.bytes.len() * 8);
         self.length = length;
         self.null_count = null_count(&self.bytes, self.offset, self.length);
         self
diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs
index 780a479e35..e97934b8d1 100644
--- a/src/bitmap/mutable.rs
+++ b/src/bitmap/mutable.rs
@@ -1,5 +1,6 @@
 use std::iter::FromIterator;
 
+use crate::bitmap::utils::merge_reversed;
 use crate::{buffer::MutableBuffer, trusted_len::TrustedLen};
 
 use super::utils::{fmt, get_bit, null_count, set, set_bit, BitmapIter};
@@ -39,6 +40,13 @@ impl MutableBitmap {
         }
     }
 
+    /// Empties the [`MutableBitmap`].
+    #[inline]
+    pub fn clear(&mut self) {
+        self.length = 0;
+        self.buffer.clear();
+    }
+
     /// Initializes a zeroed [`MutableBitmap`].
     #[inline]
     pub fn from_len_zeroed(length: usize) -> Self {
@@ -123,16 +131,63 @@ impl MutableBitmap {
         self.length = len;
     }
 
+    fn extend_set(&mut self, mut additional: usize) {
+        let offset = self.length % 8;
+        let added = if offset != 0 {
+            // offset != 0 => at least one byte in the buffer
+            let last_index = self.buffer.len() - 1;
+            let last = &mut self.buffer[last_index];
+
+            let remaining = 0b11111111u8;
+            let remaining = remaining >> 8usize.saturating_sub(additional);
+            let remaining = remaining << offset;
+            *last |= remaining;
+            std::cmp::min(additional, 8 - offset)
+        } else {
+            0
+        };
+        self.length += added;
+        additional = additional.saturating_sub(added);
+        if additional > 0 {
+            debug_assert_eq!(self.length % 8, 0);
+            let existing = self.buffer.len();
+            let required = (self.length + additional).saturating_add(7) / 8;
+            // add remaining as full bytes
+            self.buffer.extend_from_trusted_len_iter(
+                std::iter::repeat(0b11111111u8).take(required - existing),
+            );
+            self.length += additional;
+        }
+    }
+
+    fn extend_unset(&mut self, mut additional: usize) {
+        let offset = self.length % 8;
+        let added = if offset != 0 {
+            // offset != 0 => at least one byte in the buffer
+            let last_index = self.buffer.len() - 1;
+            let last = &mut self.buffer[last_index];
+            *last &= 0b11111111u8 >> (8 - offset); // unset them
+            std::cmp::min(additional, 8 - offset)
+        } else {
+            0
+        };
+        self.length += added;
+        additional = additional.saturating_sub(added);
+        if additional > 0 {
+            debug_assert_eq!(self.length % 8, 0);
+            self.buffer
+                .resize((self.length + additional).saturating_add(7) / 8, 0);
+            self.length += additional;
+        }
+    }
+
     /// Extends [`MutableBitmap`] by `additional` values of constant `value`.
     #[inline]
     pub fn extend_constant(&mut self, additional: usize, value: bool) {
         if value {
-            let iter = std::iter::repeat(true).take(additional);
-            self.extend_from_trusted_len_iter(iter);
+            self.extend_set(additional)
         } else {
-            self.buffer
-                .resize((self.length + additional).saturating_add(7) / 8, 0);
-            self.length += additional;
+            self.extend_unset(additional)
         }
     }
 
@@ -413,6 +468,42 @@ impl MutableBitmap {
         Ok(Self { buffer, length })
     }
 
+    fn extend_unaligned(&mut self, slice: &[u8], offset: usize, length: usize) {
+        let own_offset = self.length % 8;
+        // e.g.
+        // [a, b, --101010]     <- to be extended
+        // [00111111, 11010101] <- to extend
+        // [a, b, 11101010, --001111] expected result
+        let aligned_offset = offset / 8;
+        let bytes_len = length.saturating_add(7) / 8;
+        let items = &slice[aligned_offset..aligned_offset + bytes_len];
+        // self has some offset => we need to shift all `items`, and merge the first
+        let buffer = self.buffer.as_mut_slice();
+        let last = &mut buffer[buffer.len() - 1];
+
+        // --101010 | 00111111 << 6 = 11101010
+        // erase previous
+        *last &= 0b11111111u8 >> (8 - own_offset); // unset before setting
+        *last |= items[0] << own_offset;
+
+        let remaining = [items[items.len() - 1], 0];
+        let bytes = items
+            .windows(2)
+            .chain(std::iter::once(remaining.as_ref()))
+            .map(|w| merge_reversed(w[0], w[1], 8 - own_offset));
+        self.buffer.extend_from_trusted_len_iter(bytes);
+
+        self.length += length;
+    }
+
+    fn extend_aligned(&mut self, slice: &[u8], offset: usize, length: usize) {
+        let aligned_offset = offset / 8;
+        let bytes_len = length.saturating_add(7) / 8;
+        let items = &slice[aligned_offset..aligned_offset + bytes_len];
+        self.buffer.extend_from_slice(items);
+        self.length += length;
+    }
+
     /// Extends the [`MutableBitmap`] from a slice of bytes with optional offset.
     /// This is the fastest way to extend a [`MutableBitmap`].
     /// # Implementation
@@ -421,16 +512,14 @@ impl MutableBitmap {
     #[inline]
     pub fn extend_from_slice(&mut self, slice: &[u8], offset: usize, length: usize) {
         assert!(offset + length <= slice.len() * 8);
+        if length == 0 {
+            return;
+        };
         let is_aligned = self.length % 8 == 0;
         let other_is_aligned = offset % 8 == 0;
         match (is_aligned, other_is_aligned) {
-            (true, true) => {
-                let aligned_offset = offset / 8;
-                let bytes_len = length.saturating_add(7) / 8;
-                let items = &slice[aligned_offset..aligned_offset + bytes_len];
-                self.buffer.extend_from_slice(items);
-                self.length += length;
-            }
+            (true, true) => self.extend_aligned(slice, offset, length),
+            (false, true) => self.extend_unaligned(slice, offset, length),
             // todo: further optimize the other branches.
             _ => self.extend_from_trusted_len_iter(BitmapIter::new(slice, offset, length)),
         }
diff --git a/src/bitmap/utils/chunk_iterator/mod.rs b/src/bitmap/utils/chunk_iterator/mod.rs
index 8528533a21..47390955fb 100644
--- a/src/bitmap/utils/chunk_iterator/mod.rs
+++ b/src/bitmap/utils/chunk_iterator/mod.rs
@@ -7,7 +7,7 @@ pub use crate::types::BitChunk;
 pub use chunks_exact::BitChunksExact;
 
 use crate::{trusted_len::TrustedLen, types::BitChunkIter};
-use merge::merge_reversed;
+pub(crate) use merge::merge_reversed;
 
 pub trait BitChunkIterExact<B: BitChunk>: Iterator<Item = B> {
     fn remainder(&self) -> B;
diff --git a/src/bitmap/utils/mod.rs b/src/bitmap/utils/mod.rs
index 979900b234..5d866886eb 100644
--- a/src/bitmap/utils/mod.rs
+++ b/src/bitmap/utils/mod.rs
@@ -4,6 +4,7 @@ mod iterator;
 mod slice_iterator;
 mod zip_validity;
 
+pub(crate) use chunk_iterator::merge_reversed;
 pub use chunk_iterator::{BitChunk, BitChunkIterExact, BitChunks, BitChunksExact};
 pub use fmt::fmt;
 pub use iterator::BitmapIter;
diff --git a/src/bitmap/utils/slice_iterator.rs b/src/bitmap/utils/slice_iterator.rs
index 686c5f1c16..c081a79198 100644
--- a/src/bitmap/utils/slice_iterator.rs
+++ b/src/bitmap/utils/slice_iterator.rs
@@ -87,7 +87,7 @@ impl<'a> Iterator for SlicesIterator<'a> {
                 // at the beginning of a byte => try to skip it all together
                 match (self.on_region, self.current_byte) {
                     (true, &255u8) => {
-                        self.len += 8;
+                        self.len = std::cmp::min(self.max_len - self.start, self.len + 8);
                         match self.values.next() {
                             Some(v) => self.current_byte = v,
                             None => return self.finish(),
@@ -95,7 +95,7 @@ impl<'a> Iterator for SlicesIterator<'a> {
                         continue;
                     }
                     (false, &0) => {
-                        self.len += 8;
+                        self.len = std::cmp::min(self.max_len - self.start, self.len + 8);
                         match self.values.next() {
                             Some(v) => self.current_byte = v,
                             None => return self.finish(),


Test patch:
diff --git a/tests/it/array/primitive/mutable.rs b/tests/it/array/primitive/mutable.rs
index 8dbff29823..0adb4a534d 100644
--- a/tests/it/array/primitive/mutable.rs
+++ b/tests/it/array/primitive/mutable.rs
@@ -2,19 +2,60 @@ use arrow2::{
     array::*,
     bitmap::{Bitmap, MutableBitmap},
     buffer::MutableBuffer,
+    datatypes::DataType,
 };
 use std::iter::FromIterator;
 
+#[test]
+fn from_and_into_data() {
+    let a = MutablePrimitiveArray::from_data(
+        DataType::Int32,
+        MutableBuffer::from([1i32, 0]),
+        Some(MutableBitmap::from([true, false])),
+    );
+    assert_eq!(a.len(), 2);
+    let (a, b, c) = a.into_data();
+    assert_eq!(a, DataType::Int32);
+    assert_eq!(b, MutableBuffer::from([1i32, 0]));
+    assert_eq!(c, Some(MutableBitmap::from([true, false])));
+}
+
+#[test]
+fn to() {
+    let a = MutablePrimitiveArray::from_data(
+        DataType::Int32,
+        MutableBuffer::from([1i32, 0]),
+        Some(MutableBitmap::from([true, false])),
+    );
+    let a = a.to(DataType::Date32);
+    assert_eq!(a.data_type(), &DataType::Date32);
+}
+
+#[test]
+fn values_mut_slice() {
+    let mut a = MutablePrimitiveArray::from_data(
+        DataType::Int32,
+        MutableBuffer::from([1i32, 0]),
+        Some(MutableBitmap::from([true, false])),
+    );
+    let values = a.values_mut_slice();
+
+    values[0] = 10;
+    assert_eq!(a.values()[0], 10);
+}
+
 #[test]
 fn push() {
     let mut a = MutablePrimitiveArray::<i32>::new();
     a.push(Some(1));
     a.push(None);
-    assert_eq!(a.len(), 2);
+    a.push_null();
+    assert_eq!(a.len(), 3);
     assert!(a.is_valid(0));
     assert!(!a.is_valid(1));
+    assert!(!a.is_valid(2));
 
-    assert_eq!(a.values(), &MutableBuffer::from([1, 0]));
+    assert_eq!(a.values(), &MutableBuffer::from([1, 0, 0]));
 }
 
 #[test]
@@ -29,6 +70,17 @@ fn set() {
     assert!(a.is_valid(1));
 
     assert_eq!(a.values(), &MutableBuffer::from([2, 1]));
+
+    let mut a = MutablePrimitiveArray::<i32>::from_slice([1, 2]);
+
+    a.set(0, Some(2));
+    a.set(1, None);
+
+    assert_eq!(a.len(), 2);
+    assert!(a.is_valid(0));
+    assert!(!a.is_valid(1));
+
+    assert_eq!(a.values(), &MutableBuffer::from([2, 0]));
 }
 
 #[test]
@@ -72,3 +124,19 @@ fn extend_trusted_len() {
     );
     assert_eq!(a.values(), &MutableBuffer::<i32>::from([1, 2, 0, 4]));
 }
+
+#[test]
+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);
+    a.set_validity(Some(MutableBitmap::from([false, true])));
+    assert_eq!(a.validity(), &Some(MutableBitmap::from([false, true])));
+}
+
+#[test]
+fn set_values() {
+    let mut a = MutablePrimitiveArray::<i32>::from_slice([1, 2]);
+    a.set_values(MutableBuffer::from([1, 3]));
+    assert_eq!(a.values().as_slice(), [1, 3]);
+}
diff --git a/tests/it/bitmap/mutable.rs b/tests/it/bitmap/mutable.rs
index 419c0ab869..fd3d719fbf 100644
--- a/tests/it/bitmap/mutable.rs
+++ b/tests/it/bitmap/mutable.rs
@@ -1,4 +1,7 @@
-use arrow2::bitmap::{Bitmap, MutableBitmap};
+use arrow2::{
+    bitmap::{Bitmap, MutableBitmap},
+    buffer::MutableBuffer,
+};
 
 #[test]
 fn trusted_len() {
@@ -156,6 +159,31 @@ fn extend_from_bitmap() {
     assert_eq!(bitmap.as_slice()[0], 0b00101101);
 }
 
+#[test]
+fn extend_from_bitmap_offset() {
+    let other = Bitmap::from_u8_slice(&[0b00111111], 8);
+    let mut bitmap = MutableBitmap::from_buffer(MutableBuffer::from(&[1, 0, 0b00101010]), 22);
+
+    // call is optimized to perform a memcopy
+    bitmap.extend_from_bitmap(&other);
+
+    assert_eq!(bitmap.len(), 22 + 8);
+    assert_eq!(bitmap.as_slice(), &[1, 0, 0b11101010, 0b00001111]);
+
+    // more than one byte
+    let other = Bitmap::from_u8_slice(&[0b00111111, 0b00001111, 0b0001100], 20);
+    let mut bitmap = MutableBitmap::from_buffer(MutableBuffer::from(&[1, 0, 0b00101010]), 22);
+
+    // call is optimized to perform a memcopy
+    bitmap.extend_from_bitmap(&other);
+
+    assert_eq!(bitmap.len(), 22 + 20);
+    assert_eq!(
+        bitmap.as_slice(),
+        &[1, 0, 0b11101010, 0b11001111, 0b0000011, 0b0000011]
+    );
+}
+
 #[test]
 fn debug() {
     let mut b = MutableBitmap::new();
@@ -173,3 +201,77 @@ fn debug() {
     b.push(true);
     assert_eq!(format!("{:?}", b), "[0b11000001, 0b_______1]");
 }
+
+#[test]
+fn extend_set() {
+    let mut b = MutableBitmap::new();
+    b.extend_constant(6, true);
+    assert_eq!(b.as_slice(), &[0b11111111]);
+    assert_eq!(b.len(), 6);
+
+    let mut b = MutableBitmap::from(&[false]);
+    b.extend_constant(6, true);
+    assert_eq!(b.as_slice(), &[0b01111110]);
+    assert_eq!(b.len(), 1 + 6);
+
+    let mut b = MutableBitmap::from(&[false]);
+    b.extend_constant(9, true);
+    assert_eq!(b.as_slice(), &[0b11111110, 0b11111111]);
+    assert_eq!(b.len(), 1 + 9);
+
+    let mut b = MutableBitmap::from(&[false, false, false, false]);
+    b.extend_constant(2, true);
+    assert_eq!(b.as_slice(), &[0b00110000]);
+    assert_eq!(b.len(), 4 + 2);
+
+    let mut b = MutableBitmap::from(&[true, true]);
+    b.extend_constant(3, true);
+    assert_eq!(b.as_slice(), &[0b00011111]);
+    assert_eq!(b.len(), 2 + 3);
+}
+
+#[test]
+fn extend_unset() {
+    let mut b = MutableBitmap::new();
+    b.extend_constant(6, false);
+    assert_eq!(b.as_slice(), &[0b0000000]);
+    assert_eq!(b.len(), 6);
+
+    let mut b = MutableBitmap::from(&[true]);
+    b.extend_constant(6, false);
+    assert_eq!(b.as_slice(), &[0b00000001]);
+    assert_eq!(b.len(), 1 + 6);
+
+    let mut b = MutableBitmap::from(&[true]);
+    b.extend_constant(9, false);
+    assert_eq!(b.as_slice(), &[0b0000001, 0b00000000]);
+    assert_eq!(b.len(), 1 + 9);
+
+    let mut b = MutableBitmap::from(&[true, true, true, true]);
+    b.extend_constant(2, false);
+    assert_eq!(b.as_slice(), &[0b00001111]);
+    assert_eq!(b.len(), 4 + 2);
+}
+
+#[test]
+fn extend_bitmap() {
+    let mut b = MutableBitmap::from(&[true]);
+    b.extend_from_slice(&[0b00011001], 0, 6);
+    assert_eq!(b.as_slice(), &[0b00110011]);
+    assert_eq!(b.len(), 1 + 6);
+
+    let mut b = MutableBitmap::from(&[true]);
+    b.extend_from_slice(&[0b00011001, 0b00011001], 0, 9);
+    assert_eq!(b.as_slice(), &[0b00110011, 0b00110010]);
+    assert_eq!(b.len(), 1 + 9);
+
+    let mut b = MutableBitmap::from(&[true, true, true, true]);
+    b.extend_from_slice(&[0b00011001, 0b00011001], 0, 9);
+    assert_eq!(b.as_slice(), &[0b10011111, 0b10010001]);
+    assert_eq!(b.len(), 4 + 9);
+
+    let mut b = MutableBitmap::from(&[true, true, true, true, true]);
+    b.extend_from_slice(&[0b00001011], 0, 4);
+    assert_eq!(b.as_slice(), &[0b01111111, 0b00000001]);
+    assert_eq!(b.len(), 5 + 4);
+}
diff --git a/tests/it/bitmap/utils/slice_iterator.rs b/tests/it/bitmap/utils/slice_iterator.rs
index 47bfcf2fbd..ea60940c41 100644
--- a/tests/it/bitmap/utils/slice_iterator.rs
+++ b/tests/it/bitmap/utils/slice_iterator.rs
@@ -1,144 +1,169 @@
+use rand::distributions::{Bernoulli, Uniform};
+use rand::prelude::StdRng;
+use rand::Rng;
+use rand::SeedableRng;
+
 use arrow2::bitmap::utils::SlicesIterator;
 use arrow2::bitmap::Bitmap;
 
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn check_invariant() {
-        let values = (0..8).map(|i| i % 2 != 0).collect::<Bitmap>();
-        let iter = SlicesIterator::new(&values);
+#[test]
+fn check_invariant() {
+    let values = (0..8).map(|i| i % 2 != 0).collect::<Bitmap>();
+    let iter = SlicesIterator::new(&values);
 
-        let slots = iter.slots();
+    let slots = iter.slots();
 
-        let slices = iter.collect::<Vec<_>>();
+    let slices = iter.collect::<Vec<_>>();
 
-        assert_eq!(slices, vec![(1, 1), (3, 1), (5, 1), (7, 1)]);
+    assert_eq!(slices, vec![(1, 1), (3, 1), (5, 1), (7, 1)]);
 
-        let mut sum = 0;
-        for (_, len) in slices {
-            sum += len;
-        }
-        assert_eq!(sum, slots);
+    let mut sum = 0;
+    for (_, len) in slices {
+        sum += len;
     }
+    assert_eq!(sum, slots);
+}
 
-    #[test]
-    fn single_set() {
-        let values = (0..16).map(|i| i == 1).collect::<Bitmap>();
+#[test]
+fn single_set() {
+    let values = (0..16).map(|i| i == 1).collect::<Bitmap>();
 
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        let chunks = iter.collect::<Vec<_>>();
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    let chunks = iter.collect::<Vec<_>>();
 
-        assert_eq!(chunks, vec![(1, 1)]);
-        assert_eq!(count, 1);
-    }
+    assert_eq!(chunks, vec![(1, 1)]);
+    assert_eq!(count, 1);
+}
 
-    #[test]
-    fn single_unset() {
-        let values = (0..64).map(|i| i != 1).collect::<Bitmap>();
+#[test]
+fn single_unset() {
+    let values = (0..64).map(|i| i != 1).collect::<Bitmap>();
 
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        let chunks = iter.collect::<Vec<_>>();
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    let chunks = iter.collect::<Vec<_>>();
 
-        assert_eq!(chunks, vec![(0, 1), (2, 62)]);
-        assert_eq!(count, 64 - 1);
-    }
+    assert_eq!(chunks, vec![(0, 1), (2, 62)]);
+    assert_eq!(count, 64 - 1);
+}
 
-    #[test]
-    fn generic() {
-        let values = (0..130).map(|i| i % 62 != 0).collect::<Bitmap>();
+#[test]
+fn generic() {
+    let values = (0..130).map(|i| i % 62 != 0).collect::<Bitmap>();
 
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        let chunks = iter.collect::<Vec<_>>();
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    let chunks = iter.collect::<Vec<_>>();
 
-        assert_eq!(chunks, vec![(1, 61), (63, 61), (125, 5)]);
-        assert_eq!(count, 61 + 61 + 5);
-    }
+    assert_eq!(chunks, vec![(1, 61), (63, 61), (125, 5)]);
+    assert_eq!(count, 61 + 61 + 5);
+}
 
-    #[test]
-    fn incomplete_byte() {
-        let values = (0..6).map(|i| i == 1).collect::<Bitmap>();
+#[test]
+fn incomplete_byte() {
+    let values = (0..6).map(|i| i == 1).collect::<Bitmap>();
 
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        let chunks = iter.collect::<Vec<_>>();
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    let chunks = iter.collect::<Vec<_>>();
 
-        assert_eq!(chunks, vec![(1, 1)]);
-        assert_eq!(count, 1);
-    }
+    assert_eq!(chunks, vec![(1, 1)]);
+    assert_eq!(count, 1);
+}
 
-    #[test]
-    fn incomplete_byte1() {
-        let values = (0..12).map(|i| i == 9).collect::<Bitmap>();
+#[test]
+fn incomplete_byte1() {
+    let values = (0..12).map(|i| i == 9).collect::<Bitmap>();
 
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        let chunks = iter.collect::<Vec<_>>();
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    let chunks = iter.collect::<Vec<_>>();
 
-        assert_eq!(chunks, vec![(9, 1)]);
-        assert_eq!(count, 1);
-    }
+    assert_eq!(chunks, vec![(9, 1)]);
+    assert_eq!(count, 1);
+}
 
-    #[test]
-    fn end_of_byte() {
-        let values = (0..16).map(|i| i != 7).collect::<Bitmap>();
+#[test]
+fn end_of_byte() {
+    let values = (0..16).map(|i| i != 7).collect::<Bitmap>();
 
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        let chunks = iter.collect::<Vec<_>>();
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    let chunks = iter.collect::<Vec<_>>();
 
-        assert_eq!(chunks, vec![(0, 7), (8, 8)]);
-        assert_eq!(count, 15);
-    }
+    assert_eq!(chunks, vec![(0, 7), (8, 8)]);
+    assert_eq!(count, 15);
+}
 
-    #[test]
-    fn bla() {
-        let values = vec![true, true, true, true, true, true, true, false]
-            .into_iter()
-            .collect::<Bitmap>();
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        assert_eq!(values.null_count() + iter.slots(), values.len());
+#[test]
+fn bla() {
+    let values = vec![true, true, true, true, true, true, true, false]
+        .into_iter()
+        .collect::<Bitmap>();
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    assert_eq!(values.null_count() + iter.slots(), values.len());
 
-        let total = iter.into_iter().fold(0, |acc, x| acc + x.1);
+    let total = iter.into_iter().fold(0, |acc, x| acc + x.1);
 
-        assert_eq!(count, total);
-    }
+    assert_eq!(count, total);
+}
 
-    #[test]
-    fn past_end_should_not_be_returned() {
-        let values = Bitmap::from_u8_slice(&[0b11111010], 3);
-        let iter = SlicesIterator::new(&values);
-        let count = iter.slots();
-        assert_eq!(values.null_count() + iter.slots(), values.len());
+#[test]
+fn past_end_should_not_be_returned() {
+    let values = Bitmap::from_u8_slice(&[0b11111010], 3);
+    let iter = SlicesIterator::new(&values);
+    let count = iter.slots();
+    assert_eq!(values.null_count() + iter.slots(), values.len());
 
-        let total = iter.into_iter().fold(0, |acc, x| acc + x.1);
+    let total = iter.into_iter().fold(0, |acc, x| acc + x.1);
 
-        assert_eq!(count, total);
-    }
+    assert_eq!(count, total);
+}
 
-    #[test]
-    fn sliced() {
-        let values = Bitmap::from_u8_slice(&[0b11111010, 0b11111011], 16);
-        let values = values.slice(8, 2);
-        let iter = SlicesIterator::new(&values);
+#[test]
+fn sliced() {
+    let values = Bitmap::from_u8_slice(&[0b11111010, 0b11111011], 16);
+    let values = values.slice(8, 2);
+    let iter = SlicesIterator::new(&values);
 
-        let chunks = iter.collect::<Vec<_>>();
+    let chunks = iter.collect::<Vec<_>>();
 
-        // the first "11" in the second byte
-        assert_eq!(chunks, vec![(0, 2)]);
-    }
+    // the first "11" in the second byte
+    assert_eq!(chunks, vec![(0, 2)]);
+}
+
+#[test]
+fn remainder_1() {
+    let values = Bitmap::from_u8_slice(&[0, 0, 0b00000000, 0b00010101], 27);
+    let values = values.slice(22, 5);
+    let iter = SlicesIterator::new(&values);
+    let chunks = iter.collect::<Vec<_>>();
+    assert_eq!(chunks, vec![(2, 1), (4, 1)]);
+}
 
-    #[test]
-    fn remainder_1() {
-        let values = Bitmap::from_u8_slice(&[0, 0, 0b00000000, 0b00010101], 27);
-        let values = values.slice(22, 5);
-        let iter = SlicesIterator::new(&values);
-        let chunks = iter.collect::<Vec<_>>();
-        assert_eq!(chunks, vec![(2, 1), (4, 1)]);
+#[test]
+fn filter_slices() {
+    let mut rng = StdRng::seed_from_u64(42);
+    let length = 500;
+
+    let mask: Bitmap = (0..length)
+        .map(|_| {
+            let v: bool = (&mut rng).sample(Bernoulli::new(0.5).unwrap());
+            v
+        })
+        .collect();
+
+    for offset in 100usize..(length - 1) {
+        let len = (&mut rng).sample(Uniform::new(0, length - offset));
+        let mask_s = mask.clone().slice(offset, len);
+
+        let iter = SlicesIterator::new(&mask_s);
+        iter.for_each(|(start, slice_len)| {
+            if start + slice_len > len {
+                panic!("Fail")
+            }
+        });
     }
 }