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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
//! In-memory vector index implementations and the `VectorIndex` trait.
use anyhow::Result;
use std::collections::HashMap;
use crate::similarity;
use crate::Vector;
use crate::VectorId;
/// Vector index trait for efficient similarity search
pub trait VectorIndex: Send + Sync {
/// Insert a vector with associated URI
fn insert(&mut self, uri: String, vector: Vector) -> Result<()>;
/// Find k nearest neighbors
fn search_knn(&self, query: &Vector, k: usize) -> Result<Vec<(String, f32)>>;
/// Find all vectors within threshold similarity
fn search_threshold(&self, query: &Vector, threshold: f32) -> Result<Vec<(String, f32)>>;
/// Get a vector by its URI
fn get_vector(&self, uri: &str) -> Option<&Vector>;
/// Add a vector with associated ID and metadata
fn add_vector(
&mut self,
id: VectorId,
vector: Vector,
_metadata: Option<HashMap<String, String>>,
) -> Result<()> {
// Default implementation that delegates to insert
self.insert(id, vector)
}
/// Update an existing vector
fn update_vector(&mut self, id: VectorId, vector: Vector) -> Result<()> {
// Default implementation that delegates to insert
self.insert(id, vector)
}
/// Update metadata for a vector
fn update_metadata(&mut self, _id: VectorId, _metadata: HashMap<String, String>) -> Result<()> {
// Default implementation (no-op)
Ok(())
}
/// Remove a vector by its ID
fn remove_vector(&mut self, _id: VectorId) -> Result<()> {
// Default implementation (no-op)
Ok(())
}
/// Iterate all stored (id, vector) pairs.
///
/// The default returns an empty list; concrete index types that hold their
/// vectors in memory (or can reconstruct them, e.g. via decoding quantized
/// codes) should override this **and** [`VectorIndex::supports_enumeration`]
/// so callers like `VectorStore::save_to_disk` can tell real emptiness
/// apart from "this index type cannot enumerate its vectors".
fn iter_vectors(&self) -> Vec<(String, Vector)> {
Vec::new()
}
/// Whether [`VectorIndex::iter_vectors`] returns a real, complete
/// enumeration of the vectors held by this index.
///
/// Index types that override `iter_vectors` with a real implementation
/// (e.g. [`MemoryVectorIndex`], `HnswIndex`, `IvfIndex`, `PQIndex`) must
/// also override this to return `true`. Callers that need to persist or
/// otherwise fully enumerate an index (e.g. `VectorStore::save_to_disk`)
/// should check this flag and fail loudly instead of silently persisting
/// an empty snapshot when it is `false`.
fn supports_enumeration(&self) -> bool {
false
}
}
/// In-memory vector index implementation
pub struct MemoryVectorIndex {
vectors: Vec<(String, Vector)>,
similarity_config: similarity::SimilarityConfig,
}
impl MemoryVectorIndex {
/// Create a new empty in-memory vector index with default similarity config.
pub fn new() -> Self {
Self {
vectors: Vec::new(),
similarity_config: similarity::SimilarityConfig::default(),
}
}
/// Create a new in-memory vector index with a custom similarity configuration.
pub fn with_similarity_config(config: similarity::SimilarityConfig) -> Self {
Self {
vectors: Vec::new(),
similarity_config: config,
}
}
}
impl Default for MemoryVectorIndex {
fn default() -> Self {
Self::new()
}
}
impl VectorIndex for MemoryVectorIndex {
fn insert(&mut self, uri: String, vector: Vector) -> Result<()> {
// Check if vector already exists and update it
if let Some(pos) = self.vectors.iter().position(|(id, _)| id == &uri) {
self.vectors[pos] = (uri, vector);
} else {
self.vectors.push((uri, vector));
}
Ok(())
}
fn search_knn(&self, query: &Vector, k: usize) -> Result<Vec<(String, f32)>> {
let metric = self.similarity_config.primary_metric;
let query_f32 = query.as_f32();
let mut similarities: Vec<(String, f32)> = self
.vectors
.iter()
.map(|(uri, vec)| {
let vec_f32 = vec.as_f32();
let sim = metric.similarity(&query_f32, &vec_f32).unwrap_or(0.0);
(uri.clone(), sim)
})
.collect();
similarities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
similarities.truncate(k);
Ok(similarities)
}
fn search_threshold(&self, query: &Vector, threshold: f32) -> Result<Vec<(String, f32)>> {
let metric = self.similarity_config.primary_metric;
let query_f32 = query.as_f32();
let similarities: Vec<(String, f32)> = self
.vectors
.iter()
.filter_map(|(uri, vec)| {
let vec_f32 = vec.as_f32();
let sim = metric.similarity(&query_f32, &vec_f32).unwrap_or(0.0);
if sim >= threshold {
Some((uri.clone(), sim))
} else {
None
}
})
.collect();
Ok(similarities)
}
fn get_vector(&self, uri: &str) -> Option<&Vector> {
self.vectors.iter().find(|(u, _)| u == uri).map(|(_, v)| v)
}
fn update_vector(&mut self, id: VectorId, vector: Vector) -> Result<()> {
if let Some(pos) = self.vectors.iter().position(|(uri, _)| uri == &id) {
self.vectors[pos] = (id, vector);
Ok(())
} else {
Err(anyhow::anyhow!("Vector with id '{}' not found", id))
}
}
fn remove_vector(&mut self, id: VectorId) -> Result<()> {
if let Some(pos) = self.vectors.iter().position(|(uri, _)| uri == &id) {
self.vectors.remove(pos);
Ok(())
} else {
Err(anyhow::anyhow!("Vector with id '{}' not found", id))
}
}
fn iter_vectors(&self) -> Vec<(String, Vector)> {
self.vectors.clone()
}
fn supports_enumeration(&self) -> bool {
true
}
}