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