Improved performance of unary _not_ for aligned bitmaps (3x) (#365)

* Added benches for bitmap unary.

* Improved perf of binary unary op.

Cargo.toml

@@ -183,3 +183,7 @@ harness = false
 [[bench]]
 name = "concat"
 harness = false
+
+[[bench]]
+name = "bitmap_ops"
+harness = false

benches/bitmap_ops.rs

@@ -0,0 +1,37 @@
+use arrow2::bitmap::Bitmap;
+
+use criterion::{criterion_group, criterion_main, Criterion};
+
+fn bench_arrow2(lhs: &Bitmap, rhs: &Bitmap) {
+    let r = lhs | rhs;
+    assert!(r.null_count() > 0);
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    (10..=20).step_by(2).for_each(|log2_size| {
+        let size = 2usize.pow(log2_size);
+
+        let bitmap: Bitmap = (0..size).into_iter().map(|x| x % 3 == 0).collect();
+        c.bench_function(&format!("bitmap aligned not 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let r = !&bitmap;
+                assert!(r.null_count() > 0);
+            })
+        });
+        let bitmap1 = bitmap.clone().slice(1, size - 1);
+        c.bench_function(&format!("bitmap not 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let r = !&bitmap1;
+                assert!(r.null_count() > 0);
+            })
+        });
+
+        let bitmap1: Bitmap = (0..size).into_iter().map(|x| x % 4 == 0).collect();
+        c.bench_function(&format!("bitmap aligned or 2^{}", log2_size), |b| {
+            b.iter(|| bench_arrow2(&bitmap, &bitmap1))
+        });
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

src/bitmap/bitmap_ops.rs

@@ -2,7 +2,10 @@ use std::ops::{BitAnd, BitOr, Not};
 
 use crate::buffer::MutableBuffer;
 
-use super::Bitmap;
+use super::{
+    utils::{BitChunkIterExact, BitChunksExact},
+    Bitmap,
+};
 
 /// Apply a bitwise operation `op` to four inputs and return the result as a [`Bitmap`].
 pub fn quaternary<F>(a1: &Bitmap, a2: &Bitmap, a3: &Bitmap, a4: &Bitmap, op: F) -> Bitmap
@@ -100,23 +103,33 @@ where
     Bitmap::from_u8_buffer(buffer, length)
 }
 
-/// Apply a bitwise operation `op` to one input and return the result as a [`Bitmap`].
-pub fn unary<F>(lhs: &Bitmap, op: F) -> Bitmap
+fn unary_impl<F, I>(iter: I, op: F, length: usize) -> Bitmap
 where
+    I: BitChunkIterExact<u64>,
     F: Fn(u64) -> u64,
 {
-    let mut lhs_chunks = lhs.chunks();
+    let rem = op(iter.remainder());
 
-    let chunks = lhs_chunks.by_ref().map(|left| op(left));
-    let mut buffer = unsafe { MutableBuffer::from_chunk_iter_unchecked(chunks) };
+    let iterator = iter.map(|left| op(left)).chain(std::iter::once(rem));
 
-    let remainder_bytes = lhs_chunks.remainder_len().saturating_add(7) / 8;
-    let rem = op(lhs_chunks.remainder());
+    let buffer = MutableBuffer::from_trusted_len_iter(iterator);
 
-    let rem = &rem.to_ne_bytes()[..remainder_bytes];
-    buffer.extend_from_slice(rem);
+    Bitmap::from_u8_buffer(buffer.into(), length)
+}
 
-    Bitmap::from_u8_buffer(buffer, lhs.len())
+/// Apply a bitwise operation `op` to one input and return the result as a [`Bitmap`].
+pub fn unary<F>(lhs: &Bitmap, op: F) -> Bitmap
+where
+    F: Fn(u64) -> u64,
+{
+    let (slice, offset, length) = lhs.as_slice();
+    if offset == 0 {
+        let iter = BitChunksExact::<u64>::new(slice, length);
+        unary_impl(iter, op, lhs.len())
+    } else {
+        let iter = lhs.chunks::<u64>();
+        unary_impl(iter, op, lhs.len())
+    }
 }
 
 fn and(lhs: &Bitmap, rhs: &Bitmap) -> Bitmap {

src/bitmap/utils/chunk_iterator/chunks_exact.rs

@@ -1,5 +1,7 @@
 use std::{convert::TryInto, slice::ChunksExact};
 
+use crate::trusted_len::TrustedLen;
+
 use super::{BitChunk, BitChunkIterExact};
 
 /// An iterator over a slice of bytes in [`BitChunk`]s.
@@ -84,6 +86,8 @@ impl<T: BitChunk> Iterator for BitChunksExact<'_, T> {
     }
 }
 
+unsafe impl<T: BitChunk> TrustedLen for BitChunksExact<'_, T> {}
+
 impl<T: BitChunk> BitChunkIterExact<T> for BitChunksExact<'_, T> {
     #[inline]
     fn remainder(&self) -> T {

src/bitmap/utils/chunk_iterator/mod.rs

@@ -10,7 +10,7 @@ use crate::{trusted_len::TrustedLen, types::BitChunkIter};
 pub(crate) use merge::merge_reversed;
 
 /// Trait representing an exact iterator over bytes in [`BitChunk`].
-pub trait BitChunkIterExact<B: BitChunk>: Iterator<Item = B> {
+pub trait BitChunkIterExact<B: BitChunk>: TrustedLen<Item = B> {
     /// The remainder of the iterator.
     fn remainder(&self) -> B;
 }