1use std::io;
19use std::sync::RwLock;
20
21use kevy_index::{Catalog, Cursor, IndexKind, IndexSpec, IndexValue, Segment, SegmentStats, ValType};
22
23use crate::store::{Store, lock_write};
24
25pub type IndexPage = (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>);
27
28#[derive(Default)]
31pub(crate) struct IndexReg {
32 pub(crate) catalog: RwLock<(u64, Catalog)>,
33}
34
35#[derive(Default)]
38pub(crate) struct ShardSegs {
39 pub(crate) version: u64,
40 pub(crate) segs: Vec<(IndexSpec, Segment)>,
41 pub(crate) text: Vec<(IndexSpec, kevy_text::TextSegment)>,
44 pub(crate) ann: Vec<(IndexSpec, kevy_vector::Hnsw)>,
46 pub(crate) agg: Vec<(IndexSpec, kevy_index::AggSegment)>,
48}
49
50fn new_graph(spec: &IndexSpec) -> kevy_vector::Hnsw {
51 let a = spec.ann.as_ref().expect("ann spec");
52 kevy_vector::Hnsw::new(
53 a.dim as usize,
54 kevy_vector::HnswParams {
55 m: a.m as usize,
56 ef_construction: a.ef as usize,
57 distance: match a.distance {
58 1 => kevy_vector::Distance::L2,
59 2 => kevy_vector::Distance::Ip,
60 _ => kevy_vector::Distance::Cosine,
61 },
62 },
63 )
64}
65
66const SIDECAR: &str = "index-catalog.meta";
67
68impl Store {
69 pub fn idx_create(
72 &self,
73 name: &[u8],
74 prefix: &[u8],
75 field: &[u8],
76 ty: ValType,
77 kind: IndexKind,
78 ) -> io::Result<()> {
79 if prefix.is_empty() {
80 return Err(io::Error::new(io::ErrorKind::InvalidInput, "empty prefix"));
81 }
82 let spec = IndexSpec {
83 name: name.to_vec(),
84 prefix: prefix.to_vec(),
85 field: field.to_vec(),
86 ty,
87 kind,
88 max_bytes: 0,
89 ann: None,
90 group_by: None,
91 };
92 self.register_spec(spec)
93 }
94
95 fn register_spec(&self, spec: IndexSpec) -> io::Result<()> {
96 {
97 let mut g = self
98 .indexes
99 .catalog
100 .write()
101 .unwrap_or_else(std::sync::PoisonError::into_inner);
102 let (ver, cat) = &mut *g;
103 cat.create(spec)
104 .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
105 *ver += 1;
106 }
107 self.persist_index_sidecar();
108 for shard in self.shards.iter() {
110 let mut g = lock_write(shard);
111 let inner = &mut *g;
112 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
113 }
114 Ok(())
115 }
116
117 pub fn idx_create_ann(
120 &self,
121 name: &[u8],
122 prefix: &[u8],
123 field: &[u8],
124 params: kevy_index::AnnSpec,
125 ) -> io::Result<()> {
126 if params.dim == 0 || params.distance > 2 {
127 return Err(io::Error::new(io::ErrorKind::InvalidInput, "bad ann parameters"));
128 }
129 let spec = IndexSpec {
130 name: name.to_vec(),
131 prefix: prefix.to_vec(),
132 field: field.to_vec(),
133 ty: ValType::Vector,
134 kind: IndexKind::Ann,
135 max_bytes: 0,
136 ann: Some(kevy_index::AnnSpec {
137 m: if params.m == 0 { 16 } else { params.m },
138 ef: if params.ef == 0 { 200 } else { params.ef },
139 ..params
140 }),
141 group_by: None,
142 };
143 self.register_spec(spec)
144 }
145
146 pub fn idx_drop(&self, name: &[u8]) -> bool {
149 let hit = {
150 let mut g = self
151 .indexes
152 .catalog
153 .write()
154 .unwrap_or_else(std::sync::PoisonError::into_inner);
155 let (ver, cat) = &mut *g;
156 let hit = cat.drop_index(name);
157 if hit {
158 *ver += 1;
159 }
160 hit
161 };
162 if hit {
163 self.persist_index_sidecar();
164 }
165 hit
166 }
167
168 pub fn idx_query(
172 &self,
173 name: &[u8],
174 min: &IndexValue,
175 max: &IndexValue,
176 cursor: Option<&Cursor>,
177 limit: usize,
178 ) -> io::Result<IndexPage> {
179 let limit = limit.clamp(1, 100_000);
180 let mut all: Vec<(IndexValue, Vec<u8>)> = Vec::new();
181 self.for_each_segment(name, |seg| {
182 let (hits, _) = seg.range(min, max, cursor, limit);
183 all.extend(hits.into_iter().map(|(k, v)| (v, k)));
184 })?;
185 all.sort();
186 all.truncate(limit);
187 let next = if all.len() == limit {
188 all.last().map(|(v, k)| Cursor { value: v.clone(), key: k.clone() })
189 } else {
190 None
191 };
192 Ok((all.into_iter().map(|(v, k)| (k, v)).collect(), next))
193 }
194
195 pub fn idx_count(&self, name: &[u8], min: &IndexValue, max: &IndexValue) -> io::Result<u64> {
197 let mut total = 0u64;
198 self.for_each_segment(name, |seg| total += seg.count(min, max))?;
199 Ok(total)
200 }
201
202 pub fn idx_stats(&self, name: &[u8]) -> io::Result<SegmentStats> {
205 let mut sum = SegmentStats::default();
206 self.for_each_segment(name, |seg| {
207 let s = seg.stats();
208 sum.entries += s.entries;
209 sum.approx_bytes += s.approx_bytes;
210 sum.coerce_failures += s.coerce_failures;
211 sum.duplicates += s.duplicates;
212 })?;
213 Ok(sum)
214 }
215
216 pub fn idx_list(&self) -> Vec<(Vec<u8>, Vec<u8>, IndexKind)> {
218 let g = self
219 .indexes
220 .catalog
221 .read()
222 .unwrap_or_else(std::sync::PoisonError::into_inner);
223 g.1.iter()
224 .map(|(s, _)| (s.name.clone(), s.prefix.clone(), s.kind))
225 .collect()
226 }
227
228 pub fn idx_match(
231 &self,
232 name: &[u8],
233 query: &[u8],
234 limit: usize,
235 ) -> io::Result<Vec<(Vec<u8>, f64)>> {
236 let limit = limit.clamp(1, 1000);
237 let mut all: Vec<kevy_text::TextMatch> = Vec::new();
238 let mut found = false;
239 for shard in self.shards.iter() {
240 let mut g = lock_write(shard);
241 let inner = &mut *g;
242 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
243 if let Some((_, ts)) = inner.idx_segs.text.iter().find(|(s, _)| s.name == name) {
244 found = true;
245 all.extend(ts.matches(query, limit));
246 }
247 }
248 if !found {
249 return Err(io::Error::new(io::ErrorKind::NotFound, "no such text index"));
250 }
251 all.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key)));
252 all.truncate(limit);
253 Ok(all.into_iter().map(|m| (m.key, m.score)).collect())
254 }
255
256 pub fn idx_create_agg(
259 &self,
260 name: &[u8],
261 prefix: &[u8],
262 field: &[u8],
263 ty: ValType,
264 group_by: &[u8],
265 ) -> io::Result<()> {
266 if !matches!(ty, ValType::I64 | ValType::F64) || group_by.is_empty() {
267 return Err(io::Error::new(io::ErrorKind::InvalidInput, "agg requires numeric type + group field"));
268 }
269 let spec = IndexSpec {
270 name: name.to_vec(),
271 prefix: prefix.to_vec(),
272 field: field.to_vec(),
273 ty,
274 kind: IndexKind::Agg,
275 max_bytes: 0,
276 ann: None,
277 group_by: Some(group_by.to_vec()),
278 };
279 self.register_spec(spec)
280 }
281
282 pub fn idx_group(&self, name: &[u8], group: &[u8]) -> io::Result<kevy_index::GroupStats> {
284 let mut merged = kevy_index::GroupStats { count: 0, sum: 0.0, min: None, max: None };
285 let mut found = false;
286 for shard in self.shards.iter() {
287 let mut g = lock_write(shard);
288 let inner = &mut *g;
289 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
290 if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
291 found = true;
292 kevy_index::merge_group(&mut merged, &a.group(group));
293 }
294 }
295 if !found {
296 return Err(io::Error::new(io::ErrorKind::NotFound, "no such aggregate index"));
297 }
298 Ok(merged)
299 }
300
301 pub fn idx_groups(
303 &self,
304 name: &[u8],
305 by: kevy_index::AggBy,
306 limit: usize,
307 ) -> io::Result<Vec<(Vec<u8>, kevy_index::GroupStats)>> {
308 let limit = limit.clamp(1, 1000);
309 let mut merged: std::collections::HashMap<Vec<u8>, kevy_index::GroupStats> =
312 std::collections::HashMap::new();
313 let mut found = false;
314 for shard in self.shards.iter() {
315 let mut g = lock_write(shard);
316 let inner = &mut *g;
317 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
318 if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
319 found = true;
320 for (gk, st) in a.all_groups() {
321 match merged.get_mut(&gk) {
322 Some(m) => kevy_index::merge_group(m, &st),
323 None => {
324 merged.insert(gk, st);
325 }
326 }
327 }
328 }
329 }
330 if !found {
331 return Err(io::Error::new(io::ErrorKind::NotFound, "no such aggregate index"));
332 }
333 let mut ranked: Vec<(Vec<u8>, kevy_index::GroupStats)> = merged.into_iter().collect();
334 kevy_index::sort_groups(&mut ranked, by);
335 ranked.truncate(limit);
336 Ok(ranked)
337 }
338
339 pub fn idx_knn(
342 &self,
343 name: &[u8],
344 query: &[f32],
345 k: usize,
346 ef: usize,
347 ) -> io::Result<Vec<(Vec<u8>, f32)>> {
348 let k = k.clamp(1, 1000);
349 let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
350 let mut found = false;
351 for shard in self.shards.iter() {
352 let mut g = lock_write(shard);
353 let inner = &mut *g;
354 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
355 if let Some((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name == name) {
356 found = true;
357 all.extend(graph.knn(query, k, ef));
358 }
359 }
360 if !found {
361 return Err(io::Error::new(io::ErrorKind::NotFound, "no such vector index"));
362 }
363 all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
364 all.truncate(k);
365 Ok(all)
366 }
367
368 fn for_each_segment(
369 &self,
370 name: &[u8],
371 mut f: impl FnMut(&Segment),
372 ) -> io::Result<()> {
373 let mut found = false;
374 for shard in self.shards.iter() {
375 let mut g = lock_write(shard);
376 let inner = &mut *g;
377 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
378 if let Some((_, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name == name) {
379 found = true;
380 f(seg);
381 }
382 }
383 if found {
384 Ok(())
385 } else {
386 Err(io::Error::new(io::ErrorKind::NotFound, "no such index"))
387 }
388 }
389
390 fn persist_index_sidecar(&self) {
391 let Some(dir) = &self.config.data_dir else { return };
392 let g = self
393 .indexes
394 .catalog
395 .read()
396 .unwrap_or_else(std::sync::PoisonError::into_inner);
397 let tmp = dir.join("index-catalog.meta.tmp");
398 if std::fs::write(&tmp, g.1.to_sidecar()).is_ok() {
399 let _ = std::fs::rename(&tmp, dir.join(SIDECAR));
400 }
401 }
402
403 pub(crate) fn idx_boot(&self) {
406 let Some(dir) = &self.config.data_dir else { return };
407 if let Ok(text) = std::fs::read_to_string(dir.join(SIDECAR))
408 && let Some(cat) = Catalog::from_sidecar(&text)
409 && !cat.is_empty()
410 {
411 let mut g = self
412 .indexes
413 .catalog
414 .write()
415 .unwrap_or_else(std::sync::PoisonError::into_inner);
416 *g = (g.0 + 1, cat);
417 }
418 }
419}
420
421pub(crate) fn sync_segs(
425 reg: &IndexReg,
426 shard_segs: &mut ShardSegs,
427 store: &mut kevy_store::Store,
428) {
429 let g = reg
430 .catalog
431 .read()
432 .unwrap_or_else(std::sync::PoisonError::into_inner);
433 let (ver, cat) = &*g;
434 if shard_segs.version == *ver {
435 return;
436 }
437 let mut next: Vec<(IndexSpec, Segment)> = Vec::new();
438 let mut next_text: Vec<(IndexSpec, kevy_text::TextSegment)> = Vec::new();
439 let mut next_ann: Vec<(IndexSpec, kevy_vector::Hnsw)> = Vec::new();
440 let mut next_agg: Vec<(IndexSpec, kevy_index::AggSegment)> = Vec::new();
441 for (spec, _) in cat.iter() {
442 if spec.kind == kevy_index::IndexKind::Agg {
443 match shard_segs.agg.iter().position(|(s, _)| s == spec) {
444 Some(i) => next_agg.push(shard_segs.agg.swap_remove(i)),
445 None => {
446 let mut a = kevy_index::AggSegment::new();
447 let mut pat = spec.prefix.clone();
448 pat.push(b'*');
449 for key in store.collect_keys(Some(&pat), None) {
450 apply_agg_key(store, spec, &mut a, &key);
451 }
452 next_agg.push((spec.clone(), a));
453 }
454 }
455 continue;
456 }
457 if spec.kind == kevy_index::IndexKind::Ann {
458 match shard_segs.ann.iter().position(|(s, _)| s == spec) {
459 Some(i) => next_ann.push(shard_segs.ann.swap_remove(i)),
460 None => {
461 let mut g = new_graph(spec);
462 let mut pat = spec.prefix.clone();
463 pat.push(b'*');
464 for key in store.collect_keys(Some(&pat), None) {
465 apply_ann_key(store, spec, &mut g, &key);
466 }
467 next_ann.push((spec.clone(), g));
468 }
469 }
470 continue;
471 }
472 if spec.kind == kevy_index::IndexKind::Text {
473 match shard_segs.text.iter().position(|(s, _)| s == spec) {
474 Some(i) => next_text.push(shard_segs.text.swap_remove(i)),
475 None => {
476 let mut ts = kevy_text::TextSegment::new();
477 let mut pat = spec.prefix.clone();
478 pat.push(b'*');
479 for key in store.collect_keys(Some(&pat), None) {
480 apply_text_key(store, spec, &mut ts, &key);
481 }
482 next_text.push((spec.clone(), ts));
483 }
484 }
485 continue;
486 }
487 match shard_segs.segs.iter().position(|(s, _)| s == spec) {
488 Some(i) => next.push(shard_segs.segs.swap_remove(i)),
489 None => {
490 let mut seg = Segment::new();
491 let mut pat = spec.prefix.clone();
492 pat.push(b'*');
493 for key in store.collect_keys(Some(&pat), None) {
494 apply_key(store, spec, &mut seg, &key);
495 }
496 next.push((spec.clone(), seg));
497 }
498 }
499 }
500 shard_segs.segs = next;
501 shard_segs.text = next_text;
502 shard_segs.ann = next_ann;
503 shard_segs.agg = next_agg;
504 shard_segs.version = *ver;
505}
506
507fn apply_agg_key(
508 store: &mut kevy_store::Store,
509 spec: &IndexSpec,
510 a: &mut kevy_index::AggSegment,
511 key: &[u8],
512) {
513 let group_field = spec.group_by.as_deref().unwrap_or_default();
514 let group = match store.hget(key, group_field) {
515 Ok(Some(g)) => Some(g.to_vec()),
516 _ => None,
517 };
518 let val = match store.hget(key, &spec.field) {
519 Ok(Some(raw)) => {
520 let raw = raw.to_vec();
521 kevy_index::IndexValue::coerce(spec.ty, &raw)
522 }
523 _ => None,
524 };
525 match (group, val) {
526 (Some(g), Some(v)) => a.apply(key, Some((g, v)), false),
527 _ => a.apply(key, None, store.exists(&[key.to_vec()]) > 0),
528 }
529}
530
531fn apply_ann_key(
532 store: &mut kevy_store::Store,
533 spec: &IndexSpec,
534 g: &mut kevy_vector::Hnsw,
535 key: &[u8],
536) {
537 let v = match store.hget(key, &spec.field) {
538 Ok(Some(raw)) => {
539 let raw = raw.to_vec();
540 kevy_vector::parse_vector(&raw, g.dim())
541 }
542 _ => None,
543 };
544 g.apply(key, v);
545}
546
547fn apply_text_key(
548 store: &mut kevy_store::Store,
549 spec: &IndexSpec,
550 ts: &mut kevy_text::TextSegment,
551 key: &[u8],
552) {
553 match store.hget(key, &spec.field) {
554 Ok(Some(raw)) => {
555 let raw = raw.to_vec();
556 ts.apply(key, Some(&raw));
557 }
558 _ => ts.apply(key, None),
559 }
560}
561
562pub(crate) fn on_commit(
566 reg: &IndexReg,
567 shard_segs: &mut ShardSegs,
568 store: &mut kevy_store::Store,
569 parts: &[&[u8]],
570) {
571 {
572 let g = reg
574 .catalog
575 .read()
576 .unwrap_or_else(std::sync::PoisonError::into_inner);
577 if g.1.is_empty() {
578 return;
579 }
580 }
581 sync_segs(reg, shard_segs, store);
582 let verb = parts.first().copied().unwrap_or(b"");
583 if verb.eq_ignore_ascii_case(b"FLUSHALL") || verb.eq_ignore_ascii_case(b"FLUSHDB") {
584 for (_, seg) in &mut shard_segs.segs {
585 *seg = Segment::new();
586 }
587 for (_, ts) in &mut shard_segs.text {
588 *ts = kevy_text::TextSegment::new();
589 }
590 for (spec, g) in &mut shard_segs.ann {
591 *g = new_graph(spec);
592 }
593 for (_, a) in &mut shard_segs.agg {
594 *a = kevy_index::AggSegment::new();
595 }
596 return;
597 }
598 each_written_key(verb, parts, |key| {
599 for (spec, seg) in &mut shard_segs.segs {
600 if key.starts_with(&spec.prefix) {
601 apply_key(store, spec, seg, key);
602 }
603 }
604 for (spec, ts) in &mut shard_segs.text {
605 if key.starts_with(&spec.prefix) {
606 apply_text_key(store, spec, ts, key);
607 }
608 }
609 for (spec, g) in &mut shard_segs.ann {
610 if key.starts_with(&spec.prefix) {
611 apply_ann_key(store, spec, g, key);
612 }
613 }
614 for (spec, a) in &mut shard_segs.agg {
615 if key.starts_with(&spec.prefix) {
616 apply_agg_key(store, spec, a, key);
617 }
618 }
619 });
620}
621
622pub(crate) fn each_written_key_pub(verb: &[u8], parts: &[&[u8]], f: impl FnMut(&[u8])) {
626 each_written_key(verb, parts, f);
627}
628
629fn each_written_key(verb: &[u8], parts: &[&[u8]], mut f: impl FnMut(&[u8])) {
630 let up = |v: &[u8], t: &[u8]| v.eq_ignore_ascii_case(t);
631 if up(verb, b"DEL") || up(verb, b"UNLINK") {
632 for k in &parts[1..] {
633 f(k);
634 }
635 } else if up(verb, b"MSET") {
636 let mut i = 1;
637 while i + 1 < parts.len() {
638 f(parts[i]);
639 i += 2;
640 }
641 } else if up(verb, b"COPY") || up(verb, b"RENAME") || up(verb, b"RENAMENX") {
642 if let Some(k) = parts.get(1) {
643 f(k);
644 }
645 if let Some(k) = parts.get(2) {
646 f(k);
647 }
648 } else if let Some(k) = parts.get(1) {
649 f(k);
650 }
651}
652
653fn apply_key(store: &mut kevy_store::Store, spec: &IndexSpec, seg: &mut Segment, key: &[u8]) {
654 match store.hget(key, &spec.field) {
655 Ok(Some(raw)) => {
656 let raw = raw.to_vec();
657 match IndexValue::coerce(spec.ty, &raw) {
658 Some(v) => seg.apply(key, Some(v)),
659 None => seg.apply(key, None),
660 }
661 }
662 Ok(None) => {
663 if store.exists(&[key.to_vec()]) == 0 {
664 seg.remove(key);
665 } else {
666 seg.apply(key, None);
667 }
668 }
669 Err(_) => seg.remove(key),
670 }
671}