1use crate::{KevyError, KevyResult};
19use std::io;
20use std::sync::RwLock;
21
22use kevy_index::{
23 Catalog, Cursor, IndexKind, IndexSpec, IndexValue, Segment, SegmentStats, ValType,
24};
25
26use crate::store::{Store, lock_write};
27
28pub(crate) use crate::ops_index_sync::{each_written_key_pub, on_commit, sync_segs};
29
30pub type IndexPage = (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>);
32
33#[cfg(feature = "text")]
35pub type FieldSpans = (Vec<u8>, Vec<(u32, u32)>);
36#[cfg(feature = "text")]
38pub type HighlightedHit = (Vec<u8>, f64, Vec<FieldSpans>);
39
40#[cfg(feature = "text")]
45#[path = "ops_index_highlight.rs"]
46pub(crate) mod highlight;
47
48#[path = "ops_index_claused.rs"]
52pub(crate) mod claused;
53
54#[path = "ops_index_advise.rs"]
56pub(crate) mod advise;
57
58#[path = "ops_index_admin.rs"]
60mod admin;
61
62#[cfg(feature = "text")]
63#[path = "ops_index_text.rs"]
64mod text;
65
66#[cfg(feature = "text")]
68#[path = "ops_index_text_cold.rs"]
69pub(crate) mod text_cold;
70
71pub(crate) fn merge_page(mut all: Vec<(IndexValue, Vec<u8>)>, limit: usize) -> IndexPage {
77 all.sort();
78 all.truncate(limit);
79 let next = if all.len() == limit {
80 all.last().map(|(v, k)| Cursor { value: v.clone(), key: k.clone() })
81 } else {
82 None
83 };
84 (all.into_iter().map(|(v, k)| (k, v)).collect(), next)
85}
86
87#[derive(Default)]
91pub(crate) struct IndexReg {
92 pub(crate) catalog: RwLock<(u64, Catalog)>,
93 pub(crate) usage:
94 RwLock<std::collections::HashMap<Vec<u8>, std::sync::Arc<kevy_index::UsageCell>>>,
95}
96
97#[cfg(not(target_arch = "wasm32"))]
100pub(crate) type WinRef<'a> = Option<&'a kevy_window::WindowRt>;
101#[cfg(target_arch = "wasm32")]
102pub(crate) type WinRef<'a> = Option<&'a core::convert::Infallible>;
103
104#[derive(Default)]
107pub(crate) struct ShardSegs {
108 pub(crate) version: u64,
109 pub(crate) segs: Vec<(IndexSpec, Segment)>,
110 #[cfg(feature = "text")]
113 pub(crate) text: Vec<(IndexSpec, kevy_text::TextSegment)>,
114 #[cfg(feature = "vector")]
116 pub(crate) ann: Vec<(IndexSpec, kevy_vector::Hnsw)>,
117 pub(crate) agg: Vec<(IndexSpec, kevy_index::AggSegment)>,
119 #[cfg(not(target_arch = "wasm32"))]
123 pub(crate) windows: Vec<(Vec<u8>, kevy_window::WindowRt)>,
124 #[cfg(all(feature = "text", not(target_arch = "wasm32")))]
128 pub(crate) cold_text: Vec<(Vec<u8>, kevy_window::TextColdDir)>,
129 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
136 pub(crate) stats_dirty: bool,
137 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
138 pub(crate) reserved_cache: u64,
139}
140
141impl ShardSegs {
142 #[cfg(not(target_arch = "wasm32"))]
147 pub(crate) fn window_of(&self, name: &[u8]) -> Option<&kevy_window::WindowRt> {
148 self.windows.iter().find(|(n, _)| n == name).map(|(_, w)| w)
149 }
150
151 #[inline]
155 pub(crate) fn mark_stats_dirty(&mut self) {
156 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
157 {
158 self.stats_dirty = true;
159 }
160 }
161}
162
163#[cfg(feature = "persist")]
164const SIDECAR: &str = "index-catalog.meta";
165
166impl Store {
167 pub fn idx_create(
170 &self,
171 name: &[u8],
172 prefix: &[u8],
173 field: &[u8],
174 ty: ValType,
175 kind: IndexKind,
176 ) -> KevyResult<()> {
177 if prefix.is_empty() {
178 return Err(KevyError::InvalidInput("empty prefix".into()));
179 }
180 #[cfg(not(feature = "text"))]
181 if kind == IndexKind::Text {
182 return Err(KevyError::Unsupported("text indexes need the `text` feature".into()));
183 }
184 #[cfg(not(feature = "vector"))]
185 if kind == IndexKind::Ann {
186 return Err(KevyError::Unsupported("vector indexes need the `vector` feature".into()));
187 }
188 let spec = IndexSpec {
189 name: name.to_vec(),
190 prefix: prefix.to_vec(),
191 fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
192 ty,
193 kind,
194 max_bytes: 0,
195 ann: None,
196 group_by: None,
197 with_positions: false,
198 values: Vec::new(),
199 composite: None,
200 };
201 self.register_spec(spec)
202 }
203
204 pub(crate) fn register_spec(&self, spec: IndexSpec) -> KevyResult<()> {
205 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
208 crate::ops_index_sync::tier_floor_check(&self.shards)?;
209 {
210 let mut g =
211 self.indexes.catalog.write().unwrap_or_else(std::sync::PoisonError::into_inner);
212 let (ver, cat) = &mut *g;
213 cat.create(spec).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
214 *ver += 1;
215 }
216 self.persist_index_sidecar();
217 self.advise_clear();
218 self.usage_rekey();
219 for shard in self.shards.iter() {
221 let mut g = lock_write(shard);
222 let inner = &mut *g;
223 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
224 }
225 Ok(())
226 }
227
228 #[cfg(feature = "vector")]
231 pub fn idx_create_ann(
232 &self,
233 name: &[u8],
234 prefix: &[u8],
235 field: &[u8],
236 params: kevy_index::AnnSpec,
237 ) -> KevyResult<()> {
238 if params.dim == 0 || params.distance > 2 {
239 return Err(KevyError::InvalidInput("bad ann parameters".into()));
240 }
241 let spec = IndexSpec {
242 name: name.to_vec(),
243 prefix: prefix.to_vec(),
244 fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
245 ty: ValType::Vector,
246 kind: IndexKind::Ann,
247 max_bytes: 0,
248 ann: Some(kevy_index::AnnSpec {
249 m: if params.m == 0 { 16 } else { params.m },
250 ef: if params.ef == 0 { 200 } else { params.ef },
251 ..params
252 }),
253 group_by: None,
254 with_positions: false,
255 values: Vec::new(),
256 composite: None,
257 };
258 self.register_spec(spec)
259 }
260
261 pub fn idx_drop(&self, name: &[u8]) -> bool {
264 let hit = {
265 let mut g =
266 self.indexes.catalog.write().unwrap_or_else(std::sync::PoisonError::into_inner);
267 let (ver, cat) = &mut *g;
268 let hit = cat.drop_index(name);
269 if hit {
270 *ver += 1;
271 }
272 hit
273 };
274 if hit {
275 self.persist_index_sidecar();
276 self.advise_clear();
277 self.usage_rekey();
278 }
279 hit
280 }
281
282 #[cfg(feature = "text")]
291 pub fn idx_match(
292 &self,
293 name: &[u8],
294 query: &[u8],
295 limit: usize,
296 ) -> KevyResult<Vec<(Vec<u8>, f64)>> {
297 Ok(self
298 .idx_match_with(name, query, limit, crate::MatchOpts::default())?
299 .into_iter()
300 .map(|(key, score, _)| (key, score))
301 .collect())
302 }
303
304 pub fn idx_create_agg(
307 &self,
308 name: &[u8],
309 prefix: &[u8],
310 field: &[u8],
311 ty: ValType,
312 group_by: &[u8],
313 ) -> KevyResult<()> {
314 if !matches!(ty, ValType::I64 | ValType::F64) || group_by.is_empty() {
315 return Err(KevyError::InvalidInput("agg requires numeric type + group field".into()));
316 }
317 let spec = IndexSpec {
318 name: name.to_vec(),
319 prefix: prefix.to_vec(),
320 fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
321 ty,
322 kind: IndexKind::Agg,
323 max_bytes: 0,
324 ann: None,
325 group_by: Some(group_by.to_vec()),
326 with_positions: false,
327 values: Vec::new(),
328 composite: None,
329 };
330 self.register_spec(spec)
331 }
332
333 pub fn idx_group(&self, name: &[u8], group: &[u8]) -> KevyResult<kevy_index::GroupStats> {
335 let mut merged = kevy_index::GroupStats { count: 0, sum: 0.0, min: None, max: None };
336 let mut found = false;
337 for shard in self.shards.iter() {
338 let mut g = lock_write(shard);
339 let inner = &mut *g;
340 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
341 if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
342 found = true;
343 kevy_index::merge_group(&mut merged, &a.group(group));
344 }
345 }
346 if !found {
347 return Err(KevyError::NotFound("no such aggregate index".into()));
348 }
349 Ok(merged)
350 }
351
352 pub fn idx_groups(
354 &self,
355 name: &[u8],
356 by: kevy_index::AggBy,
357 limit: usize,
358 ) -> KevyResult<Vec<(Vec<u8>, kevy_index::GroupStats)>> {
359 let limit = limit.clamp(1, 1000);
360 let mut merged: std::collections::HashMap<Vec<u8>, kevy_index::GroupStats> =
363 std::collections::HashMap::new();
364 let mut found = false;
365 for shard in self.shards.iter() {
366 let mut g = lock_write(shard);
367 let inner = &mut *g;
368 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
369 if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
370 found = true;
371 for (gk, st) in a.all_groups() {
372 match merged.get_mut(&gk) {
373 Some(m) => kevy_index::merge_group(m, &st),
374 None => {
375 merged.insert(gk, st);
376 }
377 }
378 }
379 }
380 }
381 if !found {
382 return Err(KevyError::NotFound("no such aggregate index".into()));
383 }
384 let mut ranked: Vec<(Vec<u8>, kevy_index::GroupStats)> = merged.into_iter().collect();
385 kevy_index::sort_groups(&mut ranked, by);
386 ranked.truncate(limit);
387 Ok(ranked)
388 }
389
390 #[cfg(feature = "vector")]
393 pub fn idx_knn(
394 &self,
395 name: &[u8],
396 query: &[f32],
397 k: usize,
398 ef: usize,
399 ) -> KevyResult<Vec<(Vec<u8>, f32)>> {
400 let k = k.clamp(1, 1000);
401 let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
402 let mut found = false;
403 for shard in self.shards.iter() {
404 let mut g = lock_write(shard);
405 let inner = &mut *g;
406 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
407 if let Some((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name == name) {
408 found = true;
409 all.extend(graph.knn(query, k, ef));
410 }
411 }
412 if !found {
413 return Err(KevyError::NotFound("no such vector index".into()));
414 }
415 all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
416 all.truncate(k);
417 Ok(all)
418 }
419
420 #[cfg(not(feature = "persist"))]
423 fn persist_index_sidecar(&self) {}
424
425 #[cfg(not(feature = "persist"))]
426 pub(crate) fn idx_boot(&self) {}
427
428 fn for_each_segment(&self, name: &[u8], mut f: impl FnMut(&Segment)) -> KevyResult<()> {
429 let mut found = false;
430 for shard in self.shards.iter() {
431 let mut g = lock_write(shard);
432 let inner = &mut *g;
433 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
434 if let Some((_, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name == name) {
435 found = true;
436 f(seg);
437 }
438 }
439 if found { Ok(()) } else { Err(KevyError::NotFound("no such index".into())) }
440 }
441
442 #[cfg(feature = "persist")]
443 fn persist_index_sidecar(&self) {
444 let Some(dir) = &self.config.data_dir else { return };
445 let g = self.indexes.catalog.read().unwrap_or_else(std::sync::PoisonError::into_inner);
446 let tmp = dir.join("index-catalog.meta.tmp");
447 if std::fs::write(&tmp, g.1.to_sidecar()).is_ok() {
448 let _ = std::fs::rename(&tmp, dir.join(SIDECAR));
449 }
450 }
451
452 #[cfg(feature = "persist")]
455 pub(crate) fn idx_boot(&self) {
456 let Some(dir) = &self.config.data_dir else { return };
457 if let Ok(text) = std::fs::read_to_string(dir.join(SIDECAR))
458 && let Some(cat) = Catalog::from_sidecar(&text)
459 && !cat.is_empty()
460 {
461 let mut g =
462 self.indexes.catalog.write().unwrap_or_else(std::sync::PoisonError::into_inner);
463 *g = (g.0 + 1, cat);
464 }
465 }
466}