move code to DenseIndexBlock

improve benchmark
This commit is contained in:
Pascal Seitz
2022-12-08 19:46:23 +08:00
parent a36b50d825
commit 789d29cf45
2 changed files with 108 additions and 60 deletions

View File

@@ -1,3 +1,4 @@
use std::convert::TryInto;
use std::io::{self, Write};
use common::BinarySerializable;
@@ -30,13 +31,28 @@ const BLOCK_BITVEC_SIZE: usize = 8;
const BLOCK_OFFSET_SIZE: usize = 4;
const SERIALIZED_BLOCK_SIZE: usize = BLOCK_BITVEC_SIZE + BLOCK_OFFSET_SIZE;
fn count_ones(block: u64, pos_in_block: u32) -> u32 {
if pos_in_block == 63 {
block.count_ones()
#[inline]
fn count_ones(bitvec: u64, pos_in_bitvec: u32) -> u32 {
if pos_in_bitvec == 63 {
bitvec.count_ones()
} else {
let mask = (1u64 << (pos_in_block + 1)) - 1;
let masked_block = block & mask;
masked_block.count_ones()
let mask = (1u64 << (pos_in_bitvec + 1)) - 1;
let masked_bitvec = bitvec & mask;
masked_bitvec.count_ones()
}
}
#[derive(Clone, Copy)]
pub(crate) struct DenseIndexBlock {
bitvec: u64,
offset: u32,
}
impl From<[u8; SERIALIZED_BLOCK_SIZE]> for DenseIndexBlock {
fn from(data: [u8; SERIALIZED_BLOCK_SIZE]) -> Self {
let bitvec = u64::from_le_bytes(data[..BLOCK_BITVEC_SIZE].try_into().unwrap());
let offset = u32::from_le_bytes(data[BLOCK_BITVEC_SIZE..].try_into().unwrap());
Self { bitvec, offset }
}
}
@@ -49,51 +65,33 @@ impl DenseCodec {
/// Check if value at position is not null.
pub fn exists(&self, idx: u32) -> bool {
let block_pos = idx / ELEMENTS_PER_BLOCK;
let bitvec: u64 = self.block(block_pos as usize);
let bitvec = self.dense_index_block(block_pos).bitvec;
let pos_in_block = idx % ELEMENTS_PER_BLOCK;
let pos_in_bitvec = idx % ELEMENTS_PER_BLOCK;
get_bit_at(bitvec, pos_in_block)
get_bit_at(bitvec, pos_in_bitvec)
}
#[inline]
pub(crate) fn block(&self, block_pos: usize) -> u64 {
let data = &mut &self.data[block_pos as usize * SERIALIZED_BLOCK_SIZE..];
let block: u64 =
BinarySerializable::deserialize(data).expect("could not read block in null index");
block
}
#[inline]
/// Returns (bitvec, offset)
///
/// offset is the start offset of actual docids in the block.
pub(crate) fn block_and_offset(&self, block_pos: u32) -> (u64, u32) {
let data = &mut &self.data[block_pos as usize * SERIALIZED_BLOCK_SIZE..];
let block: u64 =
BinarySerializable::deserialize(data).expect("could not read block in null index");
let offset: u32 =
BinarySerializable::deserialize(data).expect("could not read block in null index");
(block, offset)
pub(crate) fn dense_index_block(&self, block_pos: u32) -> DenseIndexBlock {
dense_index_block(&self.data, block_pos)
}
/// Return the number of non-null values in an index
pub fn num_non_null_vals(&self) -> u32 {
let last_block = (self.data.len() / SERIALIZED_BLOCK_SIZE) - 1;
self.block_and_offset(last_block as u32).1
self.dense_index_block(last_block as u32).offset
}
#[inline]
/// Translate from the original index to the codec index.
pub fn translate_to_codec_idx(&self, idx: u32) -> Option<u32> {
let block_pos = idx / ELEMENTS_PER_BLOCK;
let (block, offset) = self.block_and_offset(block_pos);
let pos_in_block = idx % ELEMENTS_PER_BLOCK;
let ones_in_block = count_ones(block, pos_in_block);
if get_bit_at(block, pos_in_block) {
Some(offset + ones_in_block - 1) // -1 is ok, since idx does exist, so there's at least
// one
let index_block = self.dense_index_block(block_pos);
let pos_in_block_bit_vec = idx % ELEMENTS_PER_BLOCK;
let ones_in_block = count_ones(index_block.bitvec, pos_in_block_bit_vec);
if get_bit_at(index_block.bitvec, pos_in_block_bit_vec) {
// -1 is ok, since idx does exist, so there's at least one
Some(index_block.offset + ones_in_block - 1)
} else {
None
}
@@ -112,17 +110,17 @@ impl DenseCodec {
iter.map(move |dense_idx| {
// update block_pos to limit search scope
block_pos = find_block(dense_idx, block_pos, &self.data);
let (bitvec, offset) = self.block_and_offset(block_pos);
let index_block = self.dense_index_block(block_pos);
// The next offset is higher than dense_idx and therefore:
// dense_idx <= offset + num_set_bits in block
let mut num_set_bits = 0;
for idx_in_block in 0..ELEMENTS_PER_BLOCK {
if get_bit_at(bitvec, idx_in_block) {
for idx_in_bitvec in 0..ELEMENTS_PER_BLOCK {
if get_bit_at(index_block.bitvec, idx_in_bitvec) {
num_set_bits += 1;
}
if num_set_bits == (dense_idx - offset + 1) {
let orig_idx = block_pos * ELEMENTS_PER_BLOCK + idx_in_block as u32;
if num_set_bits == (dense_idx - index_block.offset + 1) {
let orig_idx = block_pos * ELEMENTS_PER_BLOCK + idx_in_bitvec as u32;
return orig_idx;
}
}
@@ -131,6 +129,17 @@ impl DenseCodec {
}
}
#[inline]
pub(crate) fn dense_index_block(data: &[u8], block_pos: u32) -> DenseIndexBlock {
let data_start_pos = block_pos as usize * SERIALIZED_BLOCK_SIZE;
let block_data: [u8; SERIALIZED_BLOCK_SIZE] = data
[data_start_pos..data_start_pos + SERIALIZED_BLOCK_SIZE]
.try_into()
.unwrap();
let block = block_data.into();
block
}
#[inline]
/// Finds the block position containing the dense_idx.
///
@@ -140,11 +149,8 @@ impl DenseCodec {
/// The last offset number is equal to the number of values in the index.
fn find_block(dense_idx: u32, mut block_pos: u32, data: &[u8]) -> u32 {
loop {
let data = &mut &data[BLOCK_BITVEC_SIZE + block_pos as usize * SERIALIZED_BLOCK_SIZE..];
let offset: u32 = BinarySerializable::deserialize(data)
.expect("could not read offset from block in null index");
let offset = dense_index_block(&data, block_pos).offset;
if offset > dense_idx {
// offset
return block_pos - 1;
}
block_pos += 1;
@@ -330,36 +336,76 @@ mod bench {
use super::*;
fn gen_bools() -> DenseCodec {
const NUM_VALUES: u32 = 1_000_000;
fn gen_bools(fill_ratio: f64) -> DenseCodec {
let mut out = Vec::new();
let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
// 80% of values are set
let bools: Vec<_> = (0..100_000).map(|_| rng.gen_bool(8f64 / 10f64)).collect();
let bools: Vec<_> = (0..NUM_VALUES).map(|_| rng.gen_bool(fill_ratio)).collect();
serialize_dense_codec(bools.into_iter(), &mut out).unwrap();
let codec = DenseCodec::open(OwnedBytes::new(out));
codec
}
#[bench]
fn bench_dense_codec_translate_orig_to_dense(bench: &mut Bencher) {
let codec = gen_bools();
bench.iter(|| {
let mut dense_idx: Option<u32> = None;
for idx in 0..100_000 {
dense_idx = dense_idx.or(codec.translate_to_codec_idx(idx));
fn random_range_iterator(start: u32, end: u32, step_size: u32) -> impl Iterator<Item = u32> {
let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
let mut current = start;
std::iter::from_fn(move || {
current += rng.gen_range(1..step_size + 1);
if current >= end {
None
} else {
Some(current)
}
dense_idx
});
})
}
fn walk_over_data(codec: &DenseCodec, max_step_size: u32) -> Option<u32> {
walk_over_data_from_positions(codec, random_range_iterator(0, NUM_VALUES, max_step_size))
}
fn walk_over_data_from_positions(
codec: &DenseCodec,
positions: impl Iterator<Item = u32>,
) -> Option<u32> {
let mut dense_idx: Option<u32> = None;
for idx in positions {
dense_idx = dense_idx.or(codec.translate_to_codec_idx(idx));
}
dense_idx
}
#[bench]
fn bench_dense_codec_translate_dense_to_orig(bench: &mut Bencher) {
let codec = gen_bools();
fn bench_dense_codec_translate_orig_to_dense_80percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.8f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_dense_codec_translate_orig_to_dense_10percent_full_scan(bench: &mut Bencher) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..NUM_VALUES));
}
#[bench]
fn bench_dense_codec_translate_orig_to_dense_10percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_dense_codec_translate_dense_to_orig_80percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.8f64);
let num_vals = codec.num_non_null_vals();
bench.iter(|| {
codec
.translate_codec_idx_to_original_idx(0..num_vals)
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100))
.last()
});
}

View File

@@ -2,10 +2,12 @@ pub use dense::{serialize_dense_codec, DenseCodec};
mod dense;
#[inline]
fn get_bit_at(input: u64, n: u32) -> bool {
input & (1 << n) != 0
}
#[inline]
fn set_bit_at(input: &mut u64, n: u64) {
*input |= 1 << n;
}