1use crate::analyzer::synonym::SynonymMap;
19use crate::backend::FtsBackend;
20use crate::index::writer::FtsIndex;
21
22const SYNONYM_GROUPS_COLLECTION: &str = "_synonym_groups";
24
25const INDEX_SUBKEY: &str = "_index";
27
28#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
30pub struct SynonymGroupRecord {
31 pub name: String,
32 pub terms: Vec<String>,
33 pub created_at: u64,
34}
35
36impl<B: FtsBackend> FtsIndex<B> {
37 pub fn put_synonym_group(
39 &self,
40 database_id: u64,
41 tid: u64,
42 record: &SynonymGroupRecord,
43 ) -> Result<(), B::Error> {
44 let bytes = sonic_rs::to_vec(record).unwrap_or_default();
46 self.backend.write_meta(
47 database_id,
48 tid,
49 SYNONYM_GROUPS_COLLECTION,
50 &record.name,
51 &bytes,
52 )?;
53
54 let mut names = self.read_name_index(database_id, tid)?;
56 if !names.contains(&record.name) {
57 names.push(record.name.clone());
58 self.write_name_index(database_id, tid, &names)?;
59 }
60 Ok(())
61 }
62
63 pub fn delete_synonym_group(
65 &self,
66 database_id: u64,
67 tid: u64,
68 name: &str,
69 ) -> Result<bool, B::Error> {
70 let existed = self
73 .backend
74 .read_meta(database_id, tid, SYNONYM_GROUPS_COLLECTION, name)?
75 .is_some_and(|b| !b.is_empty());
76
77 if existed {
78 self.backend
80 .write_meta(database_id, tid, SYNONYM_GROUPS_COLLECTION, name, &[])?;
81 let mut names = self.read_name_index(database_id, tid)?;
83 names.retain(|n| n != name);
84 self.write_name_index(database_id, tid, &names)?;
85 }
86 Ok(existed)
87 }
88
89 pub fn get_synonym_group(
91 &self,
92 database_id: u64,
93 tid: u64,
94 name: &str,
95 ) -> Result<Option<SynonymGroupRecord>, B::Error> {
96 match self
97 .backend
98 .read_meta(database_id, tid, SYNONYM_GROUPS_COLLECTION, name)?
99 {
100 None => Ok(None),
101 Some(bytes) if bytes.is_empty() => Ok(None),
102 Some(bytes) => Ok(sonic_rs::from_slice::<SynonymGroupRecord>(&bytes).ok()),
103 }
104 }
105
106 pub fn list_synonym_groups(
108 &self,
109 database_id: u64,
110 tid: u64,
111 ) -> Result<Vec<SynonymGroupRecord>, B::Error> {
112 let names = self.read_name_index(database_id, tid)?;
113 let mut groups = Vec::with_capacity(names.len());
114 for name in &names {
115 if let Some(rec) = self.get_synonym_group(database_id, tid, name)? {
116 groups.push(rec);
117 }
118 }
119 Ok(groups)
120 }
121
122 pub fn build_synonym_map_for_tenant(
129 &self,
130 _database_id: u64,
131 _tid: u64,
132 all_groups: &[SynonymGroupRecord],
133 ) -> SynonymMap {
134 let mut map = SynonymMap::new();
135 for group in all_groups {
136 if group.terms.len() < 2 {
137 continue;
138 }
139 let analyzed: Vec<Vec<String>> = group
142 .terms
143 .iter()
144 .map(|t| crate::analyzer::pipeline::analyze(t))
145 .collect();
146
147 for (i, my_tokens) in analyzed.iter().enumerate() {
148 let other_tokens: Vec<&str> = analyzed
149 .iter()
150 .enumerate()
151 .filter(|(j, _)| *j != i)
152 .flat_map(|(_, ts)| ts.iter().map(|s| s.as_str()))
153 .collect();
154 for my_token in my_tokens {
155 map.add(my_token, &other_tokens);
156 }
157 }
158 }
159 map
160 }
161
162 pub fn expand_query_with_synonyms(
167 &self,
168 database_id: u64,
169 tid: u64,
170 tokens: Vec<String>,
171 ) -> Result<Vec<String>, B::Error> {
172 let groups = self.list_synonym_groups(database_id, tid)?;
173 if groups.is_empty() {
174 return Ok(tokens);
175 }
176 let map = self.build_synonym_map_for_tenant(database_id, tid, &groups);
177 let expanded = map.expand(&tokens);
178 Ok(expanded)
179 }
180
181 fn read_name_index(&self, database_id: u64, tid: u64) -> Result<Vec<String>, B::Error> {
184 match self
185 .backend
186 .read_meta(database_id, tid, SYNONYM_GROUPS_COLLECTION, INDEX_SUBKEY)?
187 {
188 None => Ok(Vec::new()),
189 Some(bytes) if bytes.is_empty() => Ok(Vec::new()),
190 Some(bytes) => Ok(sonic_rs::from_slice::<Vec<String>>(&bytes).unwrap_or_default()),
191 }
192 }
193
194 fn write_name_index(
195 &self,
196 database_id: u64,
197 tid: u64,
198 names: &[String],
199 ) -> Result<(), B::Error> {
200 let bytes = sonic_rs::to_vec(names).unwrap_or_default();
201 self.backend.write_meta(
202 database_id,
203 tid,
204 SYNONYM_GROUPS_COLLECTION,
205 INDEX_SUBKEY,
206 &bytes,
207 )
208 }
209}
210
211#[cfg(test)]
212mod tests {
213 use crate::backend::memory::MemoryBackend;
214 use crate::index::writer::FtsIndex;
215
216 use super::SynonymGroupRecord;
217
218 const DB: u64 = 0;
219 const T: u64 = 1;
220
221 fn idx() -> FtsIndex<MemoryBackend> {
222 FtsIndex::new(MemoryBackend::new())
223 }
224
225 fn rec(name: &str, terms: &[&str]) -> SynonymGroupRecord {
226 SynonymGroupRecord {
227 name: name.to_string(),
228 terms: terms.iter().map(|s| s.to_string()).collect(),
229 created_at: 0,
230 }
231 }
232
233 #[test]
234 fn put_and_get() {
235 let i = idx();
236 i.put_synonym_group(DB, T, &rec("db_terms", &["database", "db", "datastore"]))
237 .unwrap();
238 let got = i.get_synonym_group(DB, T, "db_terms").unwrap().unwrap();
239 assert_eq!(got.name, "db_terms");
240 assert_eq!(got.terms.len(), 3);
241 }
242
243 #[test]
244 fn delete_removes() {
245 let i = idx();
246 i.put_synonym_group(DB, T, &rec("g1", &["a", "b"])).unwrap();
247 assert!(i.delete_synonym_group(DB, T, "g1").unwrap());
248 assert!(!i.delete_synonym_group(DB, T, "g1").unwrap());
249 assert!(i.get_synonym_group(DB, T, "g1").unwrap().is_none());
250 }
251
252 #[test]
253 fn list_reflects_puts_and_deletes() {
254 let i = idx();
255 i.put_synonym_group(DB, T, &rec("g1", &["a", "b"])).unwrap();
256 i.put_synonym_group(DB, T, &rec("g2", &["x", "y"])).unwrap();
257 let names: Vec<String> = i
258 .list_synonym_groups(DB, T)
259 .unwrap()
260 .into_iter()
261 .map(|r| r.name)
262 .collect();
263 assert_eq!(names.len(), 2);
264
265 i.delete_synonym_group(DB, T, "g1").unwrap();
266 let names2: Vec<String> = i
267 .list_synonym_groups(DB, T)
268 .unwrap()
269 .into_iter()
270 .map(|r| r.name)
271 .collect();
272 assert_eq!(names2, vec!["g2"]);
273 }
274
275 #[test]
276 fn synonym_map_bidirectional() {
277 let i = idx();
278 let recs = vec![rec("db_terms", &["db", "database", "datastore"])];
279 let map = i.build_synonym_map_for_tenant(DB, T, &recs);
280
281 let expanded = map.expand(&["db".to_string()]);
284 assert!(expanded.contains(&"databas".to_string()));
285 assert!(expanded.contains(&"datastor".to_string()));
286
287 let expanded2 = map.expand(&["databas".to_string()]);
288 assert!(expanded2.contains(&"db".to_string()));
289 assert!(expanded2.contains(&"datastor".to_string()));
290 }
291
292 #[test]
293 fn expand_query_with_synonyms_no_groups() {
294 let i = idx();
295 let tokens = vec!["hello".to_string(), "world".to_string()];
296 let expanded = i.expand_query_with_synonyms(DB, T, tokens.clone()).unwrap();
297 assert_eq!(expanded, tokens);
298 }
299
300 #[test]
301 fn expand_query_expands_matching_token() {
302 let i = idx();
303 i.put_synonym_group(DB, T, &rec("db_terms", &["db", "database", "datastore"]))
304 .unwrap();
305 let tokens = vec!["db".to_string(), "perform".to_string()];
309 let expanded = i.expand_query_with_synonyms(DB, T, tokens).unwrap();
310 assert!(expanded.contains(&"db".to_string()));
311 assert!(expanded.contains(&"databas".to_string()));
312 assert!(expanded.contains(&"datastor".to_string()));
313 assert!(expanded.contains(&"perform".to_string()));
314 }
315}