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, mut found) = (0f64, 0u64, false);
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        for shard in self.shards.iter() {
81            let mut g = lock_write(shard);
82            let inner = &mut *g;
83            sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
84            if let Some((_, ts)) = inner.idx_segs.text.iter().find(|(s, _)| s.name == name) {
85                found = true;
86                n_docs += ts.docs() as f64;
87                total_len += ts.total_len_in(scope);
88                let opts =
89                    kevy_text::QueryOpts { typo, fields: scope, ..kevy_text::QueryOpts::default() };
90                let tokdf = ts.query_df_in(text, opts);
91                // The shard's frozen buckets are pass-1 contributors
92                // like any other shard — the server seam's mirror.
93                #[cfg(not(target_arch = "wasm32"))]
94                super::text_cold::fold_cold_stats(
95                    inner.idx_segs.cold_text_of(name),
96                    &tokdf,
97                    &mut n_docs,
98                    &mut total_len,
99                    &mut df,
100                );
101                for (t, d) in tokdf {
102                    *df.entry(t).or_insert(0) += d;
103                }
104            }
105        }
106        if !found {
107            return Err(KevyError::NotFound("no such text index".into()));
108        }
109        let avgdl = if n_docs > 0.0 { total_len as f64 / n_docs } else { 0.0 };
110        Ok(kevy_text::CorpusStats { n_docs, avgdl, df })
111    }
112}