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