semtree_store/
usearch.rs

1use std::collections::HashMap;
2use std::path::Path;
3use std::sync::RwLock;
4
5use async_trait::async_trait;
6use semtree_embed::Embedding;
7use usearch::{Index, IndexOptions, MetricKind, ScalarKind};
8
9use crate::{Hit, StoreError, VectorStore};
10
11pub struct UsearchStore {
12    dimensions: usize,
13    index: Index,
14    /// maps usearch integer key → chunk id string
15    id_map: RwLock<HashMap<u64, String>>,
16    next_key: RwLock<u64>,
17}
18
19impl UsearchStore {
20    pub fn new(dimensions: usize) -> Result<Self, StoreError> {
21        let index = Self::make_index(dimensions)?;
22        Ok(Self {
23            dimensions,
24            index,
25            id_map: RwLock::new(HashMap::new()),
26            next_key: RwLock::new(0),
27        })
28    }
29
30    fn make_index(dimensions: usize) -> Result<Index, StoreError> {
31        let options = IndexOptions {
32            dimensions,
33            metric: MetricKind::Cos,
34            quantization: ScalarKind::F32,
35            ..Default::default()
36        };
37        let index = Index::new(&options).map_err(|e| StoreError::Init(e.to_string()))?;
38        index
39            .reserve(1024)
40            .map_err(|e| StoreError::Init(e.to_string()))?;
41        Ok(index)
42    }
43}
44
45#[async_trait]
46impl VectorStore for UsearchStore {
47    async fn insert(&self, id: &str, embedding: &Embedding) -> Result<(), StoreError> {
48        let key = {
49            let mut k = self.next_key.write().unwrap();
50            let current = *k;
51            *k += 1;
52            current
53        };
54        self.index
55            .add(key, embedding)
56            .map_err(|e| StoreError::Insert(e.to_string()))?;
57        self.id_map.write().unwrap().insert(key, id.to_string());
58        Ok(())
59    }
60
61    async fn search(&self, query: &Embedding, top_k: usize) -> Result<Vec<Hit>, StoreError> {
62        if self.index.size() == 0 {
63            return Ok(vec![]);
64        }
65        let results = self
66            .index
67            .search(query, top_k)
68            .map_err(|e| StoreError::Search(e.to_string()))?;
69
70        let map = self.id_map.read().unwrap();
71        let hits = results
72            .keys
73            .iter()
74            .zip(results.distances.iter())
75            .filter_map(|(key, dist)| {
76                map.get(key).map(|id| Hit {
77                    id: id.clone(),
78                    score: 1.0 - dist,
79                })
80            })
81            .collect();
82        Ok(hits)
83    }
84
85    async fn delete(&self, id: &str) -> Result<(), StoreError> {
86        let map = self.id_map.read().unwrap();
87        if let Some((&key, _)) = map.iter().find(|(_, v)| v.as_str() == id) {
88            drop(map);
89            self.index
90                .remove(key)
91                .map_err(|e| StoreError::Insert(e.to_string()))?;
92            self.id_map.write().unwrap().remove(&key);
93        }
94        Ok(())
95    }
96
97    fn save(&self, path: &Path) -> Result<(), StoreError> {
98        // save vector index
99        let index_path = path.join("index.usearch");
100        self.index
101            .save(index_path.to_str().unwrap())
102            .map_err(|e| StoreError::Init(e.to_string()))?;
103
104        // save id_map + next_key as JSON
105        let meta = serde_json::json!({
106            "dimensions": self.dimensions,
107            "next_key": *self.next_key.read().unwrap(),
108            "id_map": self.id_map.read().unwrap().iter()
109                .map(|(k, v)| (k.to_string(), v.clone()))
110                .collect::<HashMap<String, String>>(),
111        });
112        let meta_path = path.join("meta.json");
113        std::fs::write(meta_path, serde_json::to_string(&meta).unwrap())
114            .map_err(|e| StoreError::Init(e.to_string()))?;
115
116        Ok(())
117    }
118
119    fn load(&mut self, path: &Path) -> Result<(), StoreError> {
120        let meta_path = path.join("meta.json");
121        let raw =
122            std::fs::read_to_string(meta_path).map_err(|e| StoreError::Init(e.to_string()))?;
123        let meta: serde_json::Value =
124            serde_json::from_str(&raw).map_err(|e| StoreError::Init(e.to_string()))?;
125
126        let next_key: u64 = meta["next_key"].as_u64().unwrap_or(0);
127        let id_map: HashMap<u64, String> = meta["id_map"]
128            .as_object()
129            .unwrap_or(&serde_json::Map::new())
130            .iter()
131            .filter_map(|(k, v)| Some((k.parse::<u64>().ok()?, v.as_str()?.to_string())))
132            .collect();
133
134        *self.next_key.write().unwrap() = next_key;
135        *self.id_map.write().unwrap() = id_map;
136
137        let index_path = path.join("index.usearch");
138        self.index = Self::make_index(self.dimensions)?;
139        self.index
140            .load(index_path.to_str().unwrap())
141            .map_err(|e| StoreError::Init(e.to_string()))?;
142
143        Ok(())
144    }
145
146    fn len(&self) -> usize {
147        self.index.size()
148    }
149}
semtree_store/usearch.rs

semtree_store/
usearch.rs
