mongreldb_core/index/
bitmap.rs1use crate::rowid::RowId;
7use roaring::RoaringBitmap;
8use std::collections::HashMap;
9use std::sync::Arc;
10
11type BitmapLayer = HashMap<Vec<u8>, RoaringBitmap>;
12
13#[derive(Clone)]
16pub struct BitmapIndex {
17 frozen: Arc<Vec<Arc<BitmapLayer>>>,
18 active: BitmapLayer,
19}
20
21impl Default for BitmapIndex {
22 fn default() -> Self {
23 Self::new()
24 }
25}
26
27impl BitmapIndex {
28 pub fn new() -> Self {
29 Self {
30 frozen: Arc::new(Vec::new()),
31 active: HashMap::new(),
32 }
33 }
34
35 pub fn insert(&mut self, value: Vec<u8>, row_id: RowId) {
36 let id32 = u32::try_from(row_id.0)
40 .expect("bitmap index supports row_id < 2^32; shard-by-high-bits is a Phase-3 upgrade");
41 self.active.entry(value).or_default().insert(id32);
42 }
43
44 pub fn get(&self, value: &[u8]) -> RoaringBitmap {
46 let mut rows = self.active.get(value).cloned().unwrap_or_default();
47 for layer in self.frozen.iter() {
48 if let Some(layer_rows) = layer.get(value) {
49 rows |= layer_rows;
50 }
51 }
52 rows
53 }
54
55 pub fn intersect(sets: &[RoaringBitmap]) -> RoaringBitmap {
57 match sets {
58 [] => RoaringBitmap::new(),
59 [first, rest @ ..] => {
60 let mut acc = first.clone();
61 for s in rest {
62 acc &= s;
63 }
64 acc
65 }
66 }
67 }
68
69 pub fn value_count(&self) -> usize {
70 self.keys().len()
71 }
72
73 pub fn keys(&self) -> Vec<Vec<u8>> {
75 let mut keys = std::collections::HashSet::new();
76 keys.extend(self.active.keys().cloned());
77 for layer in self.frozen.iter() {
78 keys.extend(layer.keys().cloned());
79 }
80 keys.into_iter().collect()
81 }
82
83 pub fn entries(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
86 self.keys()
87 .into_iter()
88 .map(|k| {
89 let v = self.get(&k);
90 let mut bytes = Vec::new();
91 v.serialize_into(&mut bytes)
92 .expect("roaring serialize is infallible for Vec");
93 (k, bytes)
94 })
95 .collect()
96 }
97
98 pub fn from_entries(
100 entries: Vec<(Vec<u8>, Vec<u8>)>,
101 ) -> std::result::Result<Self, &'static str> {
102 let mut active = HashMap::new();
103 for (k, bytes) in entries {
104 let bm = RoaringBitmap::deserialize_from(&bytes[..]).map_err(|_| "bad bitmap bytes")?;
105 active.insert(k, bm);
106 }
107 Ok(Self {
108 frozen: Arc::new(Vec::new()),
109 active,
110 })
111 }
112
113 pub(crate) fn seal(&mut self) {
114 if self.active.is_empty() {
115 return;
116 }
117 let active = std::mem::take(&mut self.active);
118 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
119 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
120 self.consolidate();
121 }
122 }
123
124 fn consolidate(&mut self) {
125 let mut merged = HashMap::<Vec<u8>, RoaringBitmap>::new();
126 for layer in self.frozen.iter() {
127 for (key, rows) in layer.iter() {
128 *merged.entry(key.clone()).or_default() |= rows;
129 }
130 }
131 self.frozen = Arc::new(vec![Arc::new(merged)]);
132 }
133
134 #[cfg(test)]
135 pub(crate) fn frozen_layer_count(&self) -> usize {
136 self.frozen.len()
137 }
138}
139
140#[cfg(test)]
141mod tests {
142 use super::*;
143
144 #[test]
145 fn insert_get_and_intersect() {
146 let mut color = BitmapIndex::new();
147 color.insert(b"red".to_vec(), RowId(1));
148 color.insert(b"red".to_vec(), RowId(3));
149 color.insert(b"blue".to_vec(), RowId(3));
150
151 let mut region = BitmapIndex::new();
152 region.insert(b"us".to_vec(), RowId(1));
153 region.insert(b"us".to_vec(), RowId(3));
154 region.insert(b"eu".to_vec(), RowId(2));
155
156 let red = color.get(b"red");
157 let us = region.get(b"us");
158 let both = BitmapIndex::intersect(&[red, us]);
159 let ids: Vec<u32> = both.iter().collect();
160 assert_eq!(ids, vec![1, 3]);
161 }
162
163 #[test]
164 fn sealed_generations_share_bitmaps_and_consolidate() {
165 let mut writer = BitmapIndex::new();
166 for id in 0..crate::MAX_READ_GENERATION_LAYERS as u64 + 2 {
167 writer.insert(b"all".to_vec(), RowId(id));
168 writer.seal();
169 }
170 assert!(writer.frozen_layer_count() < crate::MAX_READ_GENERATION_LAYERS);
171 let generation = writer.clone();
172 writer.insert(b"new".to_vec(), RowId(99));
173 assert!(generation.get(b"new").is_empty());
174 assert!(writer.get(b"new").contains(99));
175 }
176}