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.
31///
32/// ```
33/// use kevy_embedded::{Config, IndexKind, IndexPage, IndexValType, IndexValue, Store};
34/// let s = Store::open(Config::default())?;
35/// s.idx_create(b"by_age", b"u:", b"age", IndexValType::I64, IndexKind::Range)?;
36/// for (k, age) in [(&b"u:1"[..], &b"30"[..]), (b"u:2", b"40")] {
37///     s.hset(k, &[(b"age", age)])?;
38/// }
39/// let (lo, hi) = (IndexValue::I64(0), IndexValue::I64(99));
40/// let (rows, cursor): IndexPage = s.idx_query(b"by_age", &lo, &hi, None, 1)?;
41/// assert_eq!(rows, [(b"u:1".to_vec(), IndexValue::I64(30))]);
42/// let (rest, _) = s.idx_query(b"by_age", &lo, &hi, cursor.as_ref(), 10)?; // resume
43/// assert_eq!(rest[0].0, b"u:2");
44/// # Ok::<(), kevy_embedded::KevyError>(())
45/// ```
46pub type IndexPage = (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>);
47
48/// One field's highlight: its name and the `(start, end)` match spans.
49#[cfg(feature = "text")]
50pub type FieldSpans = (Vec<u8>, Vec<(u32, u32)>);
51/// A highlighted MATCH hit: key, score, and per-field [`FieldSpans`].
52#[cfg(feature = "text")]
53pub type HighlightedHit = (Vec<u8>, f64, Vec<FieldSpans>);
54
55// `idx_match_with`, its clause options and the span-mapping helper live
56// in a child module to keep this file under the 500-LOC ceiling; the
57// text-index specifics — declaring a multi-field one, and gathering the
58// corpus statistics a global BM25 scores against — live in another.
59#[cfg(feature = "text")]
60#[path = "ops_index_highlight.rs"]
61pub(crate) mod highlight;
62
63// The clause-carrying scalar query (capacity arc G1) and the
64// [`ValueFilter`] predicate shape it shares with MATCH — independent of
65// the `text` feature: a range index filters fine without a tokenizer.
66#[path = "ops_index_claused.rs"]
67pub(crate) mod claused;
68
69// The optional query clauses both of those take; MATCH's set has its own file.
70#[cfg(feature = "text")]
71#[path = "ops_index_match_opts.rs"]
72pub(crate) mod match_opts;
73#[path = "ops_index_opts.rs"]
74pub(crate) mod opts;
75
76// The auto-declaration loop's observation face (refusal log + advice).
77#[path = "ops_index_advise.rs"]
78pub(crate) mod advise;
79
80// The read-only admin surface (stats, enumeration).
81#[path = "ops_index_admin.rs"]
82mod admin;
83
84#[cfg(feature = "text")]
85#[path = "ops_index_text.rs"]
86mod text;
87
88// The embedded MATCH's cold seams (windowed tables' frozen buckets).
89#[cfg(feature = "text")]
90#[path = "ops_index_text_cold.rs"]
91pub(crate) mod text_cold;
92
93/// Sort merged `(value, key)` hits, cut to `limit`, and derive the
94/// resume cursor. Shared by `Store::idx_query` and the transaction twin
95/// on `AtomicAllShards`, which differ only in where the segments come
96/// from — the pagination has to agree exactly or a cursor taken inside
97/// a transaction would not resume outside one.
98pub(crate) fn merge_page(mut all: Vec<(IndexValue, Vec<u8>)>, limit: usize) -> IndexPage {
99    all.sort();
100    all.truncate(limit);
101    let next = if all.len() == limit {
102        all.last().map(|(v, k)| Cursor::new(v.clone(), k.clone()))
103    } else {
104        None
105    };
106    (all.into_iter().map(|(v, k)| (k, v)).collect(), next)
107}
108
109/// A spec from `b`, its refusal reported as the catalog's are.
110pub(crate) fn built(b: kevy_index::IndexSpecBuilder) -> KevyResult<IndexSpec> {
111    Ok(b.build().map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?)
112}
113
114/// Store-level index state: catalog + a version stamp the per-shard
115/// segment lists sync against, and each declared path's usage cell
116/// (the refusal log's dual — reclaim-face raw material).
117#[derive(Debug, Default)]
118pub(crate) struct IndexReg {
119    pub(crate) catalog: RwLock<(u64, Catalog)>,
120    pub(crate) usage:
121        RwLock<std::collections::HashMap<Vec<u8>, std::sync::Arc<kevy_index::UsageCell>>>,
122}
123
124/// A shard's window runtime beside its segment — `()` stand-in on
125/// wasm, where the cold tier compiles out.
126#[cfg(not(target_arch = "wasm32"))]
127pub(crate) type WinRef<'a> = Option<&'a kevy_window::WindowRt>;
128#[cfg(target_arch = "wasm32")]
129pub(crate) type WinRef<'a> = Option<&'a core::convert::Infallible>;
130
131/// Per-shard segment list, kept inside `Inner` (guarded by the shard
132/// lock).
133#[derive(Debug, Default)]
134pub(crate) struct ShardSegs {
135    pub(crate) version: u64,
136    pub(crate) segs: Vec<(IndexSpec, Segment)>,
137    /// Inverted segments for KIND text specs (parallel list —
138    /// a spec appears in exactly one of the lists).
139    #[cfg(feature = "text")]
140    pub(crate) text: Vec<(IndexSpec, kevy_text::TextSegment)>,
141    /// HNSW graphs for KIND ann specs.
142    #[cfg(feature = "vector")]
143    pub(crate) ann: Vec<(IndexSpec, kevy_vector::Hnsw)>,
144    /// Aggregate segments for KIND agg specs.
145    pub(crate) agg: Vec<(IndexSpec, kevy_index::AggSegment)>,
146    /// Sliding-window runtimes, name-keyed — reconciled and slid by
147    /// the reaper's window tick; empty under a manual reaper (nothing
148    /// slides, everything stays hot, queries need no cold half).
149    #[cfg(not(target_arch = "wasm32"))]
150    pub(crate) windows: Vec<(Vec<u8>, kevy_window::WindowRt)>,
151    /// A windowed table's text indexes' cold directories, name-keyed —
152    /// reconciled and fed by the same window tick (the driver's
153    /// eviction batch freezes out of every same-table text index).
154    #[cfg(all(feature = "text", not(target_arch = "wasm32")))]
155    pub(crate) cold_text: Vec<(Vec<u8>, kevy_window::TextColdDir)>,
156    /// `reserved_bytes` generation cache: set by every
157    /// segment-mutating chokepoint (`on_commit` applies, list
158    /// rebuilds, FLUSH resets) via [`ShardSegs::mark_stats_dirty`];
159    /// an idle tick reads the cached sum instead of walking every
160    /// segment's stats. Compiled with the tier backend only — on
161    /// targets without it (wasm) nothing reads the cache.
162    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
163    pub(crate) stats_dirty: bool,
164    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
165    pub(crate) reserved_cache: u64,
166    /// The store record's state between drains.
167    pub(crate) drain: crate::ops_index_changes::Drain,
168}
169
170impl ShardSegs {
171    /// The window runtime for `name`, if that index is a windowed
172    /// table's window access path AND a tick has reconciled it (only
173    /// the reaper's window tick populates the list, so a manual-reaper
174    /// or memory-only store always answers `None` — nothing slid).
175    #[cfg(not(target_arch = "wasm32"))]
176    pub(crate) fn window_of(&self, name: &[u8]) -> Option<&kevy_window::WindowRt> {
177        self.windows.iter().find(|(n, _)| n == name).map(|(_, w)| w)
178    }
179
180    /// Invalidate the `reserved_bytes` cache — a no-op on targets
181    /// without the tier backend, so mutation chokepoints call it
182    /// unconditionally.
183    #[inline]
184    pub(crate) fn mark_stats_dirty(&mut self) {
185        #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
186        {
187            self.stats_dirty = true;
188        }
189    }
190}
191
192impl Store {
193    /// `IDX.CREATE` equivalent. Builds synchronously; errors on
194    /// duplicate name / cap / bad spec.
195    pub fn idx_create(
196        &self,
197        name: &[u8],
198        prefix: &[u8],
199        field: &[u8],
200        ty: ValType,
201        kind: IndexKind,
202    ) -> KevyResult<()> {
203        if prefix.is_empty() {
204            return Err(KevyError::InvalidInput("empty prefix".into()));
205        }
206        #[cfg(not(feature = "text"))]
207        if kind == IndexKind::Text {
208            return Err(KevyError::Unsupported("text indexes need the `text` feature".into()));
209        }
210        #[cfg(not(feature = "vector"))]
211        if kind == IndexKind::Ann {
212            return Err(KevyError::Unsupported("vector indexes need the `vector` feature".into()));
213        }
214        let spec = built(IndexSpec::builder(name, prefix, kind, ty).with_field(field))?;
215        self.catalog_change(|| self.register_spec(spec))
216    }
217
218    /// Admit `spec` and build it on every shard; the caller records the
219    /// change ([`Self::catalog_change`]).
220    pub(crate) fn register_spec(&self, spec: IndexSpec) -> KevyResult<()> {
221        // Tiering floor refusal: body in
222        // `ops_index_sync::tier_floor_check` (500-LOC rule).
223        #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
224        crate::ops_index_sync::tier_floor_check(&self.shards)?;
225        {
226            let mut g =
227                self.indexes.catalog.write().unwrap_or_else(std::sync::PoisonError::into_inner);
228            let (ver, cat) = &mut *g;
229            cat.create(spec).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
230            *ver += 1;
231        }
232        self.advise_clear();
233        self.usage_rekey();
234        // Build every shard's slice now (each under its own lock).
235        for shard in self.shards.iter() {
236            let mut g = lock_write(shard);
237            let inner = &mut *g;
238            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
239        }
240        Ok(())
241    }
242
243    /// Declare an ANN index (KIND ann, TYPE vector). `params.m`
244    /// / `params.ef` of 0 select the defaults (16 / 200).
245    #[cfg(feature = "vector")]
246    pub fn idx_create_ann(
247        &self,
248        name: &[u8],
249        prefix: &[u8],
250        field: &[u8],
251        params: kevy_index::AnnSpec,
252    ) -> KevyResult<()> {
253        // M = 1 has no level distribution (1/ln 1); 0 means the default
254        if params.dim == 0 || params.distance > 2 || params.m == 1 {
255            return Err(KevyError::InvalidInput("bad ann parameters".into()));
256        }
257        let ann = params
258            .with_m(if params.m == 0 { 16 } else { params.m })
259            .with_ef(if params.ef == 0 { 200 } else { params.ef });
260        let spec = IndexSpec::builder(name, prefix, IndexKind::Ann, ValType::Vector)
261            .with_field(field)
262            .with_ann(ann);
263        let spec = built(spec)?;
264        self.catalog_change(|| self.register_spec(spec))
265    }
266
267    /// `IDX.DROP` equivalent; `false` if absent. Refused on a replica
268    /// and after [`Store::shutdown`], like every write.
269    pub fn idx_drop(&self, name: &[u8]) -> KevyResult<bool> {
270        self.catalog_change(|| Ok(self.drop_index(name)))
271    }
272
273    /// Drop the index named `name` from the catalog; the caller records
274    /// the change.
275    pub(crate) fn drop_index(&self, name: &[u8]) -> bool {
276        let hit = {
277            let mut g =
278                self.indexes.catalog.write().unwrap_or_else(std::sync::PoisonError::into_inner);
279            let (ver, cat) = &mut *g;
280            let hit = cat.drop_index(name);
281            if hit {
282                *ver += 1;
283            }
284            hit
285        };
286        if hit {
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 [`kevy_text::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::builder(name, prefix, IndexKind::Agg, ty)
329            .with_field(field)
330            .with_group_by(group_by);
331        let spec = built(spec)?;
332        self.catalog_change(|| self.register_spec(spec))
333    }
334
335    /// One group's merged stats across shards.
336    pub fn idx_group(&self, name: &[u8], group: &[u8]) -> KevyResult<kevy_index::GroupStats> {
337        let mut merged = kevy_index::GroupStats::default();
338        let mut found = false;
339        for shard in self.shards.iter() {
340            let mut g = lock_write(shard);
341            let inner = &mut *g;
342            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
343            if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name() == name) {
344                found = true;
345                merged.merge(&a.group(group));
346            }
347        }
348        if !found {
349            return Err(KevyError::NotFound("no such aggregate index".into()));
350        }
351        Ok(merged)
352    }
353
354    /// Top groups merged + ranked across shards.
355    pub fn idx_groups(
356        &self,
357        name: &[u8],
358        by: kevy_index::AggBy,
359        limit: usize,
360    ) -> KevyResult<Vec<(Vec<u8>, kevy_index::GroupStats)>> {
361        let limit = limit.clamp(1, 1000);
362        // HashMap merge (same O(rows×groups) trap the server reduce
363        // had — hashing keeps it linear).
364        let mut merged: std::collections::HashMap<Vec<u8>, kevy_index::GroupStats> =
365            std::collections::HashMap::new();
366        let mut found = false;
367        for shard in self.shards.iter() {
368            let mut g = lock_write(shard);
369            let inner = &mut *g;
370            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
371            if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name() == name) {
372                found = true;
373                for (gk, st) in a.all_groups() {
374                    match merged.get_mut(&gk) {
375                        Some(m) => m.merge(&st),
376                        None => {
377                            merged.insert(gk, st);
378                        }
379                    }
380                }
381            }
382        }
383        if !found {
384            return Err(KevyError::NotFound("no such aggregate index".into()));
385        }
386        let mut ranked: Vec<(Vec<u8>, kevy_index::GroupStats)> = merged.into_iter().collect();
387        kevy_index::sort_groups(&mut ranked, by);
388        ranked.truncate(limit);
389        Ok(ranked)
390    }
391
392    /// `KNN` — nearest neighbors merged ascending across shards.
393    /// `ef` = query beam width (0 = engine default; recall knob).
394    #[cfg(feature = "vector")]
395    pub fn idx_knn(
396        &self,
397        name: &[u8],
398        query: &[f32],
399        k: usize,
400        ef: usize,
401    ) -> KevyResult<Vec<(Vec<u8>, f32)>> {
402        let k = k.clamp(1, 1000);
403        let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
404        let mut found = false;
405        for shard in self.shards.iter() {
406            let mut g = lock_write(shard);
407            let inner = &mut *g;
408            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
409            if let Some((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name() == name) {
410                found = true;
411                all.extend(graph.knn(query, k, ef));
412            }
413        }
414        if !found {
415            return Err(KevyError::NotFound("no such vector index".into()));
416        }
417        all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
418        all.truncate(k);
419        Ok(all)
420    }
421
422    fn for_each_segment(&self, name: &[u8], mut f: impl FnMut(&Segment)) -> 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 { Ok(()) } else { Err(KevyError::NotFound("no such index".into())) }
434    }
435}