Improved performance of SlicesIterator (15%) (#824)

Cargo.toml

@@ -306,3 +306,7 @@ harness = false
 [[bench]]
 name = "write_json"
 harness = false
+
+[[bench]]
+name = "slices_iterator"
+harness = false

benches/slices_iterator.rs

@@ -0,0 +1,25 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+
+use arrow2::bitmap::{utils::SlicesIterator, Bitmap};
+
+fn bench_slices(lhs: &Bitmap) {
+    let set_count = lhs.len() - lhs.null_count();
+    let slices = SlicesIterator::new(lhs);
+
+    let count = slices.fold(0usize, |acc, v| acc + v.1);
+    assert_eq!(count, set_count);
+}
+
+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 slices 2^{}", log2_size), |b| {
+            b.iter(|| bench_slices(&bitmap))
+        });
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

src/bitmap/utils/slice_iterator.rs

@@ -50,6 +50,7 @@ impl<'a> SlicesIterator<'a> {
         }
     }
 
+    #[inline]
     fn finish(&mut self) -> Option<(usize, usize)> {
         self.state = State::Finished;
         if self.on_region {
@@ -75,6 +76,7 @@ impl<'a> SlicesIterator<'a> {
 impl<'a> Iterator for SlicesIterator<'a> {
     type Item = (usize, usize);
 
+    #[inline]
     fn next(&mut self) -> Option<Self::Item> {
         loop {
             if self.state == State::Finished {
@@ -89,17 +91,15 @@ impl<'a> Iterator for SlicesIterator<'a> {
                 match (self.on_region, self.current_byte) {
                     (true, &255u8) => {
                         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(),
+                        if let Some(v) = self.values.next() {
+                            self.current_byte = v;
                         };
                         continue;
                     }
                     (false, &0) => {
                         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(),
+                        if let Some(v) = self.values.next() {
+                            self.current_byte = v;
                         };
                         continue;
                     }
@@ -114,26 +114,17 @@ impl<'a> Iterator for SlicesIterator<'a> {
                 (true, true) => self.len += 1,
                 (false, false) => self.len += 1,
                 (true, false) => {
+                    self.on_region = false;
+                    let result = (self.start, self.len);
+                    self.start += self.len;
+                    self.len = 1;
                     if self.mask == 1 {
                         // reached a new byte => try to fetch it from the iterator
-                        match self.values.next() {
-                            Some(v) => {
-                                self.on_region = false;
-                                let result = (self.start, self.len);
-                                self.start += self.len;
-                                self.len = 1;
-                                self.current_byte = v;
-                                return Some(result);
-                            }
-                            None => return self.finish(),
+                        if let Some(v) = self.values.next() {
+                            self.current_byte = v;
                         };
-                    } else {
-                        self.on_region = false;
-                        let result = (self.start, self.len);
-                        self.start += self.len;
-                        self.len = 1;
-                        return Some(result);
                     }
+                    return Some(result);
                 }
                 (false, true) => {
                     self.start += self.len;