mongreldb_core/index/
bitmap.rs1use crate::rowid::RowId;
7use roaring::RoaringBitmap;
8
9pub struct BitmapIndex {
12 map: std::collections::HashMap<Vec<u8>, RoaringBitmap>,
13}
14
15impl Default for BitmapIndex {
16 fn default() -> Self {
17 Self::new()
18 }
19}
20
21impl BitmapIndex {
22 pub fn new() -> Self {
23 Self {
24 map: std::collections::HashMap::new(),
25 }
26 }
27
28 pub fn insert(&mut self, value: Vec<u8>, row_id: RowId) {
29 let id32 = u32::try_from(row_id.0)
33 .expect("bitmap index supports row_id < 2^32; shard-by-high-bits is a Phase-3 upgrade");
34 self.map.entry(value).or_default().insert(id32);
35 }
36
37 pub fn get(&self, value: &[u8]) -> RoaringBitmap {
39 self.map.get(value).cloned().unwrap_or_default()
40 }
41
42 pub fn intersect(sets: &[RoaringBitmap]) -> RoaringBitmap {
44 match sets {
45 [] => RoaringBitmap::new(),
46 [first, rest @ ..] => {
47 let mut acc = first.clone();
48 for s in rest {
49 acc &= s;
50 }
51 acc
52 }
53 }
54 }
55
56 pub fn value_count(&self) -> usize {
57 self.map.len()
58 }
59
60 pub fn keys(&self) -> Vec<&Vec<u8>> {
62 self.map.keys().collect()
63 }
64
65 pub fn entries(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
68 self.map
69 .iter()
70 .map(|(k, v)| {
71 let mut bytes = Vec::new();
72 v.serialize_into(&mut bytes)
73 .expect("roaring serialize is infallible for Vec");
74 (k.clone(), bytes)
75 })
76 .collect()
77 }
78
79 pub fn from_entries(
81 entries: Vec<(Vec<u8>, Vec<u8>)>,
82 ) -> std::result::Result<Self, &'static str> {
83 let mut map = std::collections::HashMap::new();
84 for (k, bytes) in entries {
85 let bm = RoaringBitmap::deserialize_from(&bytes[..]).map_err(|_| "bad bitmap bytes")?;
86 map.insert(k, bm);
87 }
88 Ok(Self { map })
89 }
90}
91
92#[cfg(test)]
93mod tests {
94 use super::*;
95
96 #[test]
97 fn insert_get_and_intersect() {
98 let mut color = BitmapIndex::new();
99 color.insert(b"red".to_vec(), RowId(1));
100 color.insert(b"red".to_vec(), RowId(3));
101 color.insert(b"blue".to_vec(), RowId(3));
102
103 let mut region = BitmapIndex::new();
104 region.insert(b"us".to_vec(), RowId(1));
105 region.insert(b"us".to_vec(), RowId(3));
106 region.insert(b"eu".to_vec(), RowId(2));
107
108 let red = color.get(b"red");
109 let us = region.get(b"us");
110 let both = BitmapIndex::intersect(&[red, us]);
111 let ids: Vec<u32> = both.iter().collect();
112 assert_eq!(ids, vec![1, 3]);
113 }
114}