1use std::collections::HashMap;
22use std::sync::Arc;
23
24use parking_lot::RwLock;
25use serde::{Deserialize, Serialize};
26use thiserror::Error;
27
28use crate::distance::Distance;
29use crate::encoding::{Codec, EncodedVector, EncodingError};
30use crate::index::{HnswIndex, HnswParams, IndexError, NodeId, SearchResult};
31use crate::turbo_hnsw::TurboHnswIndex;
32use crate::turbo_index::TurboTable;
33
34pub type RowKey = Vec<u8>;
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
50#[serde(rename_all = "snake_case")]
51#[non_exhaustive]
52#[derive(Default)]
53pub enum IndexAlgorithm {
54 #[default]
56 Hnsw,
57 Flat,
59}
60
61#[derive(Clone, Debug, Serialize, Deserialize)]
63pub struct TableSchema {
64 pub name: String,
66 pub dim: u16,
68 pub codec: Codec,
70 pub distance: Distance,
72 pub hnsw: HnswParams,
76 #[serde(default)]
80 pub algorithm: IndexAlgorithm,
81}
82
83#[derive(Clone, Debug, Serialize, Deserialize)]
85pub struct VectorRow {
86 pub key: RowKey,
88 pub vector: EncodedVector,
90 pub metadata: HashMap<String, serde_json::Value>,
93 pub created_at: u64,
95 pub updated_at: u64,
98}
99
100#[derive(Debug, Error)]
102#[non_exhaustive]
103pub enum StoreError {
104 #[error("table not found: {0}")]
106 UnknownTable(String),
107 #[error("table already exists: {0}")]
109 TableExists(String),
110 #[error("dimension mismatch: table {table} expects {expected}, got {got}")]
112 DimensionMismatch {
113 table: String,
115 expected: u16,
117 got: u16,
119 },
120 #[error("row not found in table {table}: {key:?}")]
122 RowNotFound {
123 table: String,
125 key: RowKey,
127 },
128 #[error("encoding: {0}")]
130 Encoding(#[from] EncodingError),
131 #[error("index: {0}")]
133 Index(#[from] IndexError),
134 #[error("backend: {0}")]
136 Backend(String),
137}
138
139pub trait Backend: Send + Sync {
146 fn put_row(&self, table: &str, key: &[u8], row: &VectorRow) -> Result<(), StoreError>;
149
150 fn get_row(&self, table: &str, key: &[u8]) -> Result<Option<VectorRow>, StoreError>;
153
154 fn delete_row(&self, table: &str, key: &[u8]) -> Result<bool, StoreError>;
157
158 fn for_each_row(&self, table: &str, f: &mut RowVisitor<'_>) -> Result<(), StoreError>;
161
162 fn put_schema(&self, schema: &TableSchema) -> Result<(), StoreError>;
164
165 fn list_schemas(&self) -> Result<Vec<TableSchema>, StoreError>;
167}
168
169pub type RowVisitor<'a> = dyn FnMut(&[u8], &VectorRow) -> Result<(), StoreError> + 'a;
171
172#[derive(Default)]
176pub struct MemoryBackend {
177 rows: RwLock<HashMap<String, HashMap<Vec<u8>, VectorRow>>>,
178 schemas: RwLock<HashMap<String, TableSchema>>,
179}
180
181impl MemoryBackend {
182 #[must_use]
184 pub fn new() -> Self {
185 Self::default()
186 }
187}
188
189impl Backend for MemoryBackend {
190 fn put_row(&self, table: &str, key: &[u8], row: &VectorRow) -> Result<(), StoreError> {
191 let mut rows = self.rows.write();
192 let entry = rows.entry(table.to_string()).or_default();
193 entry.insert(key.to_vec(), row.clone());
194 Ok(())
195 }
196
197 fn get_row(&self, table: &str, key: &[u8]) -> Result<Option<VectorRow>, StoreError> {
198 let rows = self.rows.read();
199 Ok(rows.get(table).and_then(|m| m.get(key).cloned()))
200 }
201
202 fn delete_row(&self, table: &str, key: &[u8]) -> Result<bool, StoreError> {
203 let mut rows = self.rows.write();
204 Ok(rows.get_mut(table).is_some_and(|m| m.remove(key).is_some()))
205 }
206
207 fn for_each_row(&self, table: &str, f: &mut RowVisitor<'_>) -> Result<(), StoreError> {
208 let rows = self.rows.read();
209 if let Some(m) = rows.get(table) {
210 for (k, v) in m {
211 f(k, v)?;
212 }
213 }
214 Ok(())
215 }
216
217 fn put_schema(&self, schema: &TableSchema) -> Result<(), StoreError> {
218 self.schemas
219 .write()
220 .insert(schema.name.clone(), schema.clone());
221 Ok(())
222 }
223
224 fn list_schemas(&self) -> Result<Vec<TableSchema>, StoreError> {
225 Ok(self.schemas.read().values().cloned().collect())
226 }
227}
228
229struct TableState {
233 schema: TableSchema,
234 ann: AnnContainer,
235 key_to_node: HashMap<RowKey, NodeId>,
237 node_to_key: HashMap<NodeId, RowKey>,
240 next_node_id: NodeId,
244}
245
246enum AnnContainer {
262 Hnsw(HnswIndex),
263 TurboFlat(TurboTable),
264 TurboHnsw2(TurboHnswIndex<2>),
265 TurboHnsw3(TurboHnswIndex<3>),
266 TurboHnsw4(TurboHnswIndex<4>),
267}
268
269impl AnnContainer {
270 fn new(schema: &TableSchema) -> Result<Self, StoreError> {
271 if let Some(bits) = schema.codec.turbovec_bits() {
272 match schema.algorithm {
273 IndexAlgorithm::Flat => {
274 let table = TurboTable::new(schema.distance, schema.dim, bits)?;
275 Ok(Self::TurboFlat(table))
276 }
277 IndexAlgorithm::Hnsw => match bits {
278 2 => Ok(Self::TurboHnsw2(TurboHnswIndex::<2>::new(
279 schema.distance,
280 schema.dim,
281 schema.hnsw,
282 )?)),
283 3 => Ok(Self::TurboHnsw3(TurboHnswIndex::<3>::new(
284 schema.distance,
285 schema.dim,
286 schema.hnsw,
287 )?)),
288 4 => Ok(Self::TurboHnsw4(TurboHnswIndex::<4>::new(
289 schema.distance,
290 schema.dim,
291 schema.hnsw,
292 )?)),
293 _ => Err(StoreError::Index(IndexError::Empty)),
294 },
295 }
296 } else {
297 Ok(Self::Hnsw(HnswIndex::new(schema.distance, schema.hnsw)))
301 }
302 }
303
304 fn insert(&mut self, id: NodeId, vector: Vec<f32>) -> Result<(), IndexError> {
305 match self {
306 Self::Hnsw(idx) => idx.insert(id, vector),
307 Self::TurboFlat(t) => t.insert(id, vector),
308 Self::TurboHnsw2(t) => t.insert(id, vector),
309 Self::TurboHnsw3(t) => t.insert(id, vector),
310 Self::TurboHnsw4(t) => t.insert(id, vector),
311 }
312 }
313
314 fn delete(&mut self, id: NodeId) -> bool {
315 match self {
316 Self::Hnsw(idx) => idx.delete(id),
317 Self::TurboFlat(t) => t.delete(id),
318 Self::TurboHnsw2(t) => t.delete(id),
319 Self::TurboHnsw3(t) => t.delete(id),
320 Self::TurboHnsw4(t) => t.delete(id),
321 }
322 }
323
324 fn search(
325 &self,
326 query: &[f32],
327 k: usize,
328 ef: Option<usize>,
329 ) -> Result<Vec<SearchResult>, IndexError> {
330 match self {
331 Self::Hnsw(idx) => idx.search(query, k, ef),
332 Self::TurboFlat(t) => t.search(query, k, ef),
333 Self::TurboHnsw2(t) => t.search(query, k, ef),
334 Self::TurboHnsw3(t) => t.search(query, k, ef),
335 Self::TurboHnsw4(t) => t.search(query, k, ef),
336 }
337 }
338
339 fn len(&self) -> usize {
340 match self {
341 Self::Hnsw(idx) => idx.len(),
342 Self::TurboFlat(t) => t.len(),
343 Self::TurboHnsw2(t) => t.len(),
344 Self::TurboHnsw3(t) => t.len(),
345 Self::TurboHnsw4(t) => t.len(),
346 }
347 }
348}
349
350pub struct VectorStore {
352 backend: Arc<dyn Backend>,
353 tables: RwLock<HashMap<String, Arc<parking_lot::Mutex<TableState>>>>,
354}
355
356impl VectorStore {
357 pub fn open(backend: Arc<dyn Backend>) -> Result<Self, StoreError> {
370 let tables = RwLock::new(HashMap::new());
371 let store = Self { backend, tables };
372 let schemas = store.backend.list_schemas()?;
373 for schema in schemas {
374 store.rehydrate_table(&schema)?;
375 }
376 Ok(store)
377 }
378
379 #[must_use]
382 pub fn in_memory() -> Self {
383 Self {
384 backend: Arc::new(MemoryBackend::new()),
385 tables: RwLock::new(HashMap::new()),
386 }
387 }
388
389 pub fn create_table(&self, schema: TableSchema) -> Result<(), StoreError> {
396 let mut tables = self.tables.write();
397 if tables.contains_key(&schema.name) {
398 return Err(StoreError::TableExists(schema.name));
399 }
400 let state = TableState {
401 schema: schema.clone(),
402 ann: AnnContainer::new(&schema)?,
403 key_to_node: HashMap::new(),
404 node_to_key: HashMap::new(),
405 next_node_id: 1,
406 };
407 self.backend.put_schema(&schema)?;
408 tables.insert(
409 schema.name.clone(),
410 Arc::new(parking_lot::Mutex::new(state)),
411 );
412 Ok(())
413 }
414
415 pub fn tables(&self) -> Vec<TableSchema> {
417 self.tables
418 .read()
419 .values()
420 .map(|s| s.lock().schema.clone())
421 .collect()
422 }
423
424 pub fn upsert(
439 &self,
440 table: &str,
441 key: RowKey,
442 vector: &[f32],
443 metadata: HashMap<String, serde_json::Value>,
444 ) -> Result<(), StoreError> {
445 let state = self.table_state(table)?;
446 let mut state = state.lock();
447 let dim = u16::try_from(vector.len()).unwrap_or(u16::MAX);
448 if dim != state.schema.dim {
449 return Err(StoreError::DimensionMismatch {
450 table: table.to_string(),
451 expected: state.schema.dim,
452 got: dim,
453 });
454 }
455 let codec_encoder = state.schema.codec.encoder();
456 let encoded = codec_encoder.encode(vector)?;
457 let now = now_millis();
458 let prior = self.backend.get_row(table, &key)?;
459 let row = VectorRow {
460 key: key.clone(),
461 vector: encoded,
462 metadata,
463 created_at: prior.as_ref().map_or(now, |r| r.created_at),
464 updated_at: now,
465 };
466 self.backend.put_row(table, &key, &row)?;
467 if let Some(&old_node) = state.key_to_node.get(&key) {
468 state.ann.delete(old_node);
469 state.node_to_key.remove(&old_node);
470 }
471 let node_id = state.next_node_id;
472 state.next_node_id += 1;
473 state.ann.insert(node_id, vector.to_vec())?;
474 state.key_to_node.insert(key.clone(), node_id);
475 state.node_to_key.insert(node_id, key);
476 Ok(())
477 }
478
479 pub fn get(&self, table: &str, key: &[u8]) -> Result<Option<VectorRow>, StoreError> {
486 let _ = self.table_state(table)?;
487 self.backend.get_row(table, key)
488 }
489
490 pub fn delete(&self, table: &str, key: &[u8]) -> Result<bool, StoreError> {
498 let state = self.table_state(table)?;
499 let mut state = state.lock();
500 let removed = self.backend.delete_row(table, key)?;
501 if let Some(node_id) = state.key_to_node.remove(key) {
502 state.ann.delete(node_id);
503 state.node_to_key.remove(&node_id);
504 }
505 Ok(removed)
506 }
507
508 pub fn search(
518 &self,
519 table: &str,
520 query: &[f32],
521 k: usize,
522 ef: Option<usize>,
523 ) -> Result<Vec<(VectorRow, f32)>, StoreError> {
524 let state = self.table_state(table)?;
525 let state = state.lock();
526 let dim = u16::try_from(query.len()).unwrap_or(u16::MAX);
527 if dim != state.schema.dim {
528 return Err(StoreError::DimensionMismatch {
529 table: table.to_string(),
530 expected: state.schema.dim,
531 got: dim,
532 });
533 }
534 let hits: Vec<SearchResult> = state.ann.search(query, k, ef)?;
535 let mut out = Vec::with_capacity(hits.len());
536 for hit in hits {
537 if let Some(key) = state.node_to_key.get(&hit.id) {
538 if let Some(row) = self.backend.get_row(table, key)? {
539 out.push((row, hit.score));
540 }
541 }
542 }
543 Ok(out)
544 }
545
546 pub fn stats(&self, table: &str) -> Result<TableStats, StoreError> {
553 let state = self.table_state(table)?;
554 let state = state.lock();
555 Ok(TableStats {
556 name: state.schema.name.clone(),
557 dim: state.schema.dim,
558 codec: state.schema.codec,
559 distance: state.schema.distance,
560 live_rows: state.ann.len(),
561 tracked_rows: state.key_to_node.len(),
562 })
563 }
564
565 fn table_state(&self, table: &str) -> Result<Arc<parking_lot::Mutex<TableState>>, StoreError> {
566 self.tables
567 .read()
568 .get(table)
569 .cloned()
570 .ok_or_else(|| StoreError::UnknownTable(table.to_string()))
571 }
572
573 fn rehydrate_table(&self, schema: &TableSchema) -> Result<(), StoreError> {
574 let state = TableState {
575 schema: schema.clone(),
576 ann: AnnContainer::new(schema)?,
577 key_to_node: HashMap::new(),
578 node_to_key: HashMap::new(),
579 next_node_id: 1,
580 };
581 let cell = Arc::new(parking_lot::Mutex::new(state));
582 self.tables
583 .write()
584 .insert(schema.name.clone(), cell.clone());
585 let mut guard = cell.lock();
586 let encoder = guard.schema.codec.encoder();
587 let mut to_insert: Vec<(NodeId, RowKey, Vec<f32>)> = Vec::new();
588 let table_name = schema.name.clone();
589 let mut next = 1u64;
590 self.backend.for_each_row(&table_name, &mut |k, row| {
591 let v = encoder.decode(&row.vector)?;
592 to_insert.push((next, k.to_vec(), v));
593 next += 1;
594 Ok(())
595 })?;
596 for (node, key, v) in to_insert {
597 guard.ann.insert(node, v)?;
598 guard.key_to_node.insert(key.clone(), node);
599 guard.node_to_key.insert(node, key);
600 guard.next_node_id = node + 1;
601 }
602 Ok(())
603 }
604}
605
606#[derive(Clone, Debug, Serialize, Deserialize)]
608pub struct TableStats {
609 pub name: String,
611 pub dim: u16,
613 pub codec: Codec,
615 pub distance: Distance,
617 pub live_rows: usize,
619 pub tracked_rows: usize,
621}
622
623fn now_millis() -> u64 {
624 use std::time::{SystemTime, UNIX_EPOCH};
625 SystemTime::now()
626 .duration_since(UNIX_EPOCH)
627 .map_or(0, |d| u64::try_from(d.as_millis()).unwrap_or(u64::MAX))
628}
629
630#[cfg(test)]
631mod tests {
632 use super::*;
633 use crate::index::HnswParams;
634
635 fn schema(name: &str, dim: u16) -> TableSchema {
636 TableSchema {
637 name: name.to_string(),
638 dim,
639 codec: Codec::Int8Quantized,
640 distance: Distance::Euclidean,
641 hnsw: HnswParams::default(),
642 algorithm: IndexAlgorithm::Hnsw,
643 }
644 }
645
646 #[test]
647 fn create_and_list_tables() {
648 let store = VectorStore::in_memory();
649 store.create_table(schema("t", 4)).unwrap();
650 let tables = store.tables();
651 assert_eq!(tables.len(), 1);
652 assert_eq!(tables[0].name, "t");
653 assert_eq!(tables[0].dim, 4);
654 }
655
656 #[test]
657 fn duplicate_table_rejected() {
658 let store = VectorStore::in_memory();
659 store.create_table(schema("t", 4)).unwrap();
660 assert!(matches!(
661 store.create_table(schema("t", 4)),
662 Err(StoreError::TableExists(_))
663 ));
664 }
665
666 #[test]
667 fn upsert_get_delete_round_trip() {
668 let store = VectorStore::in_memory();
669 store.create_table(schema("t", 3)).unwrap();
670 store
671 .upsert("t", b"a".to_vec(), &[1.0, 2.0, 3.0], HashMap::new())
672 .unwrap();
673 let row = store.get("t", b"a").unwrap().expect("row present");
674 assert_eq!(row.key, b"a");
675 assert_eq!(row.vector.dim, 3);
676 assert!(store.delete("t", b"a").unwrap());
677 assert!(store.get("t", b"a").unwrap().is_none());
678 assert!(!store.delete("t", b"a").unwrap());
679 }
680
681 #[test]
682 fn dimension_mismatch_rejected() {
683 let store = VectorStore::in_memory();
684 store.create_table(schema("t", 3)).unwrap();
685 assert!(matches!(
686 store.upsert("t", b"a".to_vec(), &[1.0, 2.0], HashMap::new()),
687 Err(StoreError::DimensionMismatch { .. })
688 ));
689 }
690
691 #[test]
692 fn search_returns_nearest_first() {
693 let store = VectorStore::in_memory();
694 store.create_table(schema("t", 2)).unwrap();
695 for (k, v) in [
696 (&b"origin"[..], [0.0_f32, 0.0]),
697 (&b"unit_x"[..], [1.0, 0.0]),
698 (&b"unit_y"[..], [0.0, 1.0]),
699 (&b"diag"[..], [1.0, 1.0]),
700 ] {
701 store.upsert("t", k.to_vec(), &v, HashMap::new()).unwrap();
702 }
703 let res = store.search("t", &[0.05, 0.05], 1, None).unwrap();
704 assert_eq!(res.len(), 1);
705 assert_eq!(res[0].0.key, b"origin");
706 }
707
708 #[test]
709 fn rehydrate_rebuilds_index() {
710 let backend = Arc::new(MemoryBackend::new());
711 let store = VectorStore::open(backend.clone()).unwrap();
712 store.create_table(schema("t", 2)).unwrap();
713 for i in 0..10_u8 {
714 let k = format!("k{i}").into_bytes();
715 let v = [f32::from(i), f32::from(i) * 2.0];
716 store.upsert("t", k, &v, HashMap::new()).unwrap();
717 }
718 drop(store);
720 let reopened = VectorStore::open(backend).unwrap();
721 let stats = reopened.stats("t").unwrap();
722 assert_eq!(stats.live_rows, 10);
723 let res = reopened.search("t", &[3.0, 6.0], 1, None).unwrap();
724 assert_eq!(res[0].0.key, b"k3");
725 }
726
727 #[test]
728 fn stats_reports_live_rows() {
729 let store = VectorStore::in_memory();
730 store.create_table(schema("t", 2)).unwrap();
731 store
732 .upsert("t", b"a".to_vec(), &[1.0, 2.0], HashMap::new())
733 .unwrap();
734 store
735 .upsert("t", b"b".to_vec(), &[3.0, 4.0], HashMap::new())
736 .unwrap();
737 let s = store.stats("t").unwrap();
738 assert_eq!(s.live_rows, 2);
739 assert_eq!(s.tracked_rows, 2);
740 }
741}