nodedb_vector/multivec/
meta_embed.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! MetaEmbed-specific search helpers.
4//!
5//! MetaEmbed (ICLR 2026) stores K pre-computed Meta Token vectors per document
6//! instead of one vector per token.  At query time the Matryoshka ordering of
7//! the query Meta Tokens allows a `budget` parameter to trade recall for
8//! latency: `budget=k` is full accuracy; `budget=1` is fastest.
9//!
10//! `meta_embed_search` wires together:
11//! 1. Optional PLAID candidate pruning.
12//! 2. Budgeted MaxSim scoring over the remaining candidates.
13//! 3. Top-k selection.
14
15use nodedb_types::vector_distance::DistanceMetric;
16
17use super::plaid::PlaidPruner;
18use super::scoring::budgeted_maxsim;
19use super::storage::MultiVectorStore;
20use crate::error::{VectorError, check_dim};
21
22/// Search a `MultiVectorStore` using budgeted MaxSim with optional PLAID
23/// candidate pruning.
24///
25/// # Parameters
26/// * `store`  — the document collection.
27/// * `plaid`  — optional PLAID pruner (pass `None` to scan all docs).
28/// * `query`  — query Meta Token vectors (Matryoshka ordering).
29/// * `budget` — number of leading query tokens to use; 0 falls back to all.
30/// * `k`      — number of top documents to return.
31/// * `metric` — distance metric (Cosine recommended for MetaEmbed).
32///
33/// # Returns
34/// A `Vec<(doc_id, score)>` sorted descending by score, length ≤ `k`, or
35/// [`VectorError::DimensionMismatch`] when a query vector does not have the
36/// store dimension.
37pub fn meta_embed_search(
38    store: &MultiVectorStore,
39    plaid: Option<&PlaidPruner>,
40    query: &[Vec<f32>],
41    budget: u8,
42    k: usize,
43    metric: DistanceMetric,
44) -> Result<Vec<(u32, f32)>, VectorError> {
45    for v in query {
46        check_dim(store.dim, v.len())?;
47    }
48    if k == 0 || query.is_empty() {
49        return Ok(Vec::new());
50    }
51
52    // Effective budget: 0 means use all query vectors.
53    let effective_budget = if budget == 0 {
54        query.len() as u8
55    } else {
56        budget
57    };
58
59    // Determine candidate set.
60    let candidate_ids: Vec<u32> = match plaid {
61        Some(pruner) => pruner.candidates(query)?,
62        None => store.iter().map(|doc| doc.doc_id).collect(),
63    };
64
65    // Score each candidate.
66    let mut scored: Vec<(u32, f32)> = candidate_ids
67        .into_iter()
68        .filter_map(|doc_id| {
69            store.get(doc_id).map(|doc| {
70                let score = budgeted_maxsim(query, &doc.vectors, effective_budget, metric);
71                (doc_id, score)
72            })
73        })
74        .collect();
75
76    // Sort descending by score.
77    scored.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
78    scored.truncate(k);
79    Ok(scored)
80}
81
82// ---------------------------------------------------------------------------
83// Tests
84// ---------------------------------------------------------------------------
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::multivec::storage::{MultiVecMode, MultiVectorDoc, MultiVectorStore};
90
91    /// Build a store with `n` documents, each with one unit-vector at
92    /// dimension `i % dim`.
93    fn build_store(n: u32, dim: usize, k: u8) -> MultiVectorStore {
94        let mut store = MultiVectorStore::new(dim, MultiVecMode::MetaToken { k });
95        for i in 0..n {
96            let mut vecs: Vec<Vec<f32>> = Vec::new();
97            for j in 0..k as usize {
98                let mut v = vec![0.0f32; dim];
99                // Each Meta Token of doc i points in direction (i + j) % dim.
100                v[(i as usize + j) % dim] = 1.0;
101                vecs.push(v);
102            }
103            store
104                .insert(MultiVectorDoc {
105                    doc_id: i,
106                    vectors: vecs,
107                })
108                .unwrap();
109        }
110        store
111    }
112
113    #[test]
114    fn search_returns_at_most_k_results() {
115        let store = build_store(10, 4, 2);
116        let query = vec![vec![1.0f32, 0.0, 0.0, 0.0]];
117        let results =
118            meta_embed_search(&store, None, &query, 2, 3, DistanceMetric::Cosine).unwrap();
119        assert!(results.len() <= 3);
120    }
121
122    #[test]
123    fn search_results_sorted_descending() {
124        let store = build_store(8, 4, 2);
125        let query = vec![vec![1.0f32, 0.0, 0.0, 0.0]];
126        let results =
127            meta_embed_search(&store, None, &query, 2, 8, DistanceMetric::Cosine).unwrap();
128        for w in results.windows(2) {
129            assert!(w[0].1 >= w[1].1, "not sorted: {:?}", results);
130        }
131    }
132
133    #[test]
134    fn plaid_filtered_results_are_subset_of_unfiltered() {
135        let store = build_store(9, 2, 2);
136
137        // Train PLAID on the store.
138        let pruner = PlaidPruner::train(&store, 3, 10, 99);
139
140        let query = vec![vec![1.0f32, 0.0f32]];
141        let unfiltered =
142            meta_embed_search(&store, None, &query, 2, 9, DistanceMetric::Cosine).unwrap();
143        let filtered =
144            meta_embed_search(&store, Some(&pruner), &query, 2, 9, DistanceMetric::Cosine).unwrap();
145
146        let unfiltered_ids: std::collections::HashSet<u32> =
147            unfiltered.iter().map(|(id, _)| *id).collect();
148
149        for (id, _) in &filtered {
150            assert!(
151                unfiltered_ids.contains(id),
152                "filtered result {id} not in unfiltered set"
153            );
154        }
155    }
156
157    #[test]
158    fn search_empty_query_returns_empty() {
159        let store = build_store(5, 4, 2);
160        let results = meta_embed_search(&store, None, &[], 2, 5, DistanceMetric::Cosine).unwrap();
161        assert!(results.is_empty());
162    }
163
164    #[test]
165    fn search_k_zero_returns_empty() {
166        let store = build_store(5, 4, 2);
167        let query = vec![vec![1.0f32, 0.0, 0.0, 0.0]];
168        let results =
169            meta_embed_search(&store, None, &query, 2, 0, DistanceMetric::Cosine).unwrap();
170        assert!(results.is_empty());
171    }
172
173    #[test]
174    fn top_result_is_best_matching_doc() {
175        // Doc 0 has its first Meta Token at direction 0 (index 0).
176        // Query is also in direction 0 — doc 0 should rank first.
177        let store = build_store(4, 4, 1);
178        let query = vec![vec![1.0f32, 0.0, 0.0, 0.0]];
179        let results =
180            meta_embed_search(&store, None, &query, 1, 1, DistanceMetric::Cosine).unwrap();
181        assert_eq!(results.len(), 1);
182        assert_eq!(results[0].0, 0, "expected doc_id=0 to be top result");
183    }
184
185    #[test]
186    fn wrong_dimension_query_is_a_typed_error() {
187        let store = build_store(4, 4, 2);
188        let pruner = PlaidPruner::train(&store, 2, 3, 7);
189        let query = vec![vec![1.0f32, 0.0, 0.0]];
190        for plaid in [None, Some(&pruner)] {
191            let result = meta_embed_search(&store, plaid, &query, 1, 2, DistanceMetric::Cosine);
192            assert!(
193                matches!(
194                    result,
195                    Err(VectorError::DimensionMismatch {
196                        expected: 4,
197                        got: 3
198                    })
199                ),
200                "{result:?}"
201            );
202        }
203    }
204}
nodedb_vector/multivec/meta_embed.rs

nodedb_vector/multivec/
meta_embed.rs
