use crate::rowid::RowId;
use roaring::RoaringBitmap;
pub struct BitmapIndex {
map: std::collections::HashMap<Vec<u8>, RoaringBitmap>,
}
impl Default for BitmapIndex {
fn default() -> Self {
Self::new()
}
}
impl BitmapIndex {
pub fn new() -> Self {
Self {
map: std::collections::HashMap::new(),
}
}
pub fn insert(&mut self, value: Vec<u8>, row_id: RowId) {
let id32 = u32::try_from(row_id.0)
.expect("bitmap index supports row_id < 2^32; shard-by-high-bits is a Phase-3 upgrade");
self.map.entry(value).or_default().insert(id32);
}
pub fn get(&self, value: &[u8]) -> RoaringBitmap {
self.map.get(value).cloned().unwrap_or_default()
}
pub fn intersect(sets: &[RoaringBitmap]) -> RoaringBitmap {
match sets {
[] => RoaringBitmap::new(),
[first, rest @ ..] => {
let mut acc = first.clone();
for s in rest {
acc &= s;
}
acc
}
}
}
pub fn value_count(&self) -> usize {
self.map.len()
}
pub fn keys(&self) -> Vec<&Vec<u8>> {
self.map.keys().collect()
}
pub fn entries(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
self.map
.iter()
.map(|(k, v)| {
let mut bytes = Vec::new();
v.serialize_into(&mut bytes)
.expect("roaring serialize is infallible for Vec");
(k.clone(), bytes)
})
.collect()
}
pub fn from_entries(
entries: Vec<(Vec<u8>, Vec<u8>)>,
) -> std::result::Result<Self, &'static str> {
let mut map = std::collections::HashMap::new();
for (k, bytes) in entries {
let bm = RoaringBitmap::deserialize_from(&bytes[..]).map_err(|_| "bad bitmap bytes")?;
map.insert(k, bm);
}
Ok(Self { map })
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn insert_get_and_intersect() {
let mut color = BitmapIndex::new();
color.insert(b"red".to_vec(), RowId(1));
color.insert(b"red".to_vec(), RowId(3));
color.insert(b"blue".to_vec(), RowId(3));
let mut region = BitmapIndex::new();
region.insert(b"us".to_vec(), RowId(1));
region.insert(b"us".to_vec(), RowId(3));
region.insert(b"eu".to_vec(), RowId(2));
let red = color.get(b"red");
let us = region.get(b"us");
let both = BitmapIndex::intersect(&[red, us]);
let ids: Vec<u32> = both.iter().collect();
assert_eq!(ids, vec![1, 3]);
}
}