gonzalo_vector/
index.rs

1//! In-memory exact cosine vector index keyed by [`RecordKey`].
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5
6use async_trait::async_trait;
7use gonzalo_core::{CoreError, KeyPrefix, RecordKey, Result};
8
9use crate::{Match, VectorIndex};
10
11// ---------------------------------------------------------------------------
12// Cosine helper
13// ---------------------------------------------------------------------------
14
15pub(crate) fn cosine(a: &[f32], b: &[f32]) -> f32 {
16    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
17    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
18    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
19    // Guard against a zero norm (undefined direction) and any non-finite input
20    // (a NaN/inf component propagates into `dot`/`norm_*`). A non-finite result
21    // would otherwise scatter arbitrarily during ranking, so score it 0.0 —
22    // ranked last, never NaN.
23    if norm_a == 0.0
24        || norm_b == 0.0
25        || !dot.is_finite()
26        || !norm_a.is_finite()
27        || !norm_b.is_finite()
28    {
29        0.0
30    } else {
31        let score = dot / (norm_a * norm_b);
32        if score.is_finite() { score } else { 0.0 }
33    }
34}
35
36// ---------------------------------------------------------------------------
37// MemoryVectorIndex
38// ---------------------------------------------------------------------------
39
40/// An exact, brute-force in-memory vector index.
41///
42/// All vectors must share the same dimension once the first entry is inserted.
43/// Cosine similarity is used for scoring; results are ordered by descending
44/// score with ties broken by [`RecordKey`] ordering (lexicographic), making
45/// results deterministic.
46#[derive(Default)]
47pub struct MemoryVectorIndex {
48    store: Mutex<HashMap<RecordKey, Vec<f32>>>,
49}
50
51impl MemoryVectorIndex {
52    /// Create an empty index.
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    /// Return the current dimension of stored vectors, or `None` if the index
58    /// is empty.
59    fn stored_dim(map: &HashMap<RecordKey, Vec<f32>>) -> Option<usize> {
60        map.values().next().map(|v| v.len())
61    }
62
63    fn check_dim(map: &HashMap<RecordKey, Vec<f32>>, incoming: usize) -> Result<()> {
64        if let Some(expected) = Self::stored_dim(map)
65            && incoming != expected
66        {
67            return Err(CoreError::Backend(format!(
68                "vector dimension mismatch: expected {expected}, got {incoming}"
69            )));
70        }
71        Ok(())
72    }
73
74    /// Entries whose key satisfies `pred`.
75    ///
76    /// Predicate-shaped rather than shard-shaped so sharding does not leak into
77    /// the in-memory index; a durable backend passes its own shard test here and
78    /// clones one shard rather than the whole index.
79    pub fn collect_where(&self, pred: impl Fn(&RecordKey) -> bool) -> Vec<(RecordKey, Vec<f32>)> {
80        let map = self.store.lock().expect("mutex poisoned");
81        map.iter()
82            .filter(|(key, _)| pred(key))
83            .map(|(key, vector)| (key.clone(), vector.clone()))
84            .collect()
85    }
86}
87
88#[async_trait]
89impl VectorIndex for MemoryVectorIndex {
90    async fn upsert(&self, key: RecordKey, vector: Vec<f32>) -> Result<()> {
91        let mut map = self.store.lock().expect("mutex poisoned");
92        Self::check_dim(&map, vector.len())?;
93        map.insert(key, vector);
94        Ok(())
95    }
96
97    async fn remove(&self, key: &RecordKey) -> Result<()> {
98        let mut map = self.store.lock().expect("mutex poisoned");
99        map.remove(key);
100        Ok(())
101    }
102
103    async fn query(&self, query: &[f32], k: usize, filter: &KeyPrefix) -> Result<Vec<Match>> {
104        let map = self.store.lock().expect("mutex poisoned");
105        if !map.is_empty() {
106            Self::check_dim(&map, query.len())?;
107        }
108
109        let mut matches: Vec<Match> = map
110            .iter()
111            .filter(|(key, _)| filter.matches(key))
112            .map(|(key, vec)| Match {
113                key: key.clone(),
114                score: cosine(query, vec),
115            })
116            .collect();
117
118        // Sort descending by score; break ties by RecordKey order (ascending).
119        // Use `f32::total_cmp` so ordering is total and deterministic even if a
120        // non-finite score (NaN/inf) slips through — `partial_cmp` returns
121        // `None` for NaN and would scatter such entries arbitrarily.
122        matches.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key)));
123
124        matches.truncate(k);
125        Ok(matches)
126    }
127
128    async fn keys(&self, filter: &KeyPrefix) -> Result<Vec<RecordKey>> {
129        let map = self.store.lock().expect("mutex poisoned");
130        Ok(map.keys().filter(|k| filter.matches(k)).cloned().collect())
131    }
132}
133
134// ---------------------------------------------------------------------------
135// Tests
136// ---------------------------------------------------------------------------
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use crate::Embedder;
142
143    /// Tiny deterministic embedder: maps a string to a 4-element Vec<f32> by
144    /// bucketing each byte value into one of four bins and summing them.
145    struct BucketEmbedder;
146
147    #[async_trait]
148    impl Embedder for BucketEmbedder {
149        async fn embed(&self, text: &str) -> Result<Vec<f32>> {
150            let mut out = vec![0.0f32; 4];
151            for (i, b) in text.bytes().enumerate() {
152                out[i % 4] += b as f32;
153            }
154            Ok(out)
155        }
156    }
157
158    // ------------------------------------------------------------------
159    // Test 1: nearest-first ordering
160    // ------------------------------------------------------------------
161    #[tokio::test]
162    async fn query_returns_nearest_first() {
163        let idx = MemoryVectorIndex::new();
164
165        // Three 2-D vectors pointing in clearly different directions.
166        idx.upsert(RecordKey::new("ns", "col", "east"), vec![1.0, 0.0])
167            .await
168            .unwrap();
169        idx.upsert(RecordKey::new("ns", "col", "north"), vec![0.0, 1.0])
170            .await
171            .unwrap();
172        idx.upsert(RecordKey::new("ns", "col", "neg"), vec![-1.0, 0.0])
173            .await
174            .unwrap();
175
176        // Query pointing almost east.
177        let results = idx
178            .query(&[0.99, 0.01], 3, &KeyPrefix::default())
179            .await
180            .unwrap();
181
182        assert_eq!(results[0].key.id, "east");
183    }
184
185    // ------------------------------------------------------------------
186    // Test 2: k limits results; k > index size returns all
187    // ------------------------------------------------------------------
188    #[tokio::test]
189    async fn k_limits_and_oversized_k() {
190        let idx = MemoryVectorIndex::new();
191        for i in 0..5u8 {
192            let v = vec![i as f32, 0.0];
193            idx.upsert(RecordKey::new("ns", "col", format!("{i}")), v)
194                .await
195                .unwrap();
196        }
197
198        // Non-zero query so we don't divide by zero on the [0.0, 0.0] vector.
199        let q = vec![1.0f32, 0.0];
200
201        let limited = idx.query(&q, 2, &KeyPrefix::default()).await.unwrap();
202        assert_eq!(limited.len(), 2);
203
204        let all = idx.query(&q, 100, &KeyPrefix::default()).await.unwrap();
205        assert_eq!(all.len(), 5);
206    }
207
208    // ------------------------------------------------------------------
209    // Test 3: filter by namespace
210    // ------------------------------------------------------------------
211    #[tokio::test]
212    async fn filter_restricts_to_namespace() {
213        let idx = MemoryVectorIndex::new();
214
215        idx.upsert(RecordKey::new("alpha", "col", "a1"), vec![1.0, 0.0])
216            .await
217            .unwrap();
218        idx.upsert(RecordKey::new("alpha", "col", "a2"), vec![1.0, 0.1])
219            .await
220            .unwrap();
221        idx.upsert(RecordKey::new("beta", "col", "b1"), vec![0.0, 1.0])
222            .await
223            .unwrap();
224
225        let filter = KeyPrefix {
226            namespace: Some("alpha".into()),
227            collection: None,
228        };
229        let results = idx.query(&[1.0, 0.0], 10, &filter).await.unwrap();
230
231        assert_eq!(results.len(), 2);
232        assert!(results.iter().all(|m| m.key.namespace == "alpha"));
233    }
234
235    // ------------------------------------------------------------------
236    // Test 4: dimension mismatch returns Backend error on upsert
237    // ------------------------------------------------------------------
238    #[tokio::test]
239    async fn upsert_dimension_mismatch_is_error() {
240        let idx = MemoryVectorIndex::new();
241        idx.upsert(RecordKey::new("ns", "col", "a"), vec![1.0, 0.0])
242            .await
243            .unwrap();
244
245        let err = idx
246            .upsert(RecordKey::new("ns", "col", "b"), vec![1.0, 0.0, 0.0])
247            .await
248            .unwrap_err();
249
250        assert!(
251            matches!(err, CoreError::Backend(ref msg) if msg.contains("dimension mismatch")),
252            "unexpected error: {err}"
253        );
254    }
255
256    // ------------------------------------------------------------------
257    // Test 5: remove drops key from results
258    // ------------------------------------------------------------------
259    #[tokio::test]
260    async fn remove_drops_key() {
261        let idx = MemoryVectorIndex::new();
262        let key = RecordKey::new("ns", "col", "target");
263        idx.upsert(key.clone(), vec![1.0, 0.0]).await.unwrap();
264        idx.upsert(RecordKey::new("ns", "col", "other"), vec![0.0, 1.0])
265            .await
266            .unwrap();
267
268        idx.remove(&key).await.unwrap();
269
270        let results = idx
271            .query(&[1.0, 0.0], 10, &KeyPrefix::default())
272            .await
273            .unwrap();
274        assert!(results.iter().all(|m| m.key != key));
275    }
276
277    // ------------------------------------------------------------------
278    // Test 6: cosine of identical direction is ~1.0
279    // ------------------------------------------------------------------
280    #[tokio::test]
281    async fn cosine_identical_direction_is_one() {
282        let idx = MemoryVectorIndex::new();
283        let v = vec![3.0f32, 4.0]; // magnitude 5
284        idx.upsert(RecordKey::new("ns", "col", "a"), v.clone())
285            .await
286            .unwrap();
287
288        let results = idx.query(&v, 1, &KeyPrefix::default()).await.unwrap();
289        assert!(
290            (results[0].score - 1.0).abs() < 1e-6,
291            "score={}",
292            results[0].score
293        );
294    }
295
296    // ------------------------------------------------------------------
297    // Test 7: Embedder trait integration
298    // ------------------------------------------------------------------
299    #[tokio::test]
300    async fn embedder_trait_integration() {
301        let embedder = BucketEmbedder;
302        let idx = MemoryVectorIndex::new();
303
304        let texts = ["hello", "world", "rust"];
305        for text in &texts {
306            let vec = embedder.embed(text).await.unwrap();
307            idx.upsert(RecordKey::new("ns", "col", *text), vec)
308                .await
309                .unwrap();
310        }
311
312        // Query with the same embedding as "hello" — it should be the top hit.
313        let query_vec = embedder.embed("hello").await.unwrap();
314        let results = idx
315            .query(&query_vec, 3, &KeyPrefix::default())
316            .await
317            .unwrap();
318
319        assert!(!results.is_empty());
320        assert_eq!(results[0].key.id, "hello");
321        // Self-similarity must be ~1.0.
322        assert!(
323            (results[0].score - 1.0).abs() < 1e-6,
324            "score={}",
325            results[0].score
326        );
327    }
328
329    // ------------------------------------------------------------------
330    // Bonus: remove absent key does not error
331    // ------------------------------------------------------------------
332    #[tokio::test]
333    async fn remove_absent_key_is_ok() {
334        let idx = MemoryVectorIndex::new();
335        let key = RecordKey::new("ns", "col", "ghost");
336        assert!(idx.remove(&key).await.is_ok());
337    }
338
339    // ------------------------------------------------------------------
340    // #154: non-finite components score 0.0, never NaN
341    // ------------------------------------------------------------------
342    #[test]
343    fn cosine_nan_component_scores_zero() {
344        // A NaN in either operand must score 0.0, not NaN.
345        assert_eq!(cosine(&[f32::NAN, 1.0], &[1.0, 0.0]), 0.0);
346        assert_eq!(cosine(&[1.0, 0.0], &[f32::NAN, 1.0]), 0.0);
347        // An infinity likewise scores 0.0 rather than a non-finite value.
348        assert_eq!(cosine(&[f32::INFINITY, 1.0], &[1.0, 0.0]), 0.0);
349        assert_eq!(cosine(&[1.0, 0.0], &[f32::NEG_INFINITY, 1.0]), 0.0);
350        // The pre-existing zero-norm guard still holds.
351        assert_eq!(cosine(&[0.0, 0.0], &[1.0, 0.0]), 0.0);
352        // A genuine match still scores ~1.0.
353        assert!((cosine(&[3.0, 4.0], &[3.0, 4.0]) - 1.0).abs() < 1e-6);
354    }
355
356    // ------------------------------------------------------------------
357    // #154: a NaN-bearing vector never outranks a genuinely similar one
358    // ------------------------------------------------------------------
359    #[tokio::test]
360    async fn nan_vector_never_occupies_top_k_ahead_of_match() {
361        let idx = MemoryVectorIndex::new();
362
363        // A vector aligned with the query (should rank first) ...
364        idx.upsert(RecordKey::new("ns", "col", "good"), vec![1.0, 0.0])
365            .await
366            .unwrap();
367        // ... and a poisoned vector with a NaN component.
368        idx.upsert(RecordKey::new("ns", "col", "poison"), vec![f32::NAN, 1.0])
369            .await
370            .unwrap();
371
372        let results = idx
373            .query(&[1.0, 0.0], 2, &KeyPrefix::default())
374            .await
375            .unwrap();
376
377        assert_eq!(results.len(), 2);
378        // The genuine match ranks first; the poisoned vector scores 0.0 last.
379        assert_eq!(results[0].key.id, "good");
380        assert_eq!(results[1].key.id, "poison");
381        assert!(results.iter().all(|m| m.score.is_finite()));
382        assert_eq!(results[1].score, 0.0);
383    }
384
385    // ------------------------------------------------------------------
386    // #154: ranking is deterministic across repeated queries
387    // ------------------------------------------------------------------
388    #[tokio::test]
389    async fn ranking_is_deterministic_with_poisoned_vectors() {
390        let idx = MemoryVectorIndex::new();
391        idx.upsert(RecordKey::new("ns", "col", "a"), vec![1.0, 0.0])
392            .await
393            .unwrap();
394        idx.upsert(RecordKey::new("ns", "col", "b"), vec![f32::NAN, 1.0])
395            .await
396            .unwrap();
397        idx.upsert(RecordKey::new("ns", "col", "c"), vec![f32::INFINITY, 0.0])
398            .await
399            .unwrap();
400        idx.upsert(RecordKey::new("ns", "col", "d"), vec![0.9, 0.1])
401            .await
402            .unwrap();
403
404        let first = idx
405            .query(&[1.0, 0.0], 4, &KeyPrefix::default())
406            .await
407            .unwrap();
408        for _ in 0..10 {
409            let again = idx
410                .query(&[1.0, 0.0], 4, &KeyPrefix::default())
411                .await
412                .unwrap();
413            let ids_first: Vec<_> = first.iter().map(|m| m.key.id.clone()).collect();
414            let ids_again: Vec<_> = again.iter().map(|m| m.key.id.clone()).collect();
415            assert_eq!(ids_first, ids_again, "ordering must be deterministic");
416        }
417    }
418
419    // ------------------------------------------------------------------
420    // Bonus: query dimension mismatch is error
421    // ------------------------------------------------------------------
422    #[tokio::test]
423    async fn query_dimension_mismatch_is_error() {
424        let idx = MemoryVectorIndex::new();
425        idx.upsert(RecordKey::new("ns", "col", "a"), vec![1.0, 0.0])
426            .await
427            .unwrap();
428
429        let err = idx
430            .query(&[1.0, 0.0, 0.0], 1, &KeyPrefix::default())
431            .await
432            .unwrap_err();
433
434        assert!(
435            matches!(err, CoreError::Backend(ref msg) if msg.contains("dimension mismatch")),
436            "unexpected error: {err}"
437        );
438    }
439
440    // ------------------------------------------------------------------
441    // #323: keys, upsert_many, and collect_where
442    // ------------------------------------------------------------------
443    #[tokio::test]
444    async fn keys_lists_everything_under_the_filter() {
445        let idx = MemoryVectorIndex::new();
446        idx.upsert(RecordKey::new("ns", "a", "1"), vec![1.0])
447            .await
448            .unwrap();
449        idx.upsert(RecordKey::new("ns", "b", "2"), vec![2.0])
450            .await
451            .unwrap();
452
453        let all = idx.keys(&KeyPrefix::default()).await.unwrap();
454        assert_eq!(all.len(), 2);
455
456        let only_a = idx
457            .keys(&KeyPrefix {
458                namespace: None,
459                collection: Some("a".into()),
460            })
461            .await
462            .unwrap();
463        assert_eq!(only_a, vec![RecordKey::new("ns", "a", "1")]);
464    }
465
466    #[tokio::test]
467    async fn an_arc_delegates_every_method_including_upsert_many() {
468        use std::sync::Arc;
469        let idx: Arc<MemoryVectorIndex> = Arc::new(MemoryVectorIndex::new());
470        idx.upsert_many(vec![(RecordKey::new("ns", "c", "1"), vec![1.0])])
471            .await
472            .unwrap();
473        assert_eq!(idx.keys(&KeyPrefix::default()).await.unwrap().len(), 1);
474        idx.remove(&RecordKey::new("ns", "c", "1")).await.unwrap();
475        assert!(idx.keys(&KeyPrefix::default()).await.unwrap().is_empty());
476    }
477
478    #[tokio::test]
479    async fn upsert_many_inserts_every_item() {
480        let idx = MemoryVectorIndex::new();
481        idx.upsert_many(vec![
482            (RecordKey::new("ns", "c", "1"), vec![1.0]),
483            (RecordKey::new("ns", "c", "2"), vec![2.0]),
484        ])
485        .await
486        .unwrap();
487
488        assert_eq!(idx.keys(&KeyPrefix::default()).await.unwrap().len(), 2);
489    }
490
491    #[tokio::test]
492    async fn collect_where_returns_only_matching_entries() {
493        let idx = MemoryVectorIndex::new();
494        idx.upsert(RecordKey::new("ns", "c", "keep"), vec![1.0])
495            .await
496            .unwrap();
497        idx.upsert(RecordKey::new("ns", "c", "drop"), vec![2.0])
498            .await
499            .unwrap();
500
501        let got = idx.collect_where(|k| k.id == "keep");
502        assert_eq!(got, vec![(RecordKey::new("ns", "c", "keep"), vec![1.0])]);
503    }
504}
gonzalo_vector/index.rs

gonzalo_vector/
index.rs
