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

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;
 }