kevy_embedded/
ops_index.rs

1//! Embedded secondary-index API (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 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
28/// One page of index hits plus the cursor to resume from.
29pub type IndexPage = (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>);
30
31/// One field's highlight: its name and the `(start, end)` match spans.
32#[cfg(feature = "text")]
33pub type FieldSpans = (Vec<u8>, Vec<(u32, u32)>);
34/// A highlighted MATCH hit: key, score, and per-field [`FieldSpans`].
35#[cfg(feature = "text")]
36pub type HighlightedHit = (Vec<u8>, f64, Vec<FieldSpans>);
37
38// `idx_match_with`, its clause options and the span-mapping helper live
39// in a child module to keep this file under the 500-LOC ceiling; the
40// text-index specifics — declaring a multi-field one, and gathering the
41// corpus statistics a global BM25 scores against — live in another.
42#[cfg(feature = "text")]
43#[path = "ops_index_highlight.rs"]
44pub(crate) mod highlight;
45
46// The clause-carrying scalar query (capacity arc G1) and the
47// [`ValueFilter`] predicate shape it shares with MATCH — independent of
48// the `text` feature: a range index filters fine without a tokenizer.
49#[path = "ops_index_claused.rs"]
50pub(crate) mod claused;
51
52#[cfg(feature = "text")]
53#[path = "ops_index_text.rs"]
54mod text;
55
56/// Sort merged `(value, key)` hits, cut to `limit`, and derive the
57/// resume cursor. Shared by `Store::idx_query` and the transaction twin
58/// on `AtomicAllShards`, which differ only in where the segments come
59/// from — the pagination has to agree exactly or a cursor taken inside
60/// a transaction would not resume outside one.
61pub(crate) fn merge_page(mut all: Vec<(IndexValue, Vec<u8>)>, limit: usize) -> IndexPage {
62    all.sort();
63    all.truncate(limit);
64    let next = if all.len() == limit {
65        all.last().map(|(v, k)| Cursor { value: v.clone(), key: k.clone() })
66    } else {
67        None
68    };
69    (all.into_iter().map(|(v, k)| (k, v)).collect(), next)
70}
71
72/// Store-level index state: catalog + a version stamp the per-shard
73/// segment lists sync against.
74#[derive(Default)]
75pub(crate) struct IndexReg {
76    pub(crate) catalog: RwLock<(u64, Catalog)>,
77}
78
79/// Per-shard segment list, kept inside `Inner` (guarded by the shard
80/// lock).
81#[derive(Default)]
82pub(crate) struct ShardSegs {
83    pub(crate) version: u64,
84    pub(crate) segs: Vec<(IndexSpec, Segment)>,
85    /// Inverted segments for KIND text specs (parallel list —
86    /// a spec appears in exactly one of the lists).
87    #[cfg(feature = "text")]
88    pub(crate) text: Vec<(IndexSpec, kevy_text::TextSegment)>,
89    /// HNSW graphs for KIND ann specs.
90    #[cfg(feature = "vector")]
91    pub(crate) ann: Vec<(IndexSpec, kevy_vector::Hnsw)>,
92    /// Aggregate segments for KIND agg specs.
93    pub(crate) agg: Vec<(IndexSpec, kevy_index::AggSegment)>,
94    /// `reserved_bytes` generation cache: set by every
95    /// segment-mutating chokepoint (`on_commit` applies, list
96    /// rebuilds, FLUSH resets) via [`ShardSegs::mark_stats_dirty`];
97    /// an idle tick reads the cached sum instead of walking every
98    /// segment's stats. Compiled with the tier backend only — on
99    /// targets without it (wasm) nothing reads the cache.
100    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
101    pub(crate) stats_dirty: bool,
102    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
103    pub(crate) reserved_cache: u64,
104}
105
106impl ShardSegs {
107    /// Invalidate the `reserved_bytes` cache — a no-op on targets
108    /// without the tier backend, so mutation chokepoints call it
109    /// unconditionally.
110    #[inline]
111    pub(crate) fn mark_stats_dirty(&mut self) {
112        #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
113        {
114            self.stats_dirty = true;
115        }
116    }
117}
118
119#[cfg(feature = "persist")]
120const SIDECAR: &str = "index-catalog.meta";
121
122impl Store {
123    /// `IDX.CREATE` equivalent. Builds synchronously; errors on
124    /// duplicate name / cap / bad spec.
125    pub fn idx_create(
126        &self,
127        name: &[u8],
128        prefix: &[u8],
129        field: &[u8],
130        ty: ValType,
131        kind: IndexKind,
132    ) -> KevyResult<()> {
133        if prefix.is_empty() {
134            return Err(KevyError::InvalidInput("empty prefix".into()));
135        }
136        #[cfg(not(feature = "text"))]
137        if kind == IndexKind::Text {
138            return Err(KevyError::Unsupported("text indexes need the `text` feature".into()));
139        }
140        #[cfg(not(feature = "vector"))]
141        if kind == IndexKind::Ann {
142            return Err(KevyError::Unsupported("vector indexes need the `vector` feature".into()));
143        }
144        let spec = IndexSpec {
145            name: name.to_vec(),
146            prefix: prefix.to_vec(),
147            fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
148            ty,
149            kind,
150            max_bytes: 0,
151            ann: None,
152            group_by: None,
153            with_positions: false,
154            values: Vec::new(),
155            composite: None,
156        };
157        self.register_spec(spec)
158    }
159
160    pub(crate) fn register_spec(&self, spec: IndexSpec) -> KevyResult<()> {
161        // Tiering floor refusal: body in
162        // `ops_index_sync::tier_floor_check` (500-LOC rule).
163        #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
164        crate::ops_index_sync::tier_floor_check(&self.shards)?;
165        {
166            let mut g = self
167                .indexes
168                .catalog
169                .write()
170                .unwrap_or_else(std::sync::PoisonError::into_inner);
171            let (ver, cat) = &mut *g;
172            cat.create(spec)
173                .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
174            *ver += 1;
175        }
176        self.persist_index_sidecar();
177        // Build every shard's slice now (each under its own lock).
178        for shard in self.shards.iter() {
179            let mut g = lock_write(shard);
180            let inner = &mut *g;
181            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
182        }
183        Ok(())
184    }
185
186    /// Declare an ANN index (KIND ann, TYPE vector). `params.m`
187    /// / `params.ef` of 0 select the defaults (16 / 200).
188    #[cfg(feature = "vector")]
189    pub fn idx_create_ann(
190        &self,
191        name: &[u8],
192        prefix: &[u8],
193        field: &[u8],
194        params: kevy_index::AnnSpec,
195    ) -> KevyResult<()> {
196        if params.dim == 0 || params.distance > 2 {
197            return Err(KevyError::InvalidInput("bad ann parameters".into()));
198        }
199        let spec = IndexSpec {
200            name: name.to_vec(),
201            prefix: prefix.to_vec(),
202            fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
203            ty: ValType::Vector,
204            kind: IndexKind::Ann,
205            max_bytes: 0,
206            ann: Some(kevy_index::AnnSpec {
207                m: if params.m == 0 { 16 } else { params.m },
208                ef: if params.ef == 0 { 200 } else { params.ef },
209                ..params
210            }),
211            group_by: None,
212            with_positions: false,
213            values: Vec::new(),
214            composite: None,
215        };
216        self.register_spec(spec)
217    }
218
219    /// `IDX.DROP` equivalent; `false` if absent. On a hit the catalog
220    /// sidecar is re-persisted so the drop survives restart.
221    pub fn idx_drop(&self, name: &[u8]) -> bool {
222        let hit = {
223            let mut g = self
224                .indexes
225                .catalog
226                .write()
227                .unwrap_or_else(std::sync::PoisonError::into_inner);
228            let (ver, cat) = &mut *g;
229            let hit = cat.drop_index(name);
230            if hit {
231                *ver += 1;
232            }
233            hit
234        };
235        if hit {
236            self.persist_index_sidecar();
237        }
238        hit
239    }
240
241    /// Range / EQ query with cursor pagination: merged across shards
242    /// in `(value, key)` order. `cursor = None` starts; the returned
243    /// cursor resumes exclusively.
244    pub fn idx_query(
245        &self,
246        name: &[u8],
247        min: &IndexValue,
248        max: &IndexValue,
249        cursor: Option<&Cursor>,
250        limit: usize,
251    ) -> KevyResult<IndexPage> {
252        let limit = limit.clamp(1, 100_000);
253        let mut all: Vec<(IndexValue, Vec<u8>)> = Vec::new();
254        self.for_each_segment(name, |seg| {
255            let (hits, _) = seg.range(min, max, cursor, limit);
256            all.extend(hits.into_iter().map(|(k, v)| (v, k)));
257        })?;
258        Ok(merge_page(all, limit))
259    }
260
261    /// Count without materializing keys.
262    pub fn idx_count(&self, name: &[u8], min: &IndexValue, max: &IndexValue) -> KevyResult<u64> {
263        let mut total = 0u64;
264        self.for_each_segment(name, |seg| total += seg.count(min, max))?;
265        Ok(total)
266    }
267
268    /// Summed segment stats (entries / bytes / coerce failures /
269    /// unique-fence duplicates).
270    pub fn idx_stats(&self, name: &[u8]) -> KevyResult<SegmentStats> {
271        let mut sum = SegmentStats::default();
272        self.for_each_segment(name, |seg| {
273            let s = seg.stats();
274            sum.entries += s.entries;
275            sum.approx_bytes += s.approx_bytes;
276            sum.coerce_failures += s.coerce_failures;
277            sum.duplicates += s.duplicates;
278        })?;
279        Ok(sum)
280    }
281
282    /// Declared indexes (name, prefix, kind), declaration order.
283    pub fn idx_list(&self) -> Vec<(Vec<u8>, Vec<u8>, IndexKind)> {
284        let g = self
285            .indexes
286            .catalog
287            .read()
288            .unwrap_or_else(std::sync::PoisonError::into_inner);
289        g.1.iter()
290            .map(|(s, _)| (s.name.clone(), s.prefix.clone(), s.kind))
291            .collect()
292    }
293
294    /// `MATCH` — BM25-ranked hits merged across shards, scored against
295    /// **global** corpus statistics so a hit's rank does not depend on
296    /// which shard it landed on (see docs/text-search.md).
297    ///
298    /// Two query-time passes: the first sums each shard's `n_docs`,
299    /// `total_len` and per-query-token `df` into one [`CorpusStats`]; the
300    /// second scores every shard against it. Only the query's tokens'
301    /// df is aggregated, not a whole-corpus table — the query narrows it.
302    #[cfg(feature = "text")]
303    pub fn idx_match(
304        &self,
305        name: &[u8],
306        query: &[u8],
307        limit: usize,
308    ) -> KevyResult<Vec<(Vec<u8>, f64)>> {
309        Ok(self
310            .idx_match_with(name, query, limit, crate::MatchOpts::default())?
311            .into_iter()
312            .map(|(key, score, _)| (key, score))
313            .collect())
314    }
315
316
317    /// Declare an aggregate index (KIND agg — write-time GROUP
318    /// BY). `ty` must be numeric.
319    pub fn idx_create_agg(
320        &self,
321        name: &[u8],
322        prefix: &[u8],
323        field: &[u8],
324        ty: ValType,
325        group_by: &[u8],
326    ) -> KevyResult<()> {
327        if !matches!(ty, ValType::I64 | ValType::F64) || group_by.is_empty() {
328            return Err(KevyError::InvalidInput("agg requires numeric type + group field".into()));
329        }
330        let spec = IndexSpec {
331            name: name.to_vec(),
332            prefix: prefix.to_vec(),
333            fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
334            ty,
335            kind: IndexKind::Agg,
336            max_bytes: 0,
337            ann: None,
338            group_by: Some(group_by.to_vec()),
339            with_positions: false,
340            values: Vec::new(),
341            composite: None,
342        };
343        self.register_spec(spec)
344    }
345
346    /// One group's merged stats across shards.
347    pub fn idx_group(&self, name: &[u8], group: &[u8]) -> KevyResult<kevy_index::GroupStats> {
348        let mut merged = kevy_index::GroupStats { count: 0, sum: 0.0, min: None, max: None };
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((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
355                found = true;
356                kevy_index::merge_group(&mut merged, &a.group(group));
357            }
358        }
359        if !found {
360            return Err(KevyError::NotFound("no such aggregate index".into()));
361        }
362        Ok(merged)
363    }
364
365    /// Top groups merged + ranked across shards.
366    pub fn idx_groups(
367        &self,
368        name: &[u8],
369        by: kevy_index::AggBy,
370        limit: usize,
371    ) -> KevyResult<Vec<(Vec<u8>, kevy_index::GroupStats)>> {
372        let limit = limit.clamp(1, 1000);
373        // HashMap merge (same O(rows×groups) trap the server reduce
374        // had — hashing keeps it linear).
375        let mut merged: std::collections::HashMap<Vec<u8>, kevy_index::GroupStats> =
376            std::collections::HashMap::new();
377        let mut found = false;
378        for shard in self.shards.iter() {
379            let mut g = lock_write(shard);
380            let inner = &mut *g;
381            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
382            if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
383                found = true;
384                for (gk, st) in a.all_groups() {
385                    match merged.get_mut(&gk) {
386                        Some(m) => kevy_index::merge_group(m, &st),
387                        None => {
388                            merged.insert(gk, st);
389                        }
390                    }
391                }
392            }
393        }
394        if !found {
395            return Err(KevyError::NotFound("no such aggregate index".into()));
396        }
397        let mut ranked: Vec<(Vec<u8>, kevy_index::GroupStats)> = merged.into_iter().collect();
398        kevy_index::sort_groups(&mut ranked, by);
399        ranked.truncate(limit);
400        Ok(ranked)
401    }
402
403    /// `KNN` — nearest neighbors merged ascending across shards.
404    /// `ef` = query beam width (0 = engine default; recall knob).
405    #[cfg(feature = "vector")]
406    pub fn idx_knn(
407        &self,
408        name: &[u8],
409        query: &[f32],
410        k: usize,
411        ef: usize,
412    ) -> KevyResult<Vec<(Vec<u8>, f32)>> {
413        let k = k.clamp(1, 1000);
414        let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
415        let mut found = false;
416        for shard in self.shards.iter() {
417            let mut g = lock_write(shard);
418            let inner = &mut *g;
419            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
420            if let Some((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name == name) {
421                found = true;
422                all.extend(graph.knn(query, k, ef));
423            }
424        }
425        if !found {
426            return Err(KevyError::NotFound("no such vector index".into()));
427        }
428        all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
429        all.truncate(k);
430        Ok(all)
431    }
432
433    /// Without `persist` there is no data dir — the catalog lives only
434    /// in memory, so the sidecar halves are no-ops.
435    #[cfg(not(feature = "persist"))]
436    fn persist_index_sidecar(&self) {}
437
438    #[cfg(not(feature = "persist"))]
439    pub(crate) fn idx_boot(&self) {}
440
441    fn for_each_segment(
442        &self,
443        name: &[u8],
444        mut f: impl FnMut(&Segment),
445    ) -> KevyResult<()> {
446        let mut found = false;
447        for shard in self.shards.iter() {
448            let mut g = lock_write(shard);
449            let inner = &mut *g;
450            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
451            if let Some((_, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name == name) {
452                found = true;
453                f(seg);
454            }
455        }
456        if found {
457            Ok(())
458        } else {
459            Err(KevyError::NotFound("no such index".into()))
460        }
461    }
462
463    #[cfg(feature = "persist")]
464    fn persist_index_sidecar(&self) {
465        let Some(dir) = &self.config.data_dir else { return };
466        let g = self
467            .indexes
468            .catalog
469            .read()
470            .unwrap_or_else(std::sync::PoisonError::into_inner);
471        let tmp = dir.join("index-catalog.meta.tmp");
472        if std::fs::write(&tmp, g.1.to_sidecar()).is_ok() {
473            let _ = std::fs::rename(&tmp, dir.join(SIDECAR));
474        }
475    }
476
477    /// Boot half — load a persisted catalog (indexes rebuild lazily on
478    /// first touch via `sync_segs`).
479    #[cfg(feature = "persist")]
480    pub(crate) fn idx_boot(&self) {
481        let Some(dir) = &self.config.data_dir else { return };
482        if let Ok(text) = std::fs::read_to_string(dir.join(SIDECAR))
483            && let Some(cat) = Catalog::from_sidecar(&text)
484            && !cat.is_empty()
485        {
486            let mut g = self
487                .indexes
488                .catalog
489                .write()
490                .unwrap_or_else(std::sync::PoisonError::into_inner);
491            *g = (g.0 + 1, cat);
492        }
493    }
494}
kevy_embedded/ops_index.rs

kevy_embedded/
ops_index.rs
