kevy_embedded/
ops_index.rs

1//! v2.5 — embedded secondary-index API (RFC LOCKED; server parity
2//! minus FIELDS hydration, which exists to save wire round-trips the
3//! embedded caller doesn't have — read fields with `hget` directly).
4//!
5//! Placement: each shard's `Inner` carries its slice of every index
6//! (index-follows-key, same as the server), maintained inside
7//! `commit_write` under the shard lock the write already holds — the
8//! synchronous-derivation guarantee costs no extra locking. Key
9//! extraction from the logged argv is EXACT (a precise multi-key
10//! table, not the feed filter's fail-open heuristic): a missed update
11//! would be index drift, which derived-by-construction forbids.
12//!
13//! Backfill: `idx_create` builds synchronously, shard by shard, each
14//! under its own write lock (the hook holds the same lock, so there is
15//! no race window per shard). No `Building` state embedded — create
16//! returns when the index serves.
17
18use 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
25/// One page of index hits plus the cursor to resume from.
26pub type IndexPage = (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>);
27
28/// Store-level index state: catalog + a version stamp the per-shard
29/// segment lists sync against.
30#[derive(Default)]
31pub(crate) struct IndexReg {
32    pub(crate) catalog: RwLock<(u64, Catalog)>,
33}
34
35/// Per-shard segment list, kept inside `Inner` (guarded by the shard
36/// lock).
37#[derive(Default)]
38pub(crate) struct ShardSegs {
39    pub(crate) version: u64,
40    pub(crate) segs: Vec<(IndexSpec, Segment)>,
41    /// v2.7: inverted segments for KIND text specs (parallel list —
42    /// a spec appears in exactly one of the lists).
43    pub(crate) text: Vec<(IndexSpec, kevy_text::TextSegment)>,
44    /// v2.8: HNSW graphs for KIND ann specs.
45    pub(crate) ann: Vec<(IndexSpec, kevy_vector::Hnsw)>,
46}
47
48fn new_graph(spec: &IndexSpec) -> kevy_vector::Hnsw {
49    let a = spec.ann.as_ref().expect("ann spec");
50    kevy_vector::Hnsw::new(
51        a.dim as usize,
52        kevy_vector::HnswParams {
53            m: a.m as usize,
54            ef_construction: a.ef as usize,
55            distance: match a.distance {
56                1 => kevy_vector::Distance::L2,
57                2 => kevy_vector::Distance::Ip,
58                _ => kevy_vector::Distance::Cosine,
59            },
60        },
61    )
62}
63
64const SIDECAR: &str = "index-catalog.meta";
65
66impl Store {
67    /// `IDX.CREATE` equivalent. Builds synchronously; errors on
68    /// duplicate name / cap / bad spec.
69    pub fn idx_create(
70        &self,
71        name: &[u8],
72        prefix: &[u8],
73        field: &[u8],
74        ty: ValType,
75        kind: IndexKind,
76    ) -> io::Result<()> {
77        if prefix.is_empty() {
78            return Err(io::Error::new(io::ErrorKind::InvalidInput, "empty prefix"));
79        }
80        let spec = IndexSpec {
81            name: name.to_vec(),
82            prefix: prefix.to_vec(),
83            field: field.to_vec(),
84            ty,
85            kind,
86            max_bytes: 0,
87            ann: None,
88        };
89        self.register_spec(spec)
90    }
91
92    fn register_spec(&self, spec: IndexSpec) -> io::Result<()> {
93        {
94            let mut g = self
95                .indexes
96                .catalog
97                .write()
98                .unwrap_or_else(std::sync::PoisonError::into_inner);
99            let (ver, cat) = &mut *g;
100            cat.create(spec)
101                .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
102            *ver += 1;
103        }
104        self.persist_index_sidecar();
105        // Build every shard's slice now (each under its own lock).
106        for shard in self.shards.iter() {
107            let mut g = lock_write(shard);
108            let inner = &mut *g;
109            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
110        }
111        Ok(())
112    }
113
114    /// `IDX.DROP` equivalent; `false` if absent.
115    /// v2.8: declare an ANN index (KIND ann, TYPE vector). `params.m`
116    /// / `params.ef` of 0 select the defaults (16 / 200).
117    pub fn idx_create_ann(
118        &self,
119        name: &[u8],
120        prefix: &[u8],
121        field: &[u8],
122        params: kevy_index::AnnSpec,
123    ) -> io::Result<()> {
124        if params.dim == 0 || params.distance > 2 {
125            return Err(io::Error::new(io::ErrorKind::InvalidInput, "bad ann parameters"));
126        }
127        let spec = IndexSpec {
128            name: name.to_vec(),
129            prefix: prefix.to_vec(),
130            field: field.to_vec(),
131            ty: ValType::Vector,
132            kind: IndexKind::Ann,
133            max_bytes: 0,
134            ann: Some(kevy_index::AnnSpec {
135                m: if params.m == 0 { 16 } else { params.m },
136                ef: if params.ef == 0 { 200 } else { params.ef },
137                ..params
138            }),
139        };
140        self.register_spec(spec)
141    }
142
143    pub fn idx_drop(&self, name: &[u8]) -> bool {
144        let hit = {
145            let mut g = self
146                .indexes
147                .catalog
148                .write()
149                .unwrap_or_else(std::sync::PoisonError::into_inner);
150            let (ver, cat) = &mut *g;
151            let hit = cat.drop_index(name);
152            if hit {
153                *ver += 1;
154            }
155            hit
156        };
157        if hit {
158            self.persist_index_sidecar();
159        }
160        hit
161    }
162
163    /// Range / EQ query with cursor pagination: merged across shards
164    /// in `(value, key)` order. `cursor = None` starts; the returned
165    /// cursor resumes exclusively.
166    pub fn idx_query(
167        &self,
168        name: &[u8],
169        min: &IndexValue,
170        max: &IndexValue,
171        cursor: Option<&Cursor>,
172        limit: usize,
173    ) -> io::Result<IndexPage> {
174        let limit = limit.clamp(1, 100_000);
175        let mut all: Vec<(IndexValue, Vec<u8>)> = Vec::new();
176        self.for_each_segment(name, |seg| {
177            let (hits, _) = seg.range(min, max, cursor, limit);
178            all.extend(hits.into_iter().map(|(k, v)| (v, k)));
179        })?;
180        all.sort();
181        all.truncate(limit);
182        let next = if all.len() == limit {
183            all.last().map(|(v, k)| Cursor { value: v.clone(), key: k.clone() })
184        } else {
185            None
186        };
187        Ok((all.into_iter().map(|(v, k)| (k, v)).collect(), next))
188    }
189
190    /// Count without materializing keys.
191    pub fn idx_count(&self, name: &[u8], min: &IndexValue, max: &IndexValue) -> io::Result<u64> {
192        let mut total = 0u64;
193        self.for_each_segment(name, |seg| total += seg.count(min, max))?;
194        Ok(total)
195    }
196
197    /// Summed segment stats (entries / bytes / coerce failures /
198    /// unique-fence duplicates).
199    pub fn idx_stats(&self, name: &[u8]) -> io::Result<SegmentStats> {
200        let mut sum = SegmentStats::default();
201        self.for_each_segment(name, |seg| {
202            let s = seg.stats();
203            sum.entries += s.entries;
204            sum.approx_bytes += s.approx_bytes;
205            sum.coerce_failures += s.coerce_failures;
206            sum.duplicates += s.duplicates;
207        })?;
208        Ok(sum)
209    }
210
211    /// Declared indexes (name, prefix, kind), declaration order.
212    pub fn idx_list(&self) -> Vec<(Vec<u8>, Vec<u8>, IndexKind)> {
213        let g = self
214            .indexes
215            .catalog
216            .read()
217            .unwrap_or_else(std::sync::PoisonError::into_inner);
218        g.1.iter()
219            .map(|(s, _)| (s.name.clone(), s.prefix.clone(), s.kind))
220            .collect()
221    }
222
223    /// v2.7 `MATCH` — BM25-ranked hits merged across shards
224    /// (shard-local statistics; see docs/text-search.md).
225    pub fn idx_match(
226        &self,
227        name: &[u8],
228        query: &[u8],
229        limit: usize,
230    ) -> io::Result<Vec<(Vec<u8>, f64)>> {
231        let limit = limit.clamp(1, 1000);
232        let mut all: Vec<kevy_text::TextMatch> = Vec::new();
233        let mut found = false;
234        for shard in self.shards.iter() {
235            let mut g = lock_write(shard);
236            let inner = &mut *g;
237            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
238            if let Some((_, ts)) = inner.idx_segs.text.iter().find(|(s, _)| s.name == name) {
239                found = true;
240                all.extend(ts.matches(query, limit));
241            }
242        }
243        if !found {
244            return Err(io::Error::new(io::ErrorKind::NotFound, "no such text index"));
245        }
246        all.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key)));
247        all.truncate(limit);
248        Ok(all.into_iter().map(|m| (m.key, m.score)).collect())
249    }
250
251    /// v2.8 `KNN` — nearest neighbors merged ascending across shards.
252    /// `ef` = query beam width (0 = engine default; recall knob).
253    pub fn idx_knn(
254        &self,
255        name: &[u8],
256        query: &[f32],
257        k: usize,
258        ef: usize,
259    ) -> io::Result<Vec<(Vec<u8>, f32)>> {
260        let k = k.clamp(1, 1000);
261        let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
262        let mut found = false;
263        for shard in self.shards.iter() {
264            let mut g = lock_write(shard);
265            let inner = &mut *g;
266            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
267            if let Some((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name == name) {
268                found = true;
269                all.extend(graph.knn(query, k, ef));
270            }
271        }
272        if !found {
273            return Err(io::Error::new(io::ErrorKind::NotFound, "no such vector index"));
274        }
275        all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
276        all.truncate(k);
277        Ok(all)
278    }
279
280    fn for_each_segment(
281        &self,
282        name: &[u8],
283        mut f: impl FnMut(&Segment),
284    ) -> io::Result<()> {
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((_, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name == name) {
291                found = true;
292                f(seg);
293            }
294        }
295        if found {
296            Ok(())
297        } else {
298            Err(io::Error::new(io::ErrorKind::NotFound, "no such index"))
299        }
300    }
301
302    fn persist_index_sidecar(&self) {
303        let Some(dir) = &self.config.data_dir else { return };
304        let g = self
305            .indexes
306            .catalog
307            .read()
308            .unwrap_or_else(std::sync::PoisonError::into_inner);
309        let tmp = dir.join("index-catalog.meta.tmp");
310        if std::fs::write(&tmp, g.1.to_sidecar()).is_ok() {
311            let _ = std::fs::rename(&tmp, dir.join(SIDECAR));
312        }
313    }
314
315    /// Boot half — load a persisted catalog (indexes rebuild lazily on
316    /// first touch via `sync_segs`).
317    pub(crate) fn idx_boot(&self) {
318        let Some(dir) = &self.config.data_dir else { return };
319        if let Ok(text) = std::fs::read_to_string(dir.join(SIDECAR))
320            && let Some(cat) = Catalog::from_sidecar(&text)
321            && !cat.is_empty()
322        {
323            let mut g = self
324                .indexes
325                .catalog
326                .write()
327                .unwrap_or_else(std::sync::PoisonError::into_inner);
328            *g = (g.0 + 1, cat);
329        }
330    }
331}
332
333/// Reconcile one shard's segment list with the catalog; new indexes
334/// backfill from this shard's live keys (we hold the shard's write
335/// lock — no concurrent writes can race the scan).
336pub(crate) fn sync_segs(
337    reg: &IndexReg,
338    shard_segs: &mut ShardSegs,
339    store: &mut kevy_store::Store,
340) {
341    let g = reg
342        .catalog
343        .read()
344        .unwrap_or_else(std::sync::PoisonError::into_inner);
345    let (ver, cat) = &*g;
346    if shard_segs.version == *ver {
347        return;
348    }
349    let mut next: Vec<(IndexSpec, Segment)> = Vec::new();
350    let mut next_text: Vec<(IndexSpec, kevy_text::TextSegment)> = Vec::new();
351    let mut next_ann: Vec<(IndexSpec, kevy_vector::Hnsw)> = Vec::new();
352    for (spec, _) in cat.iter() {
353        if spec.kind == kevy_index::IndexKind::Ann {
354            match shard_segs.ann.iter().position(|(s, _)| s == spec) {
355                Some(i) => next_ann.push(shard_segs.ann.swap_remove(i)),
356                None => {
357                    let mut g = new_graph(spec);
358                    let mut pat = spec.prefix.clone();
359                    pat.push(b'*');
360                    for key in store.collect_keys(Some(&pat), None) {
361                        apply_ann_key(store, spec, &mut g, &key);
362                    }
363                    next_ann.push((spec.clone(), g));
364                }
365            }
366            continue;
367        }
368        if spec.kind == kevy_index::IndexKind::Text {
369            match shard_segs.text.iter().position(|(s, _)| s == spec) {
370                Some(i) => next_text.push(shard_segs.text.swap_remove(i)),
371                None => {
372                    let mut ts = kevy_text::TextSegment::new();
373                    let mut pat = spec.prefix.clone();
374                    pat.push(b'*');
375                    for key in store.collect_keys(Some(&pat), None) {
376                        apply_text_key(store, spec, &mut ts, &key);
377                    }
378                    next_text.push((spec.clone(), ts));
379                }
380            }
381            continue;
382        }
383        match shard_segs.segs.iter().position(|(s, _)| s == spec) {
384            Some(i) => next.push(shard_segs.segs.swap_remove(i)),
385            None => {
386                let mut seg = Segment::new();
387                let mut pat = spec.prefix.clone();
388                pat.push(b'*');
389                for key in store.collect_keys(Some(&pat), None) {
390                    apply_key(store, spec, &mut seg, &key);
391                }
392                next.push((spec.clone(), seg));
393            }
394        }
395    }
396    shard_segs.segs = next;
397    shard_segs.text = next_text;
398    shard_segs.ann = next_ann;
399    shard_segs.version = *ver;
400}
401
402fn apply_ann_key(
403    store: &mut kevy_store::Store,
404    spec: &IndexSpec,
405    g: &mut kevy_vector::Hnsw,
406    key: &[u8],
407) {
408    let v = match store.hget(key, &spec.field) {
409        Ok(Some(raw)) => {
410            let raw = raw.to_vec();
411            kevy_vector::parse_vector(&raw, g.dim())
412        }
413        _ => None,
414    };
415    g.apply(key, v);
416}
417
418fn apply_text_key(
419    store: &mut kevy_store::Store,
420    spec: &IndexSpec,
421    ts: &mut kevy_text::TextSegment,
422    key: &[u8],
423) {
424    match store.hget(key, &spec.field) {
425        Ok(Some(raw)) => {
426            let raw = raw.to_vec();
427            ts.apply(key, Some(&raw));
428        }
429        _ => ts.apply(key, None),
430    }
431}
432
433/// The write-path hook body — called from `commit_write` with the
434/// logged argv, under the shard lock. Extracts the written key(s)
435/// EXACTLY and re-derives their index entries.
436pub(crate) fn on_commit(
437    reg: &IndexReg,
438    shard_segs: &mut ShardSegs,
439    store: &mut kevy_store::Store,
440    parts: &[&[u8]],
441) {
442    {
443        // Cheap gate: empty catalog = one read-lock + is_empty.
444        let g = reg
445            .catalog
446            .read()
447            .unwrap_or_else(std::sync::PoisonError::into_inner);
448        if g.1.is_empty() {
449            return;
450        }
451    }
452    sync_segs(reg, shard_segs, store);
453    let verb = parts.first().copied().unwrap_or(b"");
454    if verb.eq_ignore_ascii_case(b"FLUSHALL") || verb.eq_ignore_ascii_case(b"FLUSHDB") {
455        for (_, seg) in &mut shard_segs.segs {
456            *seg = Segment::new();
457        }
458        for (_, ts) in &mut shard_segs.text {
459            *ts = kevy_text::TextSegment::new();
460        }
461        for (spec, g) in &mut shard_segs.ann {
462            *g = new_graph(spec);
463        }
464        return;
465    }
466    each_written_key(verb, parts, |key| {
467        for (spec, seg) in &mut shard_segs.segs {
468            if key.starts_with(&spec.prefix) {
469                apply_key(store, spec, seg, key);
470            }
471        }
472        for (spec, ts) in &mut shard_segs.text {
473            if key.starts_with(&spec.prefix) {
474                apply_text_key(store, spec, ts, key);
475            }
476        }
477        for (spec, g) in &mut shard_segs.ann {
478            if key.starts_with(&spec.prefix) {
479                apply_ann_key(store, spec, g, key);
480            }
481        }
482    });
483}
484
485/// EXACT written-key walk per verb shape (the logged effect argv is a
486/// closed set — new effect verbs must be added here; the parity test
487/// in `store_tests_index.rs` pins the list).
488pub(crate) fn each_written_key_pub(verb: &[u8], parts: &[&[u8]], f: impl FnMut(&[u8])) {
489    each_written_key(verb, parts, f);
490}
491
492fn each_written_key(verb: &[u8], parts: &[&[u8]], mut f: impl FnMut(&[u8])) {
493    let up = |v: &[u8], t: &[u8]| v.eq_ignore_ascii_case(t);
494    if up(verb, b"DEL") || up(verb, b"UNLINK") {
495        for k in &parts[1..] {
496            f(k);
497        }
498    } else if up(verb, b"MSET") {
499        let mut i = 1;
500        while i + 1 < parts.len() {
501            f(parts[i]);
502            i += 2;
503        }
504    } else if up(verb, b"COPY") || up(verb, b"RENAME") || up(verb, b"RENAMENX") {
505        if let Some(k) = parts.get(1) {
506            f(k);
507        }
508        if let Some(k) = parts.get(2) {
509            f(k);
510        }
511    } else if let Some(k) = parts.get(1) {
512        f(k);
513    }
514}
515
516fn apply_key(store: &mut kevy_store::Store, spec: &IndexSpec, seg: &mut Segment, key: &[u8]) {
517    match store.hget(key, &spec.field) {
518        Ok(Some(raw)) => {
519            let raw = raw.to_vec();
520            match IndexValue::coerce(spec.ty, &raw) {
521                Some(v) => seg.apply(key, Some(v)),
522                None => seg.apply(key, None),
523            }
524        }
525        Ok(None) => {
526            if store.exists(&[key.to_vec()]) == 0 {
527                seg.remove(key);
528            } else {
529                seg.apply(key, None);
530            }
531        }
532        Err(_) => seg.remove(key),
533    }
534}
kevy_embedded/ops_index.rs

kevy_embedded/
ops_index.rs
