diff --git a/src/common/benches/bitmap.rs b/src/common/benches/bitmap.rs index 943c394f93d48..64796d42944d4 100644 --- a/src/common/benches/bitmap.rs +++ b/src/common/benches/bitmap.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; +use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion}; use itertools::Itertools; use risingwave_common::buffer::{Bitmap, BitmapIter}; @@ -23,7 +23,27 @@ fn bench_bitmap(c: &mut Criterion) { let i = 0x123; c.bench_function("zeros", |b| b.iter(|| Bitmap::zeros(CHUNK_SIZE))); c.bench_function("ones", |b| b.iter(|| Bitmap::ones(CHUNK_SIZE))); - c.bench_function("get", |b| b.iter(|| x.is_set(i))); + c.bench_function("get", |b| { + b.iter(|| { + for _ in 0..1000 { + black_box(x.is_set(i)); + } + }) + }); + c.bench_function("get_1000", |b| { + b.iter(|| { + for _ in 0..1000 { + black_box(x.is_set(i)); + } + }) + }); + c.bench_function("get_1000_000", |b| { + b.iter(|| { + for _ in 0..1_000_000 { + black_box(x.is_set(i)); + } + }) + }); c.bench_function("and", |b| b.iter(|| &x & &y)); c.bench_function("or", |b| b.iter(|| &x | &y)); c.bench_function("not", |b| b.iter(|| !&x)); @@ -36,17 +56,16 @@ fn bench_bitmap_iter(c: &mut Criterion) { fn make_iterators(bitmaps: &[Bitmap]) -> Vec> { bitmaps.iter().map(|bitmap| bitmap.iter()).collect_vec() } - fn bench_bitmap_iter(bench_id: &str, bitmap: Bitmap, c: &mut Criterion) { + fn bench_bitmap_iter_inner(bench_id: &str, bitmap: Bitmap, c: &mut Criterion) { let bitmaps = vec![bitmap; N_CHUNKS]; let make_iters = || make_iterators(&bitmaps); - let mut result = vec![true; CHUNK_SIZE]; c.bench_function(bench_id, |b| { b.iter_batched( make_iters, |iters| { for iter in iters { - for (i, bit_flag) in iter.enumerate() { - result[i] = bit_flag; + for bit_flag in iter { + black_box(bit_flag); } } }, @@ -55,9 +74,9 @@ fn bench_bitmap_iter(c: &mut Criterion) { }); } let zeros = Bitmap::zeros(CHUNK_SIZE); - bench_bitmap_iter("zeros_iter", zeros, c); + bench_bitmap_iter_inner("zeros_iter", zeros, c); let ones = Bitmap::ones(CHUNK_SIZE); - bench_bitmap_iter("ones_iter", ones, c); + bench_bitmap_iter_inner("ones_iter", ones, c); } criterion_group!(benches, bench_bitmap, bench_bitmap_iter); diff --git a/src/common/src/buffer/bitmap.rs b/src/common/src/buffer/bitmap.rs index d75a150e2c31a..630baafef1729 100644 --- a/src/common/src/buffer/bitmap.rs +++ b/src/common/src/buffer/bitmap.rs @@ -173,13 +173,15 @@ impl BitmapBuilder { /// An immutable bitmap. Use [`BitmapBuilder`] to build it. #[derive(Clone, PartialEq, Eq)] pub struct Bitmap { - // The useful bits in the bitmap. The total number of bits will usually - // be larger than the useful bits due to byte-padding. + /// The useful bits in the bitmap. The total number of bits will usually + /// be larger than the useful bits due to byte-padding. num_bits: usize, - // The number of high bits in the bitmap. + /// The number of high bits in the bitmap. count_ones: usize, + /// Bits are stored in a compact form. + /// They are packed into `usize`s. bits: Box<[usize]>, } @@ -310,6 +312,7 @@ impl Bitmap { bits: &self.bits, idx: 0, num_bits: self.num_bits, + current_usize: 0, } } @@ -556,6 +559,27 @@ pub struct BitmapIter<'a> { bits: &'a [usize], idx: usize, num_bits: usize, + current_usize: usize, +} + +impl<'a> BitmapIter<'a> { + fn next_always_load_usize(&mut self) -> Option { + if self.idx >= self.num_bits { + return None; + } + + // Offset of the bit within the usize. + let usize_offset = self.idx % BITS; + + // Get the index of usize which the bit is located in + let usize_index = self.idx / BITS; + self.current_usize = unsafe { *self.bits.get_unchecked(usize_index) }; + + let bit_mask = 1 << usize_offset; + let bit_flag = self.current_usize & bit_mask != 0; + self.idx += 1; + Some(bit_flag) + } } impl<'a> iter::Iterator for BitmapIter<'a> { @@ -565,9 +589,21 @@ impl<'a> iter::Iterator for BitmapIter<'a> { if self.idx >= self.num_bits { return None; } - let b = unsafe { self.bits.get_unchecked(self.idx / BITS) } & (1 << (self.idx % BITS)) != 0; + + // Offset of the bit within the usize. + let usize_offset = self.idx % BITS; + + if usize_offset == 0 { + // Get the index of usize which the bit is located in + let usize_index = self.idx / BITS; + self.current_usize = unsafe { *self.bits.get_unchecked(usize_index) }; + } + + let bit_mask = 1 << usize_offset; + + let bit_flag = self.current_usize & bit_mask != 0; self.idx += 1; - Some(b) + Some(bit_flag) } fn size_hint(&self) -> (usize, Option) { @@ -577,7 +613,7 @@ impl<'a> iter::Iterator for BitmapIter<'a> { fn nth(&mut self, n: usize) -> Option { self.idx += n; - self.next() + self.next_always_load_usize() } }