Skip to main content

kevy_window/
text_query.rs

1//! The cold directory's query face — what the two MATCH passes read
2//! out of the frozen buckets. Pass 1 takes the live corpus counters;
3//! pass 2 takes a whole clause-faithful page: bare terms and phrases
4//! accumulate (the hot engine's exact clause semantics, over the
5//! frozen postings), FILTER prunes on the frozen stored values,
6//! FACET counts the filtered match set, and selection runs in the
7//! page's own order — score order, or `sorted_order` under SORT, the
8//! same rule the hot top-K and the cross-shard merge use.
9//!
10//! A child module of [`super`] (`#[path]`), so it reaches the cold
11//! directory's private shape.
12
13use std::collections::HashMap;
14
15use kevy_text::cold::{decode_fwd, posting_df, score_cold, score_cold_phrase};
16use kevy_text::{CorpusStats, sorted_order};
17
18use super::TextColdDir;
19
20/// Everything pass 2 asks of the cold directory.
21pub struct ColdPageQuery<'a> {
22    /// Bare terms, sorted and deduplicated (the hot engine's rule).
23    pub bare: Vec<Vec<u8>>,
24    /// Each phrase's token sequence.
25    pub phrases: Vec<Vec<Vec<u8>>>,
26    /// The injected global statistics both passes score with.
27    pub stats: &'a CorpusStats,
28    /// `FILTER` predicates, ANDed, over the frozen stored values.
29    pub filter: &'a [kevy_text::Filter<'a>],
30    /// `SORT`: the page order is the sort key's, not the score's.
31    pub sort: Option<&'a kevy_text::Sort<'a>>,
32    /// `DISTINCT`: collapse to the best hit per value identity.
33    pub distinct: Option<&'a kevy_text::Distinct<'a>>,
34    /// `FACET` fields to count over the (filtered) match set.
35    pub facets: &'a [kevy_text::Facet<'a>],
36    /// How deep a page the merge needs (LIMIT + OFFSET).
37    pub fetch: usize,
38}
39
40/// One cold hit: its page-order ingredients, ready to merge.
41pub struct ColdHit {
42    pub key: Vec<u8>,
43    pub score: f64,
44    /// The sort key, when the query sorts by a stored value.
45    pub okey: Option<Vec<u8>>,
46}
47
48/// The cold half of one shard's pass-2 answer.
49pub struct ColdPage {
50    /// Best `fetch` cold hits in the page's order.
51    pub hits: Vec<ColdHit>,
52    /// The returned hits' frozen stored values — what the merge reads
53    /// for sort/distinct identities and the origin's okeys/dkeys.
54    pub values: HashMap<Vec<u8>, Vec<Option<Vec<u8>>>>,
55    /// Per requested facet field, (identity, label, count) over the
56    /// filtered cold match set.
57    pub facets: Vec<Vec<kevy_text::Bucket>>,
58}
59
60impl TextColdDir {
61    /// Pass-1 contribution: summed LIVE docs/length plus per-token
62    /// live df across every cold segment (one fence descent per token
63    /// per segment; the doc/length halves are in-memory numbers, no
64    /// I/O at all).
65    pub fn cold_stats(&self, tokens: &[Vec<u8>]) -> (u64, u64, Vec<(Vec<u8>, u32)>) {
66        let n_docs: u64 = self.segs.iter().map(|c| c.n_docs).sum();
67        let total_len: u64 = self.segs.iter().map(|c| c.total_len).sum();
68        let df = tokens
69            .iter()
70            .map(|t| {
71                let frozen: u32 = self
72                    .segs
73                    .iter()
74                    .filter_map(|c| c.seg.get(t).ok().flatten())
75                    .filter_map(|p| posting_df(&p))
76                    .sum();
77                let dead = self.df_dead.get(t).copied().unwrap_or(0);
78                (t.clone(), frozen.saturating_sub(dead))
79            })
80            .collect();
81        (n_docs, total_len, df)
82    }
83
84    /// Pass-2 contribution: the clause-faithful cold page (see the
85    /// module doc for what each clause does here).
86    pub fn cold_page(&self, q: &ColdPageQuery) -> ColdPage {
87        let acc = self.accumulate(q);
88        let need_values = !q.filter.is_empty()
89            || q.sort.is_some()
90            || q.distinct.is_some()
91            || !q.facets.is_empty();
92        let mut values: HashMap<Vec<u8>, Vec<Option<Vec<u8>>>> = HashMap::new();
93        let mut cands: Vec<ColdHit> = Vec::new();
94        for (key, score) in acc {
95            let vals = if need_values {
96                let Some(v) = self.frozen_values(&key) else { continue };
97                if !passes(&v, q.filter) {
98                    continue;
99                }
100                Some(v)
101            } else {
102                None
103            };
104            let okey = q.sort.and_then(|s| {
105                vals.as_ref()?.get(s.field)?.as_deref().and_then(s.key)
106            });
107            if let Some(v) = vals {
108                values.insert(key.clone(), v);
109            }
110            cands.push(ColdHit { key, score, okey });
111        }
112        let facets = self.count_facets(q, &cands, &values);
113        order_page(&mut cands, q.sort.is_some(), q.sort.is_some_and(|s| s.desc));
114        if let Some(d) = q.distinct {
115            collapse(&mut cands, d, &values);
116        }
117        cands.truncate(q.fetch);
118        values.retain(|k, _| cands.iter().any(|c| &c.key == k));
119        ColdPage { hits: cands, values, facets }
120    }
121
122    /// Every clause's accumulated cold score, keyed by row key — the
123    /// mirror of the hot `accumulate_clauses` over the frozen postings.
124    fn accumulate(&self, q: &ColdPageQuery) -> HashMap<Vec<u8>, f64> {
125        let mut acc = HashMap::new();
126        for cs in &self.segs {
127            let dead = |k: &[u8]| self.tombs.get(k).is_some_and(|s| s.contains(&cs.seq));
128            for t in &q.bare {
129                if let Ok(Some(payload)) = cs.seg.get(t) {
130                    let _ = score_cold(&payload, t, q.stats, &dead, &mut acc);
131                }
132            }
133            for phrase in &q.phrases {
134                let payloads: Option<Vec<Vec<u8>>> =
135                    phrase.iter().map(|t| cs.seg.get(t).ok().flatten()).collect();
136                // A phrase token absent from this segment = the phrase
137                // matches nothing here (the rarest-anchor None mirror).
138                if let Some(payloads) = payloads {
139                    let _ = score_cold_phrase(&payloads, phrase, q.stats, &dead, &mut acc);
140                }
141            }
142        }
143        acc
144    }
145
146    /// One live cold document's frozen stored values, from whichever
147    /// segment holds its un-shadowed copy.
148    fn frozen_values(&self, key: &[u8]) -> Option<Vec<Option<Vec<u8>>>> {
149        let mut fwd_key = vec![0u8];
150        fwd_key.extend_from_slice(key);
151        for cs in &self.segs {
152            if self.tombs.get(key).is_some_and(|s| s.contains(&cs.seq)) {
153                continue;
154            }
155            if let Ok(Some(payload)) = cs.seg.get(&fwd_key) {
156                return decode_fwd(&payload).map(|r| r.values);
157            }
158        }
159        None
160    }
161
162    /// The cold half of each facet's count — the hot `count_facet`'s
163    /// rules (filter applies, top-K and DISTINCT do not), ordered the
164    /// same way; the shard merge sums it with the hot half.
165    fn count_facets(
166        &self,
167        q: &ColdPageQuery,
168        cands: &[ColdHit],
169        values: &HashMap<Vec<u8>, Vec<Option<Vec<u8>>>>,
170    ) -> Vec<Vec<kevy_text::Bucket>> {
171        q.facets
172            .iter()
173            .map(|f| {
174                let mut counts: HashMap<Vec<u8>, (Vec<u8>, u64)> = HashMap::new();
175                for c in cands {
176                    let Some(raw) =
177                        values.get(&c.key).and_then(|v| v.get(f.field)).and_then(Option::as_deref)
178                    else {
179                        continue;
180                    };
181                    let Some(k) = (f.key)(raw) else { continue };
182                    counts.entry(k).or_insert_with(|| (raw.to_vec(), 0)).1 += 1;
183                }
184                let mut out: Vec<kevy_text::Bucket> =
185                    counts.into_iter().map(|(k, (label, n))| (k, label, n)).collect();
186                out.sort_by(|a, b| b.2.cmp(&a.2).then_with(|| a.1.cmp(&b.1)));
187                out
188            })
189            .collect()
190    }
191}
192
193/// Order candidates by the page's rule: `sorted_order` under SORT
194/// (a document WITH a value outranks one without, in both
195/// directions), else score-descending with the row key as tiebreak.
196fn order_page(cands: &mut [ColdHit], sorted: bool, desc: bool) {
197    if sorted {
198        cands.sort_by(|a, b| {
199            sorted_order((a.okey.as_deref(), &a.key), (b.okey.as_deref(), &b.key), desc)
200        });
201    } else {
202        cands.sort_by(|a, b| {
203            b.score
204                .partial_cmp(&a.score)
205                .unwrap_or(std::cmp::Ordering::Equal)
206                .then_with(|| a.key.cmp(&b.key))
207        });
208    }
209}
210
211/// Collapse an ordered candidate list to the best hit per distinct
212/// identity. A document with no value for the field is its own group
213/// (the hot rule: it has not been shown to share anything).
214fn collapse(
215    cands: &mut Vec<ColdHit>,
216    d: &kevy_text::Distinct,
217    values: &HashMap<Vec<u8>, Vec<Option<Vec<u8>>>>,
218) {
219    let mut seen: std::collections::HashSet<Vec<u8>> = std::collections::HashSet::new();
220    cands.retain(|c| {
221        let identity = values
222            .get(&c.key)
223            .and_then(|v| v.get(d.field))
224            .and_then(Option::as_deref)
225            .and_then(d.key);
226        match identity {
227            None => true,
228            Some(id) => seen.insert(id),
229        }
230    });
231}
232
233/// The hot `passes` mirror over frozen values: every predicate must
234/// pass, and an absent value never does.
235fn passes(values: &[Option<Vec<u8>>], filter: &[kevy_text::Filter]) -> bool {
236    filter.iter().all(|f| {
237        values.get(f.field).and_then(Option::as_deref).is_some_and(|v| (f.test)(v))
238    })
239}