kevy_window/
text_query.rs1use 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
20pub struct ColdPageQuery<'a> {
22 pub bare: Vec<Vec<u8>>,
24 pub phrases: Vec<Vec<Vec<u8>>>,
26 pub stats: &'a CorpusStats,
28 pub filter: &'a [kevy_text::Filter<'a>],
30 pub sort: Option<&'a kevy_text::Sort<'a>>,
32 pub distinct: Option<&'a kevy_text::Distinct<'a>>,
34 pub facets: &'a [kevy_text::Facet<'a>],
36 pub fetch: usize,
38}
39
40pub struct ColdHit {
42 pub key: Vec<u8>,
43 pub score: f64,
44 pub okey: Option<Vec<u8>>,
46}
47
48pub struct ColdPage {
50 pub hits: Vec<ColdHit>,
52 pub values: HashMap<Vec<u8>, Vec<Option<Vec<u8>>>>,
55 pub facets: Vec<Vec<kevy_text::Bucket>>,
58}
59
60impl TextColdDir {
61 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 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 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 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 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 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
193fn 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
211fn 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
233fn 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}