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}
95
96#[cfg(feature = "persist")]
97const SIDECAR: &str = "index-catalog.meta";
98
99impl Store {
100    /// `IDX.CREATE` equivalent. Builds synchronously; errors on
101    /// duplicate name / cap / bad spec.
102    pub fn idx_create(
103        &self,
104        name: &[u8],
105        prefix: &[u8],
106        field: &[u8],
107        ty: ValType,
108        kind: IndexKind,
109    ) -> KevyResult<()> {
110        if prefix.is_empty() {
111            return Err(KevyError::InvalidInput("empty prefix".into()));
112        }
113        #[cfg(not(feature = "text"))]
114        if kind == IndexKind::Text {
115            return Err(KevyError::Unsupported("text indexes need the `text` feature".into()));
116        }
117        #[cfg(not(feature = "vector"))]
118        if kind == IndexKind::Ann {
119            return Err(KevyError::Unsupported("vector indexes need the `vector` feature".into()));
120        }
121        let spec = IndexSpec {
122            name: name.to_vec(),
123            prefix: prefix.to_vec(),
124            fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
125            ty,
126            kind,
127            max_bytes: 0,
128            ann: None,
129            group_by: None,
130            with_positions: false,
131            values: Vec::new(),
132            composite: None,
133        };
134        self.register_spec(spec)
135    }
136
137    pub(crate) fn register_spec(&self, spec: IndexSpec) -> KevyResult<()> {
138        // Tiering floor refusal: body in
139        // `ops_index_sync::tier_floor_check` (500-LOC rule).
140        #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
141        crate::ops_index_sync::tier_floor_check(&self.shards)?;
142        {
143            let mut g = self
144                .indexes
145                .catalog
146                .write()
147                .unwrap_or_else(std::sync::PoisonError::into_inner);
148            let (ver, cat) = &mut *g;
149            cat.create(spec)
150                .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
151            *ver += 1;
152        }
153        self.persist_index_sidecar();
154        // Build every shard's slice now (each under its own lock).
155        for shard in self.shards.iter() {
156            let mut g = lock_write(shard);
157            let inner = &mut *g;
158            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
159        }
160        Ok(())
161    }
162
163    /// Declare an ANN index (KIND ann, TYPE vector). `params.m`
164    /// / `params.ef` of 0 select the defaults (16 / 200).
165    #[cfg(feature = "vector")]
166    pub fn idx_create_ann(
167        &self,
168        name: &[u8],
169        prefix: &[u8],
170        field: &[u8],
171        params: kevy_index::AnnSpec,
172    ) -> KevyResult<()> {
173        if params.dim == 0 || params.distance > 2 {
174            return Err(KevyError::InvalidInput("bad ann parameters".into()));
175        }
176        let spec = IndexSpec {
177            name: name.to_vec(),
178            prefix: prefix.to_vec(),
179            fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
180            ty: ValType::Vector,
181            kind: IndexKind::Ann,
182            max_bytes: 0,
183            ann: Some(kevy_index::AnnSpec {
184                m: if params.m == 0 { 16 } else { params.m },
185                ef: if params.ef == 0 { 200 } else { params.ef },
186                ..params
187            }),
188            group_by: None,
189            with_positions: false,
190            values: Vec::new(),
191            composite: None,
192        };
193        self.register_spec(spec)
194    }
195
196    /// `IDX.DROP` equivalent; `false` if absent. On a hit the catalog
197    /// sidecar is re-persisted so the drop survives restart.
198    pub fn idx_drop(&self, name: &[u8]) -> bool {
199        let hit = {
200            let mut g = self
201                .indexes
202                .catalog
203                .write()
204                .unwrap_or_else(std::sync::PoisonError::into_inner);
205            let (ver, cat) = &mut *g;
206            let hit = cat.drop_index(name);
207            if hit {
208                *ver += 1;
209            }
210            hit
211        };
212        if hit {
213            self.persist_index_sidecar();
214        }
215        hit
216    }
217
218    /// Range / EQ query with cursor pagination: merged across shards
219    /// in `(value, key)` order. `cursor = None` starts; the returned
220    /// cursor resumes exclusively.
221    pub fn idx_query(
222        &self,
223        name: &[u8],
224        min: &IndexValue,
225        max: &IndexValue,
226        cursor: Option<&Cursor>,
227        limit: usize,
228    ) -> KevyResult<IndexPage> {
229        let limit = limit.clamp(1, 100_000);
230        let mut all: Vec<(IndexValue, Vec<u8>)> = Vec::new();
231        self.for_each_segment(name, |seg| {
232            let (hits, _) = seg.range(min, max, cursor, limit);
233            all.extend(hits.into_iter().map(|(k, v)| (v, k)));
234        })?;
235        Ok(merge_page(all, limit))
236    }
237
238    /// Count without materializing keys.
239    pub fn idx_count(&self, name: &[u8], min: &IndexValue, max: &IndexValue) -> KevyResult<u64> {
240        let mut total = 0u64;
241        self.for_each_segment(name, |seg| total += seg.count(min, max))?;
242        Ok(total)
243    }
244
245    /// Summed segment stats (entries / bytes / coerce failures /
246    /// unique-fence duplicates).
247    pub fn idx_stats(&self, name: &[u8]) -> KevyResult<SegmentStats> {
248        let mut sum = SegmentStats::default();
249        self.for_each_segment(name, |seg| {
250            let s = seg.stats();
251            sum.entries += s.entries;
252            sum.approx_bytes += s.approx_bytes;
253            sum.coerce_failures += s.coerce_failures;
254            sum.duplicates += s.duplicates;
255        })?;
256        Ok(sum)
257    }
258
259    /// Declared indexes (name, prefix, kind), declaration order.
260    pub fn idx_list(&self) -> Vec<(Vec<u8>, Vec<u8>, IndexKind)> {
261        let g = self
262            .indexes
263            .catalog
264            .read()
265            .unwrap_or_else(std::sync::PoisonError::into_inner);
266        g.1.iter()
267            .map(|(s, _)| (s.name.clone(), s.prefix.clone(), s.kind))
268            .collect()
269    }
270
271    /// `MATCH` — BM25-ranked hits merged across shards, scored against
272    /// **global** corpus statistics so a hit's rank does not depend on
273    /// which shard it landed on (see docs/text-search.md).
274    ///
275    /// Two query-time passes: the first sums each shard's `n_docs`,
276    /// `total_len` and per-query-token `df` into one [`CorpusStats`]; the
277    /// second scores every shard against it. Only the query's tokens'
278    /// df is aggregated, not a whole-corpus table — the query narrows it.
279    #[cfg(feature = "text")]
280    pub fn idx_match(
281        &self,
282        name: &[u8],
283        query: &[u8],
284        limit: usize,
285    ) -> KevyResult<Vec<(Vec<u8>, f64)>> {
286        Ok(self
287            .idx_match_with(name, query, limit, crate::MatchOpts::default())?
288            .into_iter()
289            .map(|(key, score, _)| (key, score))
290            .collect())
291    }
292
293
294    /// Declare an aggregate index (KIND agg — write-time GROUP
295    /// BY). `ty` must be numeric.
296    pub fn idx_create_agg(
297        &self,
298        name: &[u8],
299        prefix: &[u8],
300        field: &[u8],
301        ty: ValType,
302        group_by: &[u8],
303    ) -> KevyResult<()> {
304        if !matches!(ty, ValType::I64 | ValType::F64) || group_by.is_empty() {
305            return Err(KevyError::InvalidInput("agg requires numeric type + group field".into()));
306        }
307        let spec = IndexSpec {
308            name: name.to_vec(),
309            prefix: prefix.to_vec(),
310            fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
311            ty,
312            kind: IndexKind::Agg,
313            max_bytes: 0,
314            ann: None,
315            group_by: Some(group_by.to_vec()),
316            with_positions: false,
317            values: Vec::new(),
318            composite: None,
319        };
320        self.register_spec(spec)
321    }
322
323    /// One group's merged stats across shards.
324    pub fn idx_group(&self, name: &[u8], group: &[u8]) -> KevyResult<kevy_index::GroupStats> {
325        let mut merged = kevy_index::GroupStats { count: 0, sum: 0.0, min: None, max: None };
326        let mut found = false;
327        for shard in self.shards.iter() {
328            let mut g = lock_write(shard);
329            let inner = &mut *g;
330            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
331            if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
332                found = true;
333                kevy_index::merge_group(&mut merged, &a.group(group));
334            }
335        }
336        if !found {
337            return Err(KevyError::NotFound("no such aggregate index".into()));
338        }
339        Ok(merged)
340    }
341
342    /// Top groups merged + ranked across shards.
343    pub fn idx_groups(
344        &self,
345        name: &[u8],
346        by: kevy_index::AggBy,
347        limit: usize,
348    ) -> KevyResult<Vec<(Vec<u8>, kevy_index::GroupStats)>> {
349        let limit = limit.clamp(1, 1000);
350        // HashMap merge (same O(rows×groups) trap the server reduce
351        // had — hashing keeps it linear).
352        let mut merged: std::collections::HashMap<Vec<u8>, kevy_index::GroupStats> =
353            std::collections::HashMap::new();
354        let mut found = false;
355        for shard in self.shards.iter() {
356            let mut g = lock_write(shard);
357            let inner = &mut *g;
358            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
359            if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
360                found = true;
361                for (gk, st) in a.all_groups() {
362                    match merged.get_mut(&gk) {
363                        Some(m) => kevy_index::merge_group(m, &st),
364                        None => {
365                            merged.insert(gk, st);
366                        }
367                    }
368                }
369            }
370        }
371        if !found {
372            return Err(KevyError::NotFound("no such aggregate index".into()));
373        }
374        let mut ranked: Vec<(Vec<u8>, kevy_index::GroupStats)> = merged.into_iter().collect();
375        kevy_index::sort_groups(&mut ranked, by);
376        ranked.truncate(limit);
377        Ok(ranked)
378    }
379
380    /// `KNN` — nearest neighbors merged ascending across shards.
381    /// `ef` = query beam width (0 = engine default; recall knob).
382    #[cfg(feature = "vector")]
383    pub fn idx_knn(
384        &self,
385        name: &[u8],
386        query: &[f32],
387        k: usize,
388        ef: usize,
389    ) -> KevyResult<Vec<(Vec<u8>, f32)>> {
390        let k = k.clamp(1, 1000);
391        let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
392        let mut found = false;
393        for shard in self.shards.iter() {
394            let mut g = lock_write(shard);
395            let inner = &mut *g;
396            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
397            if let Some((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name == name) {
398                found = true;
399                all.extend(graph.knn(query, k, ef));
400            }
401        }
402        if !found {
403            return Err(KevyError::NotFound("no such vector index".into()));
404        }
405        all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
406        all.truncate(k);
407        Ok(all)
408    }
409
410    /// Without `persist` there is no data dir — the catalog lives only
411    /// in memory, so the sidecar halves are no-ops.
412    #[cfg(not(feature = "persist"))]
413    fn persist_index_sidecar(&self) {}
414
415    #[cfg(not(feature = "persist"))]
416    pub(crate) fn idx_boot(&self) {}
417
418    fn for_each_segment(
419        &self,
420        name: &[u8],
421        mut f: impl FnMut(&Segment),
422    ) -> KevyResult<()> {
423        let mut found = false;
424        for shard in self.shards.iter() {
425            let mut g = lock_write(shard);
426            let inner = &mut *g;
427            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
428            if let Some((_, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name == name) {
429                found = true;
430                f(seg);
431            }
432        }
433        if found {
434            Ok(())
435        } else {
436            Err(KevyError::NotFound("no such index".into()))
437        }
438    }
439
440    #[cfg(feature = "persist")]
441    fn persist_index_sidecar(&self) {
442        let Some(dir) = &self.config.data_dir else { return };
443        let g = self
444            .indexes
445            .catalog
446            .read()
447            .unwrap_or_else(std::sync::PoisonError::into_inner);
448        let tmp = dir.join("index-catalog.meta.tmp");
449        if std::fs::write(&tmp, g.1.to_sidecar()).is_ok() {
450            let _ = std::fs::rename(&tmp, dir.join(SIDECAR));
451        }
452    }
453
454    /// Boot half — load a persisted catalog (indexes rebuild lazily on
455    /// first touch via `sync_segs`).
456    #[cfg(feature = "persist")]
457    pub(crate) fn idx_boot(&self) {
458        let Some(dir) = &self.config.data_dir else { return };
459        if let Ok(text) = std::fs::read_to_string(dir.join(SIDECAR))
460            && let Some(cat) = Catalog::from_sidecar(&text)
461            && !cat.is_empty()
462        {
463            let mut g = self
464                .indexes
465                .catalog
466                .write()
467                .unwrap_or_else(std::sync::PoisonError::into_inner);
468            *g = (g.0 + 1, cat);
469        }
470    }
471}
kevy_embedded/ops_index.rs

kevy_embedded/
ops_index.rs
