kevy_text/segment_opts.rs
1//! What a query asks for, beyond its text and result limit — the clause
2//! options and the shapes they answer in. A child module of `segment`
3//! (declared via `#[path]`), re-exported from it, so the crate's public
4//! surface is unchanged by the split.
5
6use std::collections::HashMap;
7
8use super::{Distinct, Filter, Sort};
9
10/// One ranked hit.
11///
12/// ```
13/// let m = kevy_text::TextMatch::new(b"doc:1".to_vec(), 1.5);
14/// assert_eq!((m.key.as_slice(), m.score), (&b"doc:1"[..], 1.5));
15/// ```
16#[derive(Debug, Clone, PartialEq, Default)]
17#[non_exhaustive]
18pub struct TextMatch {
19 /// Row key.
20 ///
21 /// ```
22 /// let mut seg = kevy_text::TextSegment::new();
23 /// seg.apply(b"user:7", Some(b"likes rust"));
24 /// assert_eq!(seg.matches(b"rust", 10)[0].key, b"user:7");
25 /// ```
26 pub key: Vec<u8>,
27 /// Shard-local BM25 score.
28 ///
29 /// Higher is better; the page comes back best first.
30 ///
31 /// ```
32 /// let mut seg = kevy_text::TextSegment::new();
33 /// seg.apply(b"short", Some(b"rust"));
34 /// seg.apply(b"long", Some(b"rust in a much longer document"));
35 /// let hits = seg.matches(b"rust", 10);
36 /// assert_eq!(hits[0].key, b"short");
37 /// assert!(hits[0].score > hits[1].score && hits[1].score > 0.0);
38 /// ```
39 pub score: f64,
40}
41
42impl TextMatch {
43 /// A hit for `key` at `score` — how a merge of several segments'
44 /// pages, or of hot and cold hits, reports its union.
45 ///
46 /// ```
47 /// let m = kevy_text::TextMatch::new(b"k".to_vec(), 0.25);
48 /// assert_eq!(m.score, 0.25);
49 /// ```
50 #[must_use]
51 pub fn new(key: Vec<u8>, score: f64) -> Self {
52 Self { key, score }
53 }
54}
55
56/// Corpus statistics supplied from outside a segment, for scoring one
57/// shard's documents against the whole corpus rather than its own slice.
58///
59/// A cross-shard text query builds this by summing each shard's local
60/// `n_docs` / `total_len` and, for each query token, its `df`. `df`
61/// need only carry the query's tokens — the values a query actually
62/// scores with — which is why global BM25 does not need a whole-corpus
63/// df table.
64///
65/// ```
66/// use std::collections::HashMap;
67/// let stats = kevy_text::CorpusStats::new(10.0, 4.0, HashMap::from([(b"rust".to_vec(), 3)]));
68/// assert_eq!(stats.df.get(&b"rust"[..]), Some(&3));
69/// ```
70#[derive(Debug, Clone, PartialEq, Default)]
71#[non_exhaustive]
72pub struct CorpusStats {
73 /// Total documents across the corpus.
74 ///
75 /// One match among a thousand documents is rarer, so it scores higher,
76 /// than the same match among two:
77 ///
78 /// ```
79 /// use kevy_text::{CorpusStats, TextSegment};
80 /// use std::collections::HashMap;
81 /// let mut seg = TextSegment::new();
82 /// seg.apply(b"doc:1", Some(b"rust engine"));
83 /// let df = HashMap::from([(b"rust".to_vec(), 1)]);
84 /// let small = CorpusStats::new(2.0, 2.0, df.clone());
85 /// let large = CorpusStats::new(1000.0, 2.0, df);
86 /// let score = |s| seg.matches_scored(b"rust", 1, Some(s))[0].score;
87 /// assert!(score(&large) > score(&small));
88 /// ```
89 pub n_docs: f64,
90 /// Mean document length (unweighted tokens) across the corpus.
91 ///
92 /// The same document is short against a corpus of long documents,
93 /// and BM25 rewards a match in a short one:
94 ///
95 /// ```
96 /// use kevy_text::{CorpusStats, TextSegment};
97 /// use std::collections::HashMap;
98 /// let mut seg = TextSegment::new();
99 /// seg.apply(b"doc:1", Some(b"rust engine"));
100 /// let df = HashMap::from([(b"rust".to_vec(), 1)]);
101 /// let terse = CorpusStats::new(10.0, 2.0, df.clone());
102 /// let verbose = CorpusStats::new(10.0, 50.0, df);
103 /// let score = |s| seg.matches_scored(b"rust", 1, Some(s))[0].score;
104 /// assert!(score(&verbose) > score(&terse));
105 /// ```
106 pub avgdl: f64,
107 /// Global document frequency per query token; a token missing here
108 /// falls back to the segment's local list length.
109 ///
110 /// ```
111 /// use kevy_text::{CorpusStats, TextSegment};
112 /// use std::collections::HashMap;
113 /// let mut seg = TextSegment::new();
114 /// seg.apply(b"doc:1", Some(b"rust engine"));
115 /// let rare = CorpusStats::new(100.0, 2.0, HashMap::from([(b"rust".to_vec(), 1)]));
116 /// let common = CorpusStats::new(100.0, 2.0, HashMap::from([(b"rust".to_vec(), 90)]));
117 /// let score = |s| seg.matches_scored(b"rust", 1, Some(s))[0].score;
118 /// assert!(score(&rare) > score(&common));
119 /// // absent from `df`: this segment's own count (1) stands in
120 /// let fallback = CorpusStats::new(100.0, 2.0, HashMap::new());
121 /// assert_eq!(score(&fallback), score(&rare));
122 /// ```
123 pub df: HashMap<Vec<u8>, u32>,
124}
125
126impl CorpusStats {
127 /// Statistics for a corpus of `n_docs` documents of mean length
128 /// `avgdl`, with the document frequency of each query token.
129 ///
130 /// ```
131 /// let s = kevy_text::CorpusStats::new(2.0, 3.5, Default::default());
132 /// assert_eq!((s.n_docs, s.avgdl), (2.0, 3.5));
133 /// ```
134 #[must_use]
135 pub fn new(n_docs: f64, avgdl: f64, df: HashMap<Vec<u8>, u32>) -> Self {
136 Self { n_docs, avgdl, df }
137 }
138}
139
140/// Which way an order runs: ascending or descending.
141///
142/// ```
143/// use std::cmp::Ordering;
144/// use kevy_text::SortOrder;
145/// assert_eq!(SortOrder::Asc.apply(1.cmp(&2)), Ordering::Less);
146/// assert_eq!(SortOrder::Desc.apply(1.cmp(&2)), Ordering::Greater);
147/// assert_eq!(SortOrder::default(), SortOrder::Asc);
148/// ```
149#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
150pub enum SortOrder {
151 /// Smallest first.
152 ///
153 /// ```
154 /// # use kevy_text::{QueryOpts, SegmentShape, Sort, SortOrder, TextSegment};
155 /// # let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
156 /// # for (k, text, price) in [("a", "red shoe", "30"), ("b", "red hat", "10"), ("c", "red cap", "20")] {
157 /// # seg.apply_doc(k.as_bytes(), Some(&[(text.as_bytes().to_vec(), 1.0)]), &[Some(price.as_bytes())]);
158 /// # }
159 /// // a "red shoe" 30, b "red hat" 10, c "red cap" 20
160 /// let price = |v: &[u8]| Some(v.to_vec());
161 /// let cheapest = QueryOpts::default().with_sort(Sort::new(0, &price).with_order(SortOrder::Asc));
162 /// let page = seg.matches_query_with(b"red", 2, cheapest);
163 /// assert_eq!(page.iter().map(|m| &m.key[..]).collect::<Vec<_>>(), [&b"b"[..], b"c"]);
164 /// ```
165 #[default]
166 Asc,
167 /// Largest first.
168 ///
169 /// ```
170 /// # use kevy_text::{QueryOpts, SegmentShape, Sort, SortOrder, TextSegment};
171 /// # let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
172 /// # for (k, text, price) in [("a", "red shoe", "30"), ("b", "red hat", "10"), ("c", "red cap", "20")] {
173 /// # seg.apply_doc(k.as_bytes(), Some(&[(text.as_bytes().to_vec(), 1.0)]), &[Some(price.as_bytes())]);
174 /// # }
175 /// // a "red shoe" 30, b "red hat" 10, c "red cap" 20
176 /// let price = |v: &[u8]| Some(v.to_vec());
177 /// let dearest = QueryOpts::default().with_sort(Sort::new(0, &price).with_order(SortOrder::Desc));
178 /// let page = seg.matches_query_with(b"red", 2, dearest);
179 /// assert_eq!(page.iter().map(|m| &m.key[..]).collect::<Vec<_>>(), [&b"a"[..], b"c"]);
180 /// ```
181 Desc,
182}
183
184impl SortOrder {
185 /// `ord` as this direction sees it: unchanged ascending, reversed
186 /// descending.
187 ///
188 /// ```
189 /// use kevy_text::SortOrder;
190 /// let mut v = [3, 1, 2];
191 /// v.sort_by(|a, b| SortOrder::Desc.apply(a.cmp(b)));
192 /// assert_eq!(v, [3, 2, 1]);
193 /// ```
194 #[must_use]
195 pub fn apply(self, ord: std::cmp::Ordering) -> std::cmp::Ordering {
196 match self {
197 Self::Asc => ord,
198 Self::Desc => ord.reverse(),
199 }
200 }
201}
202
203/// One value bucket: the identity a cross-shard merge sums by, a spelling
204/// of it that occurs in the corpus, and how many documents matched with
205/// it.
206///
207/// ```
208/// use kevy_text::{Bucket, Facet, QueryOpts, SegmentShape, TextSegment};
209/// let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
210/// seg.apply_doc(b"a", Some(&[(b"red shoe".to_vec(), 1.0)]), &[Some(b"ACME")]);
211/// seg.apply_doc(b"b", Some(&[(b"red hat".to_vec(), 1.0)]), &[Some(b"acme")]);
212/// // identity folds case; the label is a spelling that really occurs
213/// let brand = |v: &[u8]| Some(v.to_ascii_lowercase());
214/// let r = seg.matches_query_faceted(b"red", 10, QueryOpts::default(), &[Facet::new(0, &brand)]);
215/// let (identity, label, count): &Bucket = &r.facets[0][0];
216/// assert_eq!((&identity[..], count), (&b"acme"[..], &2));
217/// assert!(label == b"ACME" || label == b"acme");
218/// ```
219pub type Bucket = (Vec<u8>, Vec<u8>, u64);
220
221/// One faceted query's answer: the page, and a count per value for each
222/// requested field.
223///
224/// ```
225/// use kevy_text::{Facet, QueryOpts, SegmentShape, TextSegment};
226/// let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
227/// for (k, text, color) in [("a", "wool hat", "red"), ("b", "wool scarf", "red"), ("c", "wool sock", "blue")] {
228/// seg.apply_doc(k.as_bytes(), Some(&[(text.as_bytes().to_vec(), 1.0)]), &[Some(color.as_bytes())]);
229/// }
230/// let color = |v: &[u8]| Some(v.to_vec());
231/// let r = seg.matches_query_faceted(b"wool", 1, QueryOpts::default(), &[Facet::new(0, &color)]);
232/// // one hit on the page, but the counts cover all three matches
233/// assert_eq!(r.hits.len(), 1);
234/// assert_eq!(r.facets[0], [(b"red".to_vec(), b"red".to_vec(), 2), (b"blue".to_vec(), b"blue".to_vec(), 1)]);
235/// ```
236#[derive(Debug, Clone, PartialEq, Default)]
237#[non_exhaustive]
238pub struct FacetedMatches {
239 /// The ranked page, exactly what an unfaceted query would return.
240 ///
241 /// ```
242 /// use kevy_text::{Facet, QueryOpts, SegmentShape, TextSegment};
243 /// let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
244 /// seg.apply_doc(b"a", Some(&[(b"wool hat".to_vec(), 1.0)]), &[Some(b"red")]);
245 /// seg.apply_doc(b"b", Some(&[(b"wool wool sock".to_vec(), 1.0)]), &[Some(b"blue")]);
246 /// let color = |v: &[u8]| Some(v.to_vec());
247 /// let r = seg.matches_query_faceted(b"wool", 10, QueryOpts::default(), &[Facet::new(0, &color)]);
248 /// assert_eq!(r.hits, seg.matches_query_with(b"wool", 10, QueryOpts::default()));
249 /// ```
250 pub hits: Vec<TextMatch>,
251 /// Per requested facet field, `(identity, label, count)` over the
252 /// whole match set. The identity is what a cross-shard merge sums by;
253 /// the label is what it reports.
254 ///
255 /// One list per requested facet, in request order, most frequent
256 /// value first:
257 ///
258 /// ```
259 /// use kevy_text::{Facet, QueryOpts, SegmentShape, TextSegment};
260 /// let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(2));
261 /// seg.apply_doc(b"a", Some(&[(b"wool hat".to_vec(), 1.0)]), &[Some(b"red"), Some(b"S")]);
262 /// seg.apply_doc(b"b", Some(&[(b"wool sock".to_vec(), 1.0)]), &[Some(b"blue"), Some(b"S")]);
263 /// let raw = |v: &[u8]| Some(v.to_vec());
264 /// let facets = [Facet::new(1, &raw), Facet::new(0, &raw)];
265 /// let r = seg.matches_query_faceted(b"wool", 10, QueryOpts::default(), &facets);
266 /// assert_eq!(r.facets.len(), 2);
267 /// assert_eq!(r.facets[0], [(b"S".to_vec(), b"S".to_vec(), 2)]);
268 /// assert_eq!(r.facets[1].len(), 2);
269 /// ```
270 pub facets: Vec<Vec<Bucket>>,
271}
272
273/// Everything a MATCH query carries beyond its text and result limit.
274///
275/// Grouping them keeps the query entry point from growing a parameter per
276/// clause, and gives every clause one place to be defaulted from
277/// ([`QueryOpts::default`] is the plain, exact, unscoped query). Set what
278/// differs with the `with_*` builders or by assigning the fields.
279///
280/// ```
281/// let q = kevy_text::QueryOpts::default().with_typo(1).with_fields(&[0]);
282/// assert_eq!((q.typo, q.fields), (1, &[0usize][..]));
283/// ```
284#[derive(Debug, Clone, Copy, Default)]
285#[non_exhaustive]
286pub struct QueryOpts<'a> {
287 /// Corpus-wide BM25 statistics — the second pass of a cross-shard
288 /// query. `None` scores against this segment's own slice.
289 ///
290 /// ```
291 /// use kevy_text::{CorpusStats, QueryOpts, TextSegment};
292 /// use std::collections::HashMap;
293 /// let mut seg = TextSegment::new();
294 /// seg.apply(b"doc:1", Some(b"rust engine"));
295 /// let corpus = CorpusStats::new(50.0, 4.0, HashMap::from([(b"rust".to_vec(), 5)]));
296 /// let mut opts = QueryOpts::default();
297 /// opts.stats = Some(&corpus);
298 /// let global = seg.matches_query_with(b"rust", 1, opts)[0].score;
299 /// assert_eq!(global, seg.matches_scored(b"rust", 1, Some(&corpus))[0].score);
300 /// assert_ne!(global, seg.matches_query_with(b"rust", 1, QueryOpts::default())[0].score);
301 /// ```
302 pub stats: Option<&'a CorpusStats>,
303 /// Edit distance allowed on bare terms (`TYPO n`); 0 = exact.
304 ///
305 /// ```
306 /// use kevy_text::{QueryOpts, TextSegment};
307 /// let mut seg = TextSegment::new();
308 /// seg.apply(b"doc:1", Some(b"kevy storage"));
309 /// assert!(seg.matches_query_with(b"storge", 10, QueryOpts::default()).is_empty());
310 /// let mut fuzzy = QueryOpts::default();
311 /// fuzzy.typo = 1;
312 /// assert_eq!(seg.matches_query_with(b"storge", 10, fuzzy)[0].key, b"doc:1");
313 /// ```
314 pub typo: u32,
315 /// Field positions the query is restricted to (`IN <field…>`); empty
316 /// = every field.
317 ///
318 /// ```
319 /// use kevy_text::{QueryOpts, SegmentShape, TextSegment};
320 /// let mut seg = TextSegment::with_shape(SegmentShape::default().with_fields(2));
321 /// // field 0 = title, field 1 = body
322 /// seg.apply_fields(b"doc:1", Some(&[(b"kevy".to_vec(), 1.0), (b"about redis".to_vec(), 1.0)]));
323 /// assert_eq!(seg.matches_query_with(b"redis", 10, QueryOpts::default()).len(), 1);
324 /// let mut titles = QueryOpts::default();
325 /// titles.fields = &[0];
326 /// assert!(seg.matches_query_with(b"redis", 10, titles).is_empty());
327 /// ```
328 pub fields: &'a [usize],
329 /// `FILTER`: non-scoring predicates, ANDed. Applied before the top-K
330 /// — filtering afterwards would return fewer hits than exist.
331 ///
332 /// ```
333 /// # use kevy_text::{Filter, QueryOpts, SegmentShape, TextSegment};
334 /// # let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
335 /// # for (k, text, price) in [("a", "red shoe", "30"), ("b", "red hat", "10"), ("c", "red cap", "20")] {
336 /// # seg.apply_doc(k.as_bytes(), Some(&[(text.as_bytes().to_vec(), 1.0)]), &[Some(price.as_bytes())]);
337 /// # }
338 /// // a "red shoe" 30, b "red hat" 10, c "red cap" 20
339 /// let under_25 = |v: &[u8]| v < &b"25"[..];
340 /// let filter = [Filter::new(0, &under_25)];
341 /// let mut opts = QueryOpts::default();
342 /// opts.filter = &filter;
343 /// let mut keys: Vec<_> = seg.matches_query_with(b"red", 10, opts).into_iter().map(|m| m.key).collect();
344 /// keys.sort();
345 /// assert_eq!(keys, [b"b".to_vec(), b"c".to_vec()]);
346 /// ```
347 pub filter: &'a [Filter<'a>],
348 /// `SORT`: select by a stored value instead of by score. Selecting,
349 /// not re-ordering: a document that wins on the sort key must be
350 /// chosen even when its score would never have reached the page.
351 ///
352 /// ```
353 /// use kevy_text::{QueryOpts, SegmentShape, Sort, TextSegment};
354 /// let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
355 /// seg.apply_doc(b"best", Some(&[(b"red red red".to_vec(), 1.0)]), &[Some(b"2024")]);
356 /// seg.apply_doc(b"oldest", Some(&[(b"red and many other words".to_vec(), 1.0)]), &[Some(b"1999")]);
357 /// assert_eq!(seg.matches_query_with(b"red", 1, QueryOpts::default())[0].key, b"best");
358 /// let year = |v: &[u8]| Some(v.to_vec());
359 /// let mut opts = QueryOpts::default();
360 /// opts.sort = Some(Sort::new(0, &year));
361 /// assert_eq!(seg.matches_query_with(b"red", 1, opts)[0].key, b"oldest");
362 /// ```
363 pub sort: Option<Sort<'a>>,
364 /// `DISTINCT`: at most one hit per value of a stored field. Applied
365 /// during selection, so the page is filled with `limit` DISTINCT
366 /// documents rather than `limit` documents that then collapse.
367 ///
368 /// ```
369 /// use kevy_text::{Distinct, QueryOpts, SegmentShape, TextSegment};
370 /// let mut seg = TextSegment::with_shape(SegmentShape::default().with_values(1));
371 /// seg.apply_doc(b"a1", Some(&[(b"red red".to_vec(), 1.0)]), &[Some(b"acme")]);
372 /// seg.apply_doc(b"a2", Some(&[(b"red shoe".to_vec(), 1.0)]), &[Some(b"acme")]);
373 /// seg.apply_doc(b"z1", Some(&[(b"red hat and scarf".to_vec(), 1.0)]), &[Some(b"zeta")]);
374 /// let brand = |v: &[u8]| Some(v.to_vec());
375 /// let mut opts = QueryOpts::default();
376 /// opts.distinct = Some(Distinct::new(0, &brand));
377 /// // two slots, filled by two different brands
378 /// let page = seg.matches_query_with(b"red", 2, opts);
379 /// assert_eq!(page.iter().map(|m| &m.key[..]).collect::<Vec<_>>(), [&b"a1"[..], b"z1"]);
380 /// ```
381 pub distinct: Option<Distinct<'a>>,
382}
383
384impl<'a> QueryOpts<'a> {
385 /// Score against corpus-wide `stats` rather than this segment's slice.
386 ///
387 /// ```
388 /// let s = kevy_text::CorpusStats::default();
389 /// assert!(kevy_text::QueryOpts::default().with_stats(&s).stats.is_some());
390 /// ```
391 #[must_use]
392 pub fn with_stats(mut self, stats: &'a CorpusStats) -> Self {
393 self.stats = Some(stats);
394 self
395 }
396
397 /// Allow `typo` edits on bare terms.
398 ///
399 /// ```
400 /// assert_eq!(kevy_text::QueryOpts::default().with_typo(2).typo, 2);
401 /// ```
402 #[must_use]
403 pub fn with_typo(mut self, typo: u32) -> Self {
404 self.typo = typo;
405 self
406 }
407
408 /// Restrict the query to these field positions.
409 ///
410 /// ```
411 /// assert_eq!(kevy_text::QueryOpts::default().with_fields(&[1, 2]).fields, &[1, 2]);
412 /// ```
413 #[must_use]
414 pub fn with_fields(mut self, fields: &'a [usize]) -> Self {
415 self.fields = fields;
416 self
417 }
418
419 /// Keep only documents every one of `filter` passes.
420 ///
421 /// ```
422 /// let any = |_: &[u8]| true;
423 /// let f = [kevy_text::Filter::new(0, &any)];
424 /// assert_eq!(kevy_text::QueryOpts::default().with_filter(&f).filter.len(), 1);
425 /// ```
426 #[must_use]
427 pub fn with_filter(mut self, filter: &'a [Filter<'a>]) -> Self {
428 self.filter = filter;
429 self
430 }
431
432 /// Select by a stored value instead of by score.
433 ///
434 /// ```
435 /// let key = |v: &[u8]| Some(v.to_vec());
436 /// let q = kevy_text::QueryOpts::default().with_sort(kevy_text::Sort::new(0, &key));
437 /// assert!(q.sort.is_some());
438 /// ```
439 #[must_use]
440 pub fn with_sort(mut self, sort: Sort<'a>) -> Self {
441 self.sort = Some(sort);
442 self
443 }
444
445 /// At most one hit per value of a stored field.
446 ///
447 /// ```
448 /// let key = |v: &[u8]| Some(v.to_vec());
449 /// let q = kevy_text::QueryOpts::default().with_distinct(kevy_text::Distinct::new(0, &key));
450 /// assert!(q.distinct.is_some());
451 /// ```
452 #[must_use]
453 pub fn with_distinct(mut self, distinct: Distinct<'a>) -> Self {
454 self.distinct = Some(distinct);
455 self
456 }
457}