mongreldb_core/index/
sparse.rs

1//! Sparse inverted index for SPLADE-style learned-sparse retrieval.
2//!
3//! Each document is a sparse vector `(token → weight)`; the index is an inverted
4//! list `token → [(row_id, weight)]`. A query (also a sparse vector) scores
5//! documents by sparse dot product over shared tokens and returns the top-k.
6//! Real SPLADE produces these sparse vectors from a trained model; here any
7//! tokenizer works (the demo uses a hashing trick), so the retrieval machinery
8//! is model-agnostic — plug in real SPLADE weights as the sparse vectors.
9//!
10//! Like the other indexes, results resolve to the shared [`crate::rowid::RowId`]
11//! space, so `sparse_match ∩ fm_contains ∩ bitmap_eq` composes in one query.
12
13use crate::rowid::RowId;
14use std::collections::HashMap;
15
16/// Inverted index over weighted sparse vectors, keyed by token id.
17pub struct SparseIndex {
18    postings: HashMap<u32, Vec<(RowId, f32)>>,
19}
20
21impl SparseIndex {
22    pub fn new() -> Self {
23        Self {
24            postings: HashMap::new(),
25        }
26    }
27
28    /// Insert a document's sparse vector (`terms` need not be sorted; duplicate
29    /// tokens within one doc accumulate).
30    pub fn insert(&mut self, terms: &[(u32, f32)], row_id: RowId) {
31        for &(token, weight) in terms {
32            self.postings
33                .entry(token)
34                .or_default()
35                .push((row_id, weight));
36        }
37    }
38
39    /// Top-k row ids by sparse dot product with `query` (highest score first).
40    pub fn search(&self, query: &[(u32, f32)], k: usize) -> Vec<(RowId, f64)> {
41        self.search_filtered(query, k, |_| true)
42    }
43
44    pub fn search_filtered(
45        &self,
46        query: &[(u32, f32)],
47        k: usize,
48        allowed: impl Fn(RowId) -> bool,
49    ) -> Vec<(RowId, f64)> {
50        let mut scores: HashMap<u64, f64> = HashMap::new();
51        for &(token, q_weight) in query {
52            if let Some(list) = self.postings.get(&token) {
53                for &(rid, d_weight) in list {
54                    if allowed(rid) {
55                        *scores.entry(rid.0).or_insert(0.0) +=
56                            f64::from(q_weight) * f64::from(d_weight);
57                    }
58                }
59            }
60        }
61        let mut ranked: Vec<(RowId, f64)> = scores
62            .into_iter()
63            .map(|(rid, score)| (RowId(rid), score))
64            .collect();
65        ranked.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
66        ranked.truncate(k);
67        ranked
68    }
69
70    pub fn candidate_row_ids(&self, query: &[(u32, f32)]) -> Vec<RowId> {
71        let mut row_ids = std::collections::HashSet::new();
72        for (token, _) in query {
73            if let Some(postings) = self.postings.get(token) {
74                row_ids.extend(postings.iter().map(|(row_id, _)| *row_id));
75            }
76        }
77        row_ids.into_iter().collect()
78    }
79
80    pub fn is_empty(&self) -> bool {
81        self.postings.is_empty()
82    }
83
84    /// Snapshot the inverted lists for checkpointing to `_idx/global.idx`.
85    pub fn entries(&self) -> Vec<(u32, Vec<(RowId, f32)>)> {
86        self.postings
87            .iter()
88            .map(|(t, list)| (*t, list.clone()))
89            .collect()
90    }
91
92    /// Rebuild from a snapshot produced by [`SparseIndex::entries`].
93    pub fn from_entries(entries: Vec<(u32, Vec<(RowId, f32)>)>) -> Self {
94        let mut postings = HashMap::new();
95        for (t, list) in entries {
96            postings.insert(t, list);
97        }
98        Self { postings }
99    }
100}
101
102impl Default for SparseIndex {
103    fn default() -> Self {
104        Self::new()
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn ranks_by_sparse_overlap() {
114        let mut idx = SparseIndex::new();
115        // doc 0: {a:2, b:1}; doc 1: {a:1, c:3}; doc 2: {b:5}
116        idx.insert(&[(1, 2.0), (2, 1.0)], RowId(0));
117        idx.insert(&[(1, 1.0), (3, 3.0)], RowId(1));
118        idx.insert(&[(2, 5.0)], RowId(2));
119        // query {a:1, b:1}: doc0 = 2*1+1*1=3, doc1 = 1*1=1, doc2 = 5*1=5
120        let top = idx.search(&[(1, 1.0), (2, 1.0)], 3);
121        assert_eq!(top[0], (RowId(2), 5.0));
122        assert_eq!(top[1], (RowId(0), 3.0));
123        assert_eq!(top[2], (RowId(1), 1.0));
124    }
125}
mongreldb_core/index/sparse.rs

mongreldb_core/index/
sparse.rs
