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