Skip to main content

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