mongreldb-core 0.32.0

MongrelDB core: log-structured columnar store with sub-ms writes, learned indexes, and an AI-native access layer.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

//! Sparse inverted index for SPLADE-style learned-sparse retrieval.
//!
//! Each document is a sparse vector `(token → weight)`; the index is an inverted
//! list `token → [(row_id, weight)]`. A query (also a sparse vector) scores
//! documents by sparse dot product over shared tokens and returns the top-k.
//! Real SPLADE produces these sparse vectors from a trained model; here any
//! tokenizer works (the demo uses a hashing trick), so the retrieval machinery
//! is model-agnostic — plug in real SPLADE weights as the sparse vectors.
//!
//! Like the other indexes, results resolve to the shared [`crate::rowid::RowId`]
//! space, so `sparse_match ∩ fm_contains ∩ bitmap_eq` composes in one query.

use crate::rowid::RowId;
use std::collections::HashMap;

/// Inverted index over weighted sparse vectors, keyed by token id.
pub struct SparseIndex {
    postings: HashMap<u32, Vec<(RowId, f32)>>,
}

impl SparseIndex {
    pub fn new() -> Self {
        Self {
            postings: HashMap::new(),
        }
    }

    /// Insert a document's sparse vector (`terms` need not be sorted; duplicate
    /// tokens within one doc accumulate).
    pub fn insert(&mut self, terms: &[(u32, f32)], row_id: RowId) {
        for &(token, weight) in terms {
            self.postings
                .entry(token)
                .or_default()
                .push((row_id, weight));
        }
    }

    /// Top-k row ids by sparse dot product with `query` (highest score first).
    pub fn search(&self, query: &[(u32, f32)], k: usize) -> Vec<(RowId, f32)> {
        let mut scores: HashMap<u64, f32> = HashMap::new();
        for &(token, q_weight) in query {
            if let Some(list) = self.postings.get(&token) {
                for &(rid, d_weight) in list {
                    *scores.entry(rid.0).or_insert(0.0) += q_weight * d_weight;
                }
            }
        }
        let mut ranked: Vec<(RowId, f32)> = scores
            .into_iter()
            .map(|(rid, score)| (RowId(rid), score))
            .collect();
        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        ranked.truncate(k);
        ranked
    }

    pub fn is_empty(&self) -> bool {
        self.postings.is_empty()
    }

    /// Snapshot the inverted lists for checkpointing to `_idx/global.idx`.
    pub fn entries(&self) -> Vec<(u32, Vec<(RowId, f32)>)> {
        self.postings
            .iter()
            .map(|(t, list)| (*t, list.clone()))
            .collect()
    }

    /// Rebuild from a snapshot produced by [`SparseIndex::entries`].
    pub fn from_entries(entries: Vec<(u32, Vec<(RowId, f32)>)>) -> Self {
        let mut postings = HashMap::new();
        for (t, list) in entries {
            postings.insert(t, list);
        }
        Self { postings }
    }
}

impl Default for SparseIndex {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn ranks_by_sparse_overlap() {
        let mut idx = SparseIndex::new();
        // doc 0: {a:2, b:1}; doc 1: {a:1, c:3}; doc 2: {b:5}
        idx.insert(&[(1, 2.0), (2, 1.0)], RowId(0));
        idx.insert(&[(1, 1.0), (3, 3.0)], RowId(1));
        idx.insert(&[(2, 5.0)], RowId(2));
        // query {a:1, b:1}: doc0 = 2*1+1*1=3, doc1 = 1*1=1, doc2 = 5*1=5
        let top = idx.search(&[(1, 1.0), (2, 1.0)], 3);
        assert_eq!(top[0], (RowId(2), 5.0));
        assert_eq!(top[1], (RowId(0), 3.0));
        assert_eq!(top[2], (RowId(1), 1.0));
    }
}