kevy_embedded/
ops_index.rs

1//! v2.5 — embedded secondary-index API (RFC LOCKED; 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 std::io;
19use std::sync::RwLock;
20
21use kevy_index::{Catalog, Cursor, IndexKind, IndexSpec, IndexValue, Segment, SegmentStats, ValType};
22
23use crate::store::{Store, lock_write};
24
25pub(crate) use crate::ops_index_sync::{each_written_key_pub, on_commit, sync_segs};
26
27/// One page of index hits plus the cursor to resume from.
28pub type IndexPage = (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>);
29
30/// Store-level index state: catalog + a version stamp the per-shard
31/// segment lists sync against.
32#[derive(Default)]
33pub(crate) struct IndexReg {
34    pub(crate) catalog: RwLock<(u64, Catalog)>,
35}
36
37/// Per-shard segment list, kept inside `Inner` (guarded by the shard
38/// lock).
39#[derive(Default)]
40pub(crate) struct ShardSegs {
41    pub(crate) version: u64,
42    pub(crate) segs: Vec<(IndexSpec, Segment)>,
43    /// v2.7: inverted segments for KIND text specs (parallel list —
44    /// a spec appears in exactly one of the lists).
45    pub(crate) text: Vec<(IndexSpec, kevy_text::TextSegment)>,
46    /// v2.8: HNSW graphs for KIND ann specs.
47    pub(crate) ann: Vec<(IndexSpec, kevy_vector::Hnsw)>,
48    /// v3.1: aggregate segments for KIND agg specs.
49    pub(crate) agg: Vec<(IndexSpec, kevy_index::AggSegment)>,
50}
51
52const SIDECAR: &str = "index-catalog.meta";
53
54impl Store {
55    /// `IDX.CREATE` equivalent. Builds synchronously; errors on
56    /// duplicate name / cap / bad spec.
57    pub fn idx_create(
58        &self,
59        name: &[u8],
60        prefix: &[u8],
61        field: &[u8],
62        ty: ValType,
63        kind: IndexKind,
64    ) -> io::Result<()> {
65        if prefix.is_empty() {
66            return Err(io::Error::new(io::ErrorKind::InvalidInput, "empty prefix"));
67        }
68        let spec = IndexSpec {
69            name: name.to_vec(),
70            prefix: prefix.to_vec(),
71            field: field.to_vec(),
72            ty,
73            kind,
74            max_bytes: 0,
75            ann: None,
76            group_by: None,
77        };
78        self.register_spec(spec)
79    }
80
81    fn register_spec(&self, spec: IndexSpec) -> io::Result<()> {
82        {
83            let mut g = self
84                .indexes
85                .catalog
86                .write()
87                .unwrap_or_else(std::sync::PoisonError::into_inner);
88            let (ver, cat) = &mut *g;
89            cat.create(spec)
90                .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
91            *ver += 1;
92        }
93        self.persist_index_sidecar();
94        // Build every shard's slice now (each under its own lock).
95        for shard in self.shards.iter() {
96            let mut g = lock_write(shard);
97            let inner = &mut *g;
98            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
99        }
100        Ok(())
101    }
102
103    /// v2.8: declare an ANN index (KIND ann, TYPE vector). `params.m`
104    /// / `params.ef` of 0 select the defaults (16 / 200).
105    pub fn idx_create_ann(
106        &self,
107        name: &[u8],
108        prefix: &[u8],
109        field: &[u8],
110        params: kevy_index::AnnSpec,
111    ) -> io::Result<()> {
112        if params.dim == 0 || params.distance > 2 {
113            return Err(io::Error::new(io::ErrorKind::InvalidInput, "bad ann parameters"));
114        }
115        let spec = IndexSpec {
116            name: name.to_vec(),
117            prefix: prefix.to_vec(),
118            field: field.to_vec(),
119            ty: ValType::Vector,
120            kind: IndexKind::Ann,
121            max_bytes: 0,
122            ann: Some(kevy_index::AnnSpec {
123                m: if params.m == 0 { 16 } else { params.m },
124                ef: if params.ef == 0 { 200 } else { params.ef },
125                ..params
126            }),
127            group_by: None,
128        };
129        self.register_spec(spec)
130    }
131
132    /// `IDX.DROP` equivalent; `false` if absent. On a hit the catalog
133    /// sidecar is re-persisted so the drop survives restart.
134    pub fn idx_drop(&self, name: &[u8]) -> bool {
135        let hit = {
136            let mut g = self
137                .indexes
138                .catalog
139                .write()
140                .unwrap_or_else(std::sync::PoisonError::into_inner);
141            let (ver, cat) = &mut *g;
142            let hit = cat.drop_index(name);
143            if hit {
144                *ver += 1;
145            }
146            hit
147        };
148        if hit {
149            self.persist_index_sidecar();
150        }
151        hit
152    }
153
154    /// Range / EQ query with cursor pagination: merged across shards
155    /// in `(value, key)` order. `cursor = None` starts; the returned
156    /// cursor resumes exclusively.
157    pub fn idx_query(
158        &self,
159        name: &[u8],
160        min: &IndexValue,
161        max: &IndexValue,
162        cursor: Option<&Cursor>,
163        limit: usize,
164    ) -> io::Result<IndexPage> {
165        let limit = limit.clamp(1, 100_000);
166        let mut all: Vec<(IndexValue, Vec<u8>)> = Vec::new();
167        self.for_each_segment(name, |seg| {
168            let (hits, _) = seg.range(min, max, cursor, limit);
169            all.extend(hits.into_iter().map(|(k, v)| (v, k)));
170        })?;
171        all.sort();
172        all.truncate(limit);
173        let next = if all.len() == limit {
174            all.last().map(|(v, k)| Cursor { value: v.clone(), key: k.clone() })
175        } else {
176            None
177        };
178        Ok((all.into_iter().map(|(v, k)| (k, v)).collect(), next))
179    }
180
181    /// Count without materializing keys.
182    pub fn idx_count(&self, name: &[u8], min: &IndexValue, max: &IndexValue) -> io::Result<u64> {
183        let mut total = 0u64;
184        self.for_each_segment(name, |seg| total += seg.count(min, max))?;
185        Ok(total)
186    }
187
188    /// Summed segment stats (entries / bytes / coerce failures /
189    /// unique-fence duplicates).
190    pub fn idx_stats(&self, name: &[u8]) -> io::Result<SegmentStats> {
191        let mut sum = SegmentStats::default();
192        self.for_each_segment(name, |seg| {
193            let s = seg.stats();
194            sum.entries += s.entries;
195            sum.approx_bytes += s.approx_bytes;
196            sum.coerce_failures += s.coerce_failures;
197            sum.duplicates += s.duplicates;
198        })?;
199        Ok(sum)
200    }
201
202    /// Declared indexes (name, prefix, kind), declaration order.
203    pub fn idx_list(&self) -> Vec<(Vec<u8>, Vec<u8>, IndexKind)> {
204        let g = self
205            .indexes
206            .catalog
207            .read()
208            .unwrap_or_else(std::sync::PoisonError::into_inner);
209        g.1.iter()
210            .map(|(s, _)| (s.name.clone(), s.prefix.clone(), s.kind))
211            .collect()
212    }
213
214    /// v2.7 `MATCH` — BM25-ranked hits merged across shards
215    /// (shard-local statistics; see docs/text-search.md).
216    pub fn idx_match(
217        &self,
218        name: &[u8],
219        query: &[u8],
220        limit: usize,
221    ) -> io::Result<Vec<(Vec<u8>, f64)>> {
222        let limit = limit.clamp(1, 1000);
223        let mut all: Vec<kevy_text::TextMatch> = Vec::new();
224        let mut found = false;
225        for shard in self.shards.iter() {
226            let mut g = lock_write(shard);
227            let inner = &mut *g;
228            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
229            if let Some((_, ts)) = inner.idx_segs.text.iter().find(|(s, _)| s.name == name) {
230                found = true;
231                all.extend(ts.matches(query, limit));
232            }
233        }
234        if !found {
235            return Err(io::Error::new(io::ErrorKind::NotFound, "no such text index"));
236        }
237        all.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key)));
238        all.truncate(limit);
239        Ok(all.into_iter().map(|m| (m.key, m.score)).collect())
240    }
241
242    /// v3.1: declare an aggregate index (KIND agg — write-time GROUP
243    /// BY). `ty` must be numeric.
244    pub fn idx_create_agg(
245        &self,
246        name: &[u8],
247        prefix: &[u8],
248        field: &[u8],
249        ty: ValType,
250        group_by: &[u8],
251    ) -> io::Result<()> {
252        if !matches!(ty, ValType::I64 | ValType::F64) || group_by.is_empty() {
253            return Err(io::Error::new(io::ErrorKind::InvalidInput, "agg requires numeric type + group field"));
254        }
255        let spec = IndexSpec {
256            name: name.to_vec(),
257            prefix: prefix.to_vec(),
258            field: field.to_vec(),
259            ty,
260            kind: IndexKind::Agg,
261            max_bytes: 0,
262            ann: None,
263            group_by: Some(group_by.to_vec()),
264        };
265        self.register_spec(spec)
266    }
267
268    /// v3.1: one group's merged stats across shards.
269    pub fn idx_group(&self, name: &[u8], group: &[u8]) -> io::Result<kevy_index::GroupStats> {
270        let mut merged = kevy_index::GroupStats { count: 0, sum: 0.0, min: None, max: None };
271        let mut found = false;
272        for shard in self.shards.iter() {
273            let mut g = lock_write(shard);
274            let inner = &mut *g;
275            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
276            if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
277                found = true;
278                kevy_index::merge_group(&mut merged, &a.group(group));
279            }
280        }
281        if !found {
282            return Err(io::Error::new(io::ErrorKind::NotFound, "no such aggregate index"));
283        }
284        Ok(merged)
285    }
286
287    /// v3.1: top groups merged + ranked across shards.
288    pub fn idx_groups(
289        &self,
290        name: &[u8],
291        by: kevy_index::AggBy,
292        limit: usize,
293    ) -> io::Result<Vec<(Vec<u8>, kevy_index::GroupStats)>> {
294        let limit = limit.clamp(1, 1000);
295        // HashMap merge (same O(rows×groups) trap the server reduce
296        // had — hashing keeps it linear).
297        let mut merged: std::collections::HashMap<Vec<u8>, kevy_index::GroupStats> =
298            std::collections::HashMap::new();
299        let mut found = false;
300        for shard in self.shards.iter() {
301            let mut g = lock_write(shard);
302            let inner = &mut *g;
303            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
304            if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name == name) {
305                found = true;
306                for (gk, st) in a.all_groups() {
307                    match merged.get_mut(&gk) {
308                        Some(m) => kevy_index::merge_group(m, &st),
309                        None => {
310                            merged.insert(gk, st);
311                        }
312                    }
313                }
314            }
315        }
316        if !found {
317            return Err(io::Error::new(io::ErrorKind::NotFound, "no such aggregate index"));
318        }
319        let mut ranked: Vec<(Vec<u8>, kevy_index::GroupStats)> = merged.into_iter().collect();
320        kevy_index::sort_groups(&mut ranked, by);
321        ranked.truncate(limit);
322        Ok(ranked)
323    }
324
325    /// v2.8 `KNN` — nearest neighbors merged ascending across shards.
326    /// `ef` = query beam width (0 = engine default; recall knob).
327    pub fn idx_knn(
328        &self,
329        name: &[u8],
330        query: &[f32],
331        k: usize,
332        ef: usize,
333    ) -> io::Result<Vec<(Vec<u8>, f32)>> {
334        let k = k.clamp(1, 1000);
335        let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
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((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name == name) {
342                found = true;
343                all.extend(graph.knn(query, k, ef));
344            }
345        }
346        if !found {
347            return Err(io::Error::new(io::ErrorKind::NotFound, "no such vector index"));
348        }
349        all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
350        all.truncate(k);
351        Ok(all)
352    }
353
354    fn for_each_segment(
355        &self,
356        name: &[u8],
357        mut f: impl FnMut(&Segment),
358    ) -> io::Result<()> {
359        let mut found = false;
360        for shard in self.shards.iter() {
361            let mut g = lock_write(shard);
362            let inner = &mut *g;
363            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
364            if let Some((_, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name == name) {
365                found = true;
366                f(seg);
367            }
368        }
369        if found {
370            Ok(())
371        } else {
372            Err(io::Error::new(io::ErrorKind::NotFound, "no such index"))
373        }
374    }
375
376    fn persist_index_sidecar(&self) {
377        let Some(dir) = &self.config.data_dir else { return };
378        let g = self
379            .indexes
380            .catalog
381            .read()
382            .unwrap_or_else(std::sync::PoisonError::into_inner);
383        let tmp = dir.join("index-catalog.meta.tmp");
384        if std::fs::write(&tmp, g.1.to_sidecar()).is_ok() {
385            let _ = std::fs::rename(&tmp, dir.join(SIDECAR));
386        }
387    }
388
389    /// Boot half — load a persisted catalog (indexes rebuild lazily on
390    /// first touch via `sync_segs`).
391    pub(crate) fn idx_boot(&self) {
392        let Some(dir) = &self.config.data_dir else { return };
393        if let Ok(text) = std::fs::read_to_string(dir.join(SIDECAR))
394            && let Some(cat) = Catalog::from_sidecar(&text)
395            && !cat.is_empty()
396        {
397            let mut g = self
398                .indexes
399                .catalog
400                .write()
401                .unwrap_or_else(std::sync::PoisonError::into_inner);
402            *g = (g.0 + 1, cat);
403        }
404    }
405}
kevy_embedded/ops_index.rs

kevy_embedded/
ops_index.rs
