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