Skip to main content

mongreldb_core/index/
bitmap.rs

1//! Roaring-bitmap secondary index — `value bytes → row-id set`.
2//!
3//! Best for low-cardinality columns (equality, IN, GROUP BY). Multiple indexes
4//! intersect with cheap SIMD bitmap ops in the shared [`RowId`] space.
5
6use crate::rowid::RowId;
7use roaring::RoaringBitmap;
8use std::collections::HashMap;
9use std::sync::Arc;
10
11type BitmapLayer = HashMap<Vec<u8>, RoaringBitmap>;
12
13/// `value → row-id set`. Values are type-aware encoded bytes (lexicographically
14/// comparable), matching the encoding used for page min/max.
15#[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        // Roaring bitmaps address u32. The Phase-3 upgrade shards bitmaps by
37        // the high 32 bits to cover the full u64 row-id space; until then we
38        // require row ids < 2^32.
39        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    /// The row-id set for `value` (empty if absent).
45    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    /// Intersection of several sets — the workhorse of multi-condition queries.
56    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    /// All distinct values (keys) in this index — Phase 17.2 broadcast join.
74    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    /// Snapshot `(value_bytes → serialized RoaringBitmap)` pairs for
84    /// checkpointing to `_idx/global.idx`.
85    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    /// Rebuild from a snapshot produced by [`BitmapIndex::entries`].
99    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}