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