Skip to main content

kevy_embedded/
ops_index_text.rs

1//! Text-index specifics for the embedded API: declaring a multi-field
2//! text index, and gathering the corpus statistics a global BM25 scores
3//! against. Split from `ops_index.rs` for the 500-LOC house rule (a
4//! `#[path]` child module, so it reaches `Store`'s private fields).
5
6use kevy_index::{IndexKind, IndexSpec, ValType};
7
8use super::sync_segs;
9use crate::store::{Store, lock_write};
10use crate::{KevyError, KevyResult};
11
12impl Store {
13    /// A multi-field text index over several weighted hash fields — what
14    /// the wire creates with `ON FIELDS f1 f2 WEIGHTS 2 1`, and what `IN`
15    /// scopes to. A weight scales that field's term frequencies; 1.0 is
16    /// neutral.
17    ///
18    /// `positions` records token offsets so phrase queries can verify
19    /// adjacency, at the cost of the positional side-channel's memory.
20    /// `values` names hash fields stored per document with the type their
21    /// bytes compare as — what `FILTER` reads. An index that never
22    /// filters declares none.
23    pub fn idx_create_text(
24        &self,
25        name: &[u8],
26        prefix: &[u8],
27        fields: &[(&[u8], f32)],
28        positions: bool,
29        values: &[(&[u8], ValType)],
30    ) -> KevyResult<()> {
31        if prefix.is_empty() {
32            return Err(KevyError::InvalidInput("empty prefix".into()));
33        }
34        if fields.is_empty() {
35            return Err(KevyError::InvalidInput("a text index needs at least one field".into()));
36        }
37        let spec = IndexSpec {
38            name: name.to_vec(),
39            prefix: prefix.to_vec(),
40            fields: fields
41                .iter()
42                .map(|(f, w)| kevy_index::FieldSpec { name: f.to_vec(), weight: *w })
43                .collect(),
44            ty: ValType::Str,
45            kind: IndexKind::Text,
46            max_bytes: 0,
47            ann: None,
48            group_by: None,
49            with_positions: positions,
50            values: values
51                .iter()
52                .map(|(n, ty)| kevy_index::ValueSpec { name: n.to_vec(), ty: *ty })
53                .collect(),
54            composite: None,
55        };
56        self.register_spec(spec)
57    }
58
59    /// Corpus-wide BM25 statistics for one query, over its field scope.
60    ///
61    /// Every shard reports its own document count, its token total and,
62    /// per query term, its local document frequency; summing them is what
63    /// makes the scores comparable across shards. Restricted to a field
64    /// scope, each of those numbers describes those fields rather than
65    /// whole documents, so a scoped query's `avgdl` is the average length
66    /// *of the named fields*.
67    pub(crate) fn text_corpus_stats_in(
68        &self,
69        name: &[u8],
70        text: &[u8],
71        typo: u32,
72        scope: &[usize],
73    ) -> KevyResult<kevy_text::CorpusStats> {
74        let (mut n_docs, mut total_len) = (0f64, 0u64);
75        // Accumulated per shard from `query_df_in`, which expands `word*`
76        // prefixes against that shard's dictionary — so the df map ends up
77        // keyed by the union of every shard's query terms and prefix
78        // expansions, each summed to a global df.
79        let mut df: std::collections::HashMap<Vec<u8>, u32> = std::collections::HashMap::new();
80        let mut found = false;
81        for shard in self.shards.iter() {
82            let mut g = lock_write(shard);
83            let inner = &mut *g;
84            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
85            if let Some((_, ts)) = inner.idx_segs.text.iter().find(|(s, _)| s.name == name) {
86                found = true;
87                n_docs += ts.docs() as f64;
88                total_len += ts.total_len_in(scope);
89                let opts = kevy_text::QueryOpts {
90                    stats: None,
91                    typo,
92                    fields: scope,
93                    filter: &[],
94                    sort: None,
95                    distinct: None,
96                };
97                for (t, d) in ts.query_df_in(text, opts) {
98                    *df.entry(t).or_insert(0) += d;
99                }
100            }
101        }
102        if !found {
103            return Err(KevyError::NotFound("no such text index".into()));
104        }
105        let avgdl = if n_docs > 0.0 { total_len as f64 / n_docs } else { 0.0 };
106        Ok(kevy_text::CorpusStats { n_docs, avgdl, df })
107    }
108}