1use super::hnsw::HNSW;
2use super::kv::FjallStore;
3use fjall::{KeyspaceCreateOptions, OptimisticTxDatabase, OptimisticTxKeyspace};
4
5pub struct Query {
6 pub ids: OptimisticTxKeyspace, pub names: OptimisticTxKeyspace, pub descriptions: OptimisticTxKeyspace, pub scales: OptimisticTxKeyspace, pub aabbs: OptimisticTxKeyspace, pub hnsw: HNSW<FjallStore>,
12}
13
14use anyhow::Result;
15use indexmap::IndexMap;
16use smol_str::SmolStr;
17
18use dynamic::{Dynamic, map};
19
20#[derive(Debug, Clone, PartialEq)]
21pub struct SearchDescriptionResult {
22 pub name: SmolStr,
23 pub description: SmolStr,
24 pub scale: SmolStr,
25 pub aabb: SmolStr,
26 pub distance: f32,
27 pub similarity: f32,
28}
29
30fn similarity_from_l2_distance(distance: f32) -> f32 {
31 1.0 / (1.0 + distance)
32}
33
34fn description_results_to_index_map(results: Vec<SearchDescriptionResult>) -> IndexMap<SmolStr, Dynamic> {
35 let mut map = IndexMap::with_capacity(results.len());
36 for item in results {
37 map.insert(
38 item.name,
39 map!(
40 "description" => item.description,
41 "scale" => item.scale,
42 "aabb" => item.aabb,
43 "distance" => item.distance,
44 "similarity" => item.similarity
45 ),
46 );
47 }
48 map
49}
50
51impl Query {
52 pub fn new(name: &str) -> Result<Self> {
53 let db = OptimisticTxDatabase::builder(name).open()?;
54 let ids = db.keyspace(&format!("{}_ids", name), KeyspaceCreateOptions::default)?;
55 let names = db.keyspace(&format!("{}_names", name), KeyspaceCreateOptions::default)?;
56 let descriptions = db.keyspace(&format!("{}_descriptions", name), KeyspaceCreateOptions::default)?;
57 let scales = db.keyspace(&format!("{}_scales", name), KeyspaceCreateOptions::default)?;
58 let aabbs = db.keyspace(&format!("{}_aabbs", name), KeyspaceCreateOptions::default)?;
59 let hnsw = HNSW::<FjallStore>::new(FjallStore::open(&db, &format!("{}_hnsw", name))?, 20, 200, 16, super::Dist::L2);
60 Ok(Query { ids, names, descriptions, scales, aabbs, hnsw })
61 }
62
63 pub fn rename(&self, name: &str, new_name: &str) -> Result<()> {
64 if let Ok(Some(old)) = self.names.get(name) {
65 let id = u64::from_le_bytes(old.as_ref().try_into()?);
66 self.ids.insert(id.to_le_bytes(), new_name)?;
67 self.names.insert(new_name.to_string(), id.to_le_bytes())?;
68 if let Some(description) = self.descriptions.get(name)? {
69 self.descriptions.insert(new_name.to_string(), description)?;
70 self.descriptions.remove(name)?;
71 }
72 if let Some(scale) = self.scales.get(name)? {
73 self.scales.insert(new_name.to_string(), scale)?;
74 self.scales.remove(name)?;
75 }
76 if let Some(aabb) = self.aabbs.get(name)? {
77 self.aabbs.insert(new_name.to_string(), aabb)?;
78 self.aabbs.remove(name)?;
79 }
80 self.names.remove(name)?;
81 }
82 Ok(())
83 }
84
85 pub fn add(&self, name: &str, arrow: Vec<f32>) -> Result<u64> {
86 self.add_with_description(name, arrow, None)
87 }
88
89 pub fn add_with_description(&self, name: &str, arrow: Vec<f32>, description: Option<&str>) -> Result<u64> {
90 self.add_with_metadata(name, arrow, description, None, None)
91 }
92
93 pub fn add_with_metadata(&self, name: &str, arrow: Vec<f32>, description: Option<&str>, scale: Option<&str>, aabb: Option<&str>) -> Result<u64> {
94 if let Ok(Some(old)) = self.names.get(name) {
95 let id = u64::from_le_bytes(old.as_ref().try_into()?);
96 self.hnsw.set_arrow(id, arrow)?;
97 self.update_metadata(name, description, scale, aabb)?;
98 Ok(id)
99 } else {
100 let id = self.hnsw.insert(arrow)?;
101 self.ids.insert(id.to_le_bytes(), name)?;
102 self.names.insert(name.to_string(), id.to_le_bytes())?;
103 self.update_metadata(name, description, scale, aabb)?;
104 Ok(id)
105 }
106 }
107
108 pub fn get_description(&self, name: &str) -> Result<Option<SmolStr>> {
109 self.get_text(&self.descriptions, name)
110 }
111
112 pub fn get_scale(&self, name: &str) -> Result<Option<SmolStr>> {
113 self.get_text(&self.scales, name)
114 }
115
116 pub fn get_aabb(&self, name: &str) -> Result<Option<SmolStr>> {
117 self.get_text(&self.aabbs, name)
118 }
119
120 pub fn search(&self, arrow: Vec<f32>, number: usize) -> Result<Dynamic> {
121 let ids = self.hnsw.search(arrow, number)?;
122 let mut results = Vec::new();
123 for (id, _) in ids {
124 let name = self.ids.get(id.to_le_bytes())?.unwrap();
125 let s = SmolStr::new(std::str::from_utf8(&name)?);
126 results.push(s.into());
127 }
128 log::info!("hnsw names {:?}", results);
129 Ok(Dynamic::list(results))
130 }
131
132 pub fn search_description_map(&self, arrow: Vec<f32>, number: usize) -> Result<IndexMap<SmolStr, Dynamic>> {
133 let matches = self.search_description_results(arrow, number)?;
134 Ok(description_results_to_index_map(matches))
135 }
136
137 pub fn search_description_results(&self, arrow: Vec<f32>, number: usize) -> Result<Vec<SearchDescriptionResult>> {
138 let ids = self.hnsw.search(arrow, number)?;
139 let mut results = Vec::with_capacity(ids.len());
140 for (id, distance) in ids {
141 let name = self.ids.get(id.to_le_bytes())?.unwrap();
142 let name = SmolStr::new(std::str::from_utf8(&name)?);
143 let description = self.get_description(name.as_str())?.unwrap_or_default();
144 let scale = self.get_scale(name.as_str())?.unwrap_or_default();
145 let aabb = self.get_aabb(name.as_str())?.unwrap_or_default();
146 results.push(SearchDescriptionResult { name, description, scale, aabb, distance, similarity: similarity_from_l2_distance(distance) });
147 }
148 Ok(results)
149 }
150
151 fn get_text(&self, keyspace: &OptimisticTxKeyspace, name: &str) -> Result<Option<SmolStr>> {
152 let value = keyspace.get(name)?;
153 value.map(|buf| std::str::from_utf8(&buf).map(SmolStr::new)).transpose().map_err(Into::into)
154 }
155
156 fn update_metadata(&self, name: &str, description: Option<&str>, scale: Option<&str>, aabb: Option<&str>) -> Result<()> {
157 Self::set_optional_text(&self.descriptions, name, description)?;
158 Self::set_optional_text(&self.scales, name, scale)?;
159 Self::set_optional_text(&self.aabbs, name, aabb)?;
160 Ok(())
161 }
162
163 fn set_optional_text(keyspace: &OptimisticTxKeyspace, name: &str, value: Option<&str>) -> Result<()> {
164 if let Some(value) = value {
165 keyspace.insert(name.to_string(), value)?;
166 } else {
167 let _ = keyspace.remove(name)?;
168 }
169 Ok(())
170 }
171}
172
173#[cfg(test)]
174mod tests {
175 use super::*;
176
177 #[test]
178 fn description_results_to_index_map_preserves_order_and_similarity() {
179 let rows = vec![
180 SearchDescriptionResult { name: "first".into(), description: "first desc".into(), scale: "".into(), aabb: "".into(), distance: 0.25, similarity: 0.8 },
181 SearchDescriptionResult { name: "second".into(), description: "second desc".into(), scale: "".into(), aabb: "".into(), distance: 1.0, similarity: 0.5 },
182 ];
183
184 let results = description_results_to_index_map(rows);
185 let keys = results.keys().map(|key| key.as_str()).collect::<Vec<_>>();
186 let similarity = results.get("first").and_then(|item| item.get_dynamic("similarity")).and_then(|value| value.as_float()).unwrap();
187
188 assert_eq!(keys, vec!["first", "second"]);
189 assert!((similarity - 0.8).abs() < 0.0001);
190 }
191}