Skip to main content

zeppelin_embed/fts/prune/
mod.rs

1//! Dynamic pruning: block-max MAXSCORE and WAND, counter-gated.
2//!
3//! # What this buys, and what it must never cost
4//!
5//! Skipping is the whole wall-clock gap between a naive engine and
6//! Lucene/tantivy: documents evaluated drop from 3.8M to 21.9K per GOV2
7//! query, and latency from 225.7 ms to 27.9 ms (`research/02a:281`). Those
8//! are x86 numbers on 25M documents; the ratio is the transferable part,
9//! not the constants.
10//!
11//! What it must never cost is a single changed result. Every function here
12//! decides only *which documents get scored*. Both paths read merged postings
13//! through [`crate::fts::sealed::TermStream`] and score them with
14//! [`crate::fts::bm25::TermScorer::score`], so `prop_pruned_topk_equals_
15//! exhaustive` is a statement about skipping decisions rather than about two
16//! scorers agreeing by luck.
17//!
18//! **If a bound is uncertain, evaluate.** A slow answer is a bug report; a
19//! wrong answer is a lost user.
20//!
21//! # Strategy selection
22//!
23//! The current counter-derived rule selects block-max MAXSCORE for three terms
24//! or fewer and block-max WAND for four or more, independent of `k`. See
25//! [`select_strategy`](crate::fts::prune::select_strategy) for the measurements
26//! and recalibration history and `docs/14-pruning-contracts.md` for the
27//! executable contract.
28
29pub mod bounds;
30pub mod maxscore;
31pub mod wand;
32pub(crate) mod weighted;
33
34use crate::fts::bm25::{Bm25Params, Df, TermScorer};
35use crate::fts::index::{IndexError, LexicalIndex};
36use crate::fts::sealed::TermStream;
37use crate::fts::search::ControlledSearchError;
38use crate::fts::search::{GlobalDocId, ScoredDoc, SearchCounters, SearchResult, TermQuery};
39use crate::meta::DocBitmap;
40
41pub use bounds::{BlockBound, build_block_bounds, impact_of, term_upper_bound};
42
43/// Postings per pruning block.
44///
45/// Matches the persisted geometry so the in-memory bounds and a sealed
46/// segment's stored maxima describe the same blocks.
47pub const PRUNE_BLOCK_SIZE: usize = 64;
48
49/// Lists shorter than one block skip the pruning machinery entirely.
50///
51/// This is the dominant case in small corpora (`research/02a:328`). It is a
52/// speed constant, re-tunable from counters, not a correctness threshold:
53/// the fast path scores the same documents the pruned path would.
54pub const SHORT_LIST_POSTINGS: usize = PRUNE_BLOCK_SIZE;
55
56/// Which pruning strategy ran.
57#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
58pub enum Strategy {
59    /// Exhaustive OR, no pruning. The oracle.
60    Exhaustive,
61    /// Block-max WAND: pivot selection over sorted cursors.
62    #[default]
63    BlockMaxWand,
64    /// Block-max MAXSCORE: essential/non-essential partition.
65    BlockMaxMaxscore,
66}
67
68/// Term count at and above which the rule prefers WAND. Three terms or
69/// fewer go to MAXSCORE; see [`select_strategy`] for the measurement.
70pub const MIN_TERMS_FOR_WAND: usize = 4;
71
72/// Returns the strategy the measured rule selects.
73///
74/// **MAXSCORE at three terms or fewer, WAND at four or more, at any `k`.**
75///
76/// # Provenance: measured here, and re-derived when the engine changed
77///
78/// The original rule — WAND when `terms <= 4 && k <= 10` — came from GOV2
79/// on x86 at 25 million documents (`research/02a:282`). Recalibrating it
80/// on this engine's own counters first moved the boundary to two terms;
81/// fixing MAXSCORE's probe order (descending by bound, with a
82/// metadata-only block bound tested before every seek) then made MAXSCORE
83/// uniformly cheaper and moved the crossover again, from three terms to
84/// four. This is exactly why
85/// `strategy_rule_matches_the_counters_it_was_calibrated_from` re-derives
86/// the rule on every run rather than trusting a number in a comment.
87///
88/// What the counters say now, at 100,000 documents:
89///
90/// - **At two and three terms MAXSCORE wins**: 23,594 cost units against
91///   WAND's 26,819 at three terms, `k` 10. With few cursors the pivot has
92///   little to choose between, and MAXSCORE no longer pays a seek per
93///   cursor per candidate.
94/// - **At four terms and above WAND wins** at the corpus sizes that
95///   matter, on the posting counter that is the inner loop.
96/// - **`k` still does not discriminate at scale.** The 2,000-document
97///   corpus at `k` of 100 prefers MAXSCORE at every term count, but both
98///   strategies cost about 2,000 units there — sub-millisecond either way
99///   — and a rule fitted to that regime would misroute the large corpora
100///   where the choice is worth something.
101///
102/// **Both strategies return identical results** — `prune_equivalence`
103/// gates that — so this is a cost decision with no quality risk.
104///
105/// **Owed:** confirmation on BEIR. This corpus is synthetic, with uniform
106/// term frequencies and lengths. Replacing a constant measured on someone
107/// else's hardware and collection with one measured on ours is a strict
108/// improvement, but it is not the same as measuring the real thing.
109#[must_use]
110pub const fn select_strategy(term_count: usize, k: usize) -> Strategy {
111    // `k` is accepted and deliberately unused: the counters say it does not
112    // discriminate, and dropping the parameter would be a breaking API
113    // change made for a cosmetic reason.
114    let _ = k;
115    if term_count >= MIN_TERMS_FOR_WAND {
116        Strategy::BlockMaxWand
117    } else {
118        Strategy::BlockMaxMaxscore
119    }
120}
121
122/// One query term prepared for pruning within one segment.
123///
124/// The cursor is a view over the sealed streams, not a copy of them. It
125/// holds no entry array: `advance` and `seek` walk block metadata and
126/// decode only the blocks the traversal actually lands in.
127#[derive(Clone, Debug)]
128pub struct TermCursor<'segment> {
129    /// The merged sealed postings for this term across the weighted fields.
130    pub stream: TermStream<'segment>,
131    /// Store-wide document frequency.
132    pub df: Df,
133    /// This term's hoisted scoring constants.
134    ///
135    /// Built once per (term, segment) so the inner loops never recompute
136    /// `idf` or `avgdl`. See [`crate::fts::bm25::TermScorer`].
137    pub scorer: TermScorer,
138    /// Largest contribution this term can make to any document.
139    pub upper_bound: f64,
140}
141
142impl TermCursor<'_> {
143    /// Returns the row at the cursor, or `None` when exhausted.
144    #[must_use]
145    pub fn current(&self) -> Option<u32> {
146        self.stream.current_row()
147    }
148
149    /// Returns the merged term frequency at the cursor.
150    ///
151    /// The first ask against a block pays that block's tf decode; rows
152    /// that are skipped rather than scored never pay it.
153    #[must_use]
154    pub fn current_tf(&mut self) -> Option<u32> {
155        self.stream.current_tf()
156    }
157
158    /// Advances one merged posting.
159    pub fn advance(&mut self) {
160        self.stream.advance();
161    }
162
163    /// Rewinds to the first merged posting.
164    pub fn reset(&mut self) {
165        self.stream.reset();
166    }
167
168    /// Advances to the first row at or after `row`.
169    ///
170    /// Skip keys ascend, so the first block that can hold `row` is a binary
171    /// search over metadata rows; the blocks between are never decoded.
172    pub fn seek(&mut self, row: u32) {
173        self.stream.seek(row);
174    }
175
176    /// Returns the bound of the block the cursor currently sits in.
177    ///
178    /// Read from the stored impact pairs and evaluated against this query's
179    /// live statistics; see [`crate::fts::sealed`].
180    #[must_use]
181    pub fn current_block_max(&self) -> f64 {
182        self.stream.block_bound(&self.scorer)
183    }
184
185    /// Returns true when the cursor has passed the last entry.
186    #[must_use]
187    pub fn exhausted(&self) -> bool {
188        self.stream.exhausted()
189    }
190}
191
192/// A bounded top-k heap with the pinned tie-break.
193///
194/// Ordering is descending score then ascending document id, matching
195/// [`crate::fts::search`] exactly. The threshold this exposes is what every
196/// pruning decision is made against.
197#[derive(Clone, Debug, Default)]
198pub struct TopK {
199    capacity: usize,
200    entries: Vec<(GlobalDocId, f64)>,
201}
202
203impl TopK {
204    /// Creates a heap holding at most `capacity` entries.
205    #[must_use]
206    pub fn new(capacity: usize) -> Self {
207        Self {
208            capacity,
209            entries: Vec::with_capacity(capacity.saturating_add(1)),
210        }
211    }
212
213    /// The pinned ranking order: descending score, then ascending id.
214    ///
215    /// One definition, used by both the insertion and the tests that check
216    /// it against a full sort. Two spellings of a comparator is how a
217    /// tie-break drifts.
218    fn rank(left: &(GlobalDocId, f64), right: &(GlobalDocId, f64)) -> std::cmp::Ordering {
219        right
220            .1
221            .partial_cmp(&left.1)
222            .unwrap_or(std::cmp::Ordering::Equal)
223            .then(left.0.cmp(&right.0))
224    }
225
226    /// Offers one scored document.
227    ///
228    /// Insertion into a sorted array, not a sort on every offer. For `k` of
229    /// order one hundred that is strictly less work, and it reproduces the
230    /// pinned tie-break exactly: the comparator is the same one, and the
231    /// scan stops at the first entry the newcomer does not precede, which is
232    /// where a stable sort would have placed it.
233    pub fn offer(&mut self, doc: GlobalDocId, score: f64) {
234        self.offer_observed(doc, score, |_| {});
235    }
236
237    fn offer_observed(&mut self, doc: GlobalDocId, score: f64, observe: impl FnOnce(usize)) {
238        if self.capacity == 0 {
239            return;
240        }
241        let entry = (doc, score);
242        // A full heap whose worst entry already outranks the newcomer would
243        // have sorted it into the truncated tail. Dropping it here is the
244        // same decision, taken without touching the array.
245        if self.entries.len() >= self.capacity
246            && self
247                .entries
248                .last()
249                .is_some_and(|last| Self::rank(&entry, last) != std::cmp::Ordering::Less)
250        {
251            return;
252        }
253        let mut slot = self.entries.len();
254        while slot > 0 {
255            let Some(previous) = self.entries.get(slot.saturating_sub(1)) else {
256                break;
257            };
258            if Self::rank(&entry, previous) == std::cmp::Ordering::Less {
259                slot = slot.saturating_sub(1);
260            } else {
261                break;
262            }
263        }
264        self.entries.insert(slot, entry);
265        observe(self.entries.len());
266        self.entries.truncate(self.capacity);
267    }
268
269    /// Returns the current k-th best score, or zero while not yet full.
270    ///
271    /// Returning zero rather than negative infinity while the heap is
272    /// filling is deliberate: BM25 scores are non-negative, so zero already
273    /// admits everything, and it keeps the threshold monotone.
274    #[must_use]
275    pub fn threshold(&self) -> f64 {
276        if self.entries.len() < self.capacity {
277            return 0.0;
278        }
279        self.entries.last().map_or(0.0, |(_, score)| *score)
280    }
281
282    /// Returns true when the heap holds `capacity` entries.
283    #[must_use]
284    pub fn is_full(&self) -> bool {
285        self.entries.len() >= self.capacity
286    }
287
288    /// Consumes the heap into ranked results.
289    #[must_use]
290    pub fn into_hits(self) -> Vec<ScoredDoc> {
291        self.entries
292            .into_iter()
293            .map(|(doc, score)| ScoredDoc { doc, score })
294            .collect()
295    }
296}
297
298/// Runs a pruned search and returns the top `k`.
299///
300/// The result is required to equal [`crate::fts::search::search`] exactly.
301///
302/// # Errors
303///
304/// Returns [`IndexError::Stats`] when the index holds no documents.
305pub fn search_pruned(
306    index: &LexicalIndex,
307    query: &TermQuery,
308    k: usize,
309    params: Bm25Params,
310    strategy: Strategy,
311) -> Result<SearchResult, IndexError> {
312    search_pruned_inner(index, query, k, params, strategy, None, None, || {
313        Ok::<(), std::convert::Infallible>(())
314    })
315    .map_err(ControlledSearchError::into_index_error)
316}
317
318pub(crate) fn search_pruned_filtered_controlled<E>(
319    index: &LexicalIndex,
320    query: &TermQuery,
321    k: usize,
322    params: Bm25Params,
323    strategy: Strategy,
324    allow_lists: &[&DocBitmap],
325    check: impl FnMut() -> Result<(), E>,
326) -> Result<SearchResult, ControlledSearchError<E>> {
327    search_pruned_inner(
328        index,
329        query,
330        k,
331        params,
332        strategy,
333        Some(allow_lists),
334        None,
335        check,
336    )
337}
338
339#[cfg(test)]
340pub(crate) fn search_pruned_prepared_filtered(
341    index: &LexicalIndex,
342    prepared: &crate::fts::search::PreparedTermQuery,
343    k: usize,
344    strategy: Strategy,
345    allow_lists: &[&DocBitmap],
346) -> Result<SearchResult, IndexError> {
347    search_pruned_inner(
348        index,
349        prepared.query(),
350        k,
351        prepared.params(),
352        strategy,
353        Some(allow_lists),
354        Some(prepared),
355        || Ok::<(), std::convert::Infallible>(()),
356    )
357    .map_err(ControlledSearchError::into_index_error)
358}
359
360pub(crate) fn search_pruned_prepared_controlled<E>(
361    index: &LexicalIndex,
362    prepared: &crate::fts::search::PreparedTermQuery,
363    k: usize,
364    allow_lists: &[&DocBitmap],
365    check: impl FnMut() -> Result<(), E>,
366) -> Result<SearchResult, ControlledSearchError<E>> {
367    search_pruned_inner(
368        index,
369        prepared.query(),
370        k,
371        prepared.params(),
372        select_strategy(prepared.query().terms.len(), k),
373        Some(allow_lists),
374        Some(prepared),
375        check,
376    )
377}
378
379#[allow(clippy::too_many_arguments)]
380fn search_pruned_inner<E>(
381    index: &LexicalIndex,
382    query: &TermQuery,
383    k: usize,
384    params: Bm25Params,
385    strategy: Strategy,
386    allow_lists: Option<&[&DocBitmap]>,
387    prepared: Option<&crate::fts::search::PreparedTermQuery>,
388    mut check: impl FnMut() -> Result<(), E>,
389) -> Result<SearchResult, ControlledSearchError<E>> {
390    let mut work =
391        crate::fts::control::WorkCheck::new(|| check().map_err(ControlledSearchError::Control));
392    work.check_now()?;
393    if strategy == Strategy::Exhaustive {
394        return allow_lists
395            .map_or_else(
396                || crate::fts::search::search(index, query, k, params),
397                |allow_lists| {
398                    crate::fts::search::search_allow_list_driven(
399                        index,
400                        query,
401                        k,
402                        params,
403                        allow_lists,
404                    )
405                },
406            )
407            .map_err(ControlledSearchError::Index);
408    }
409    let stats = match prepared {
410        Some(prepared) => prepared.stats(),
411        None => index.corpus_stats()?,
412    };
413    let mut counters = SearchCounters::default();
414    let mut heap = TopK::new(k.max(1));
415
416    // Store-wide document frequency, computed once per term.
417    let frequencies = match prepared {
418        Some(prepared) => std::borrow::Cow::Borrowed(prepared.frequencies()),
419        None => {
420            let fields = query.fields.fields_controlled(&mut work)?;
421            let mut frequencies = Vec::with_capacity(query.terms.len());
422            for term in &query.terms {
423                work.step()?;
424                frequencies.push(
425                    index
426                        .prepared_document_frequency_controlled(term, &fields, || work.check_now())
427                        .map_err(|error| match error {
428                            ControlledSearchError::Index(error) => {
429                                ControlledSearchError::Index(error)
430                            }
431                            ControlledSearchError::Control(error) => error,
432                        })?,
433                );
434            }
435            std::borrow::Cow::Owned(frequencies)
436        }
437    };
438
439    for (ordinal, segment) in index.segments().iter().enumerate() {
440        work.check_now()?;
441        let segment_index = u32::try_from(ordinal).unwrap_or(u32::MAX);
442        let allow_list = allow_lists.and_then(|lists| lists.get(ordinal)).copied();
443        // Borrowed outright in the flat single-field case; see
444        // `crate::fts::search::weighted_lengths`. This used to be an
445        // O(row_count) rebuild on every query.
446        let lengths =
447            crate::fts::search::weighted_lengths_controlled(segment, &query.fields, &mut work)?;
448
449        let mut cursors: Vec<TermCursor<'_>> = Vec::with_capacity(query.terms.len());
450        for (slot, term) in query.terms.iter().enumerate() {
451            work.step()?;
452            let df = Df(frequencies.get(slot).copied().unwrap_or(0));
453            if df.0 == 0 {
454                continue;
455            }
456            let Some(stream) =
457                TermStream::open_controlled(segment, term, &query.fields, &mut work)?
458            else {
459                continue;
460            };
461            if stream.exhausted() {
462                continue;
463            }
464            let scorer = prepared
465                .and_then(|prepared| prepared.scorer(slot))
466                .unwrap_or_else(|| TermScorer::new(df, &stats, params));
467            // Read, never rebuilt. The term's bound comes from the stored
468            // per-field overall impact pairs; the
469            // previous form scored every posting in order to bound it, and
470            // that is why pruning was slower than the scan it exists to
471            // avoid.
472            let upper_bound = stream.upper_bound_controlled(&scorer, &mut work)?;
473            cursors.push(TermCursor {
474                stream,
475                df,
476                scorer,
477                upper_bound,
478            });
479        }
480        if cursors.is_empty() {
481            continue;
482        }
483
484        // Short lists skip the pruning machinery entirely. The count comes
485        // from block metadata, so asking costs no decode.
486        let total: usize = cursors
487            .iter()
488            .map(|cursor| cursor.stream.posting_count())
489            .sum();
490        if total <= SHORT_LIST_POSTINGS {
491            maxscore::score_all_controlled(
492                &mut cursors,
493                &lengths,
494                segment_index,
495                &mut heap,
496                &mut counters,
497                allow_list,
498                || work.check_now(),
499            )?;
500        } else {
501            match strategy {
502                Strategy::BlockMaxWand => wand::run_controlled(
503                    &mut cursors,
504                    &lengths,
505                    segment_index,
506                    &mut heap,
507                    &mut counters,
508                    allow_list,
509                    || work.check_now(),
510                )?,
511                Strategy::BlockMaxMaxscore => maxscore::run_controlled(
512                    &mut cursors,
513                    &lengths,
514                    segment_index,
515                    &mut heap,
516                    &mut counters,
517                    allow_list,
518                    || work.check_now(),
519                )?,
520                Strategy::Exhaustive => {}
521            }
522        }
523        for cursor in &cursors {
524            counters.blocks_decoded = counters
525                .blocks_decoded
526                .saturating_add(cursor.stream.blocks_decoded());
527            counters.blocks_skipped = counters
528                .blocks_skipped
529                .saturating_add(cursor.stream.blocks_skipped());
530        }
531    }
532
533    work.check_now()?;
534    Ok(SearchResult {
535        hits: heap.into_hits(),
536        counters,
537    })
538}
539
540#[cfg(test)]
541#[allow(
542    clippy::expect_used,
543    clippy::indexing_slicing,
544    clippy::panic,
545    clippy::unwrap_used
546)]
547mod tests {
548    use super::*;
549    use crate::fts::index::DEFAULT_FIELD;
550    use crate::fts::index::{Document, SegmentIndex};
551    use crate::fts::tokenizer::{Analyzer, Profile};
552
553    fn analyzer() -> Analyzer {
554        Analyzer::new(Profile::Code.config()).expect("valid config")
555    }
556
557    fn index_of(texts: &[String]) -> LexicalIndex {
558        let analyzer = analyzer();
559        let mut segment = SegmentIndex::new();
560        for text in texts {
561            segment
562                .push_document(&analyzer, &Document::with_text(text))
563                .expect("indexable");
564        }
565        let mut index = LexicalIndex::new();
566        index.push_segment(segment).expect("seals");
567        index
568    }
569
570    #[test]
571    fn astra_18_pruned_lexical_cancels_inside_long_traversal() {
572        let texts = vec!["alpha beta gamma delta".to_owned(); 4_096];
573        let index = index_of(&texts);
574        let query = TermQuery::flat(
575            [
576                b"alpha".to_vec(),
577                b"beta".to_vec(),
578                b"gamma".to_vec(),
579                b"delta".to_vec(),
580            ]
581            .into(),
582            &[DEFAULT_FIELD],
583        );
584        let stats = index.corpus_stats().expect("stats");
585        let segment = index.segments().first().expect("one segment");
586        let lengths = crate::fts::search::weighted_lengths(segment, &query.fields);
587        for strategy in [Strategy::BlockMaxWand, Strategy::BlockMaxMaxscore] {
588            let mut cursors: Vec<_> = query
589                .terms
590                .iter()
591                .map(|term| {
592                    let df = Df(index.document_frequency(term, &[DEFAULT_FIELD]));
593                    let scorer = TermScorer::new(df, &stats, Bm25Params::default());
594                    let stream =
595                        TermStream::open(segment, term, &query.fields).expect("term stream");
596                    let upper_bound = stream.upper_bound(&scorer);
597                    TermCursor {
598                        stream,
599                        df,
600                        scorer,
601                        upper_bound,
602                    }
603                })
604                .collect();
605            let mut checks = 0;
606            let mut check = || {
607                checks += 1;
608                if checks >= 3 {
609                    Err("cancelled")
610                } else {
611                    Ok(())
612                }
613            };
614            let mut counters = SearchCounters::default();
615            let mut heap = TopK::new(4_096);
616            let result = match strategy {
617                Strategy::BlockMaxWand => wand::run_controlled(
618                    &mut cursors,
619                    &lengths,
620                    0,
621                    &mut heap,
622                    &mut counters,
623                    None,
624                    &mut check,
625                ),
626                Strategy::BlockMaxMaxscore => maxscore::run_controlled(
627                    &mut cursors,
628                    &lengths,
629                    0,
630                    &mut heap,
631                    &mut counters,
632                    None,
633                    &mut check,
634                ),
635                Strategy::Exhaustive => unreachable!(),
636            };
637            println!(
638                "{strategy:?}: checks={checks}, postings={}, rows={}",
639                counters.postings_decoded, counters.docs_evaluated
640            );
641            assert_eq!(result, Err("cancelled"), "must check during traversal");
642            assert!(
643                counters.postings_decoded <= 128,
644                "cancellation must bound actual work, not only stop publication"
645            );
646        }
647    }
648
649    #[test]
650    fn strategy_selection_follows_the_measured_rule() {
651        assert_eq!(select_strategy(1, 10), Strategy::BlockMaxMaxscore);
652        assert_eq!(select_strategy(3, 10), Strategy::BlockMaxMaxscore);
653        assert_eq!(select_strategy(4, 10), Strategy::BlockMaxWand);
654        assert_eq!(select_strategy(6, 10), Strategy::BlockMaxWand);
655        // `k` no longer discriminates, at either end of the range.
656        assert_eq!(select_strategy(3, 1_000), Strategy::BlockMaxMaxscore);
657        assert_eq!(select_strategy(4, 1_000), Strategy::BlockMaxWand);
658        assert_eq!(select_strategy(0, 0), Strategy::BlockMaxMaxscore);
659    }
660
661    #[test]
662    fn short_list_query_takes_no_pruning_decision() {
663        // The fast path reads its one block -- a posting cannot be scored
664        // without being decoded -- but it makes no skipping decision and it
665        // scores every match. `blocks_decoded` counts real decodes now that
666        // the query path reads the sealed format, so "zero" would be a claim
667        // about a path that no longer exists.
668        let texts: Vec<String> = (0..8).map(|index| format!("alpha doc{index}")).collect();
669        let index = index_of(&texts);
670        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
671        for strategy in [Strategy::BlockMaxWand, Strategy::BlockMaxMaxscore] {
672            let result =
673                search_pruned(&index, &query, 10, Bm25Params::default(), strategy).expect("scores");
674            assert_eq!(
675                result.counters.blocks_decoded, 1,
676                "eight postings live in exactly one block"
677            );
678            assert_eq!(
679                result.counters.blocks_skipped, 0,
680                "the short-list fast path must take no skipping decision"
681            );
682            assert_eq!(result.counters.docs_evaluated, 8);
683            assert_eq!(result.hits.len(), 8);
684        }
685    }
686
687    #[test]
688    fn the_top_k_heap_uses_the_pinned_tie_break() {
689        let mut heap = TopK::new(3);
690        for row in [2_u32, 0, 1] {
691            heap.offer(GlobalDocId { segment: 0, row }, 5.0);
692        }
693        let rows: Vec<u32> = heap.into_hits().iter().map(|hit| hit.doc.row).collect();
694        assert_eq!(rows, vec![0, 1, 2], "ties must break on ascending id");
695    }
696
697    #[test]
698    fn the_threshold_is_zero_until_the_heap_is_full() {
699        let mut heap = TopK::new(2);
700        assert_eq!(heap.threshold(), 0.0);
701        assert!(!heap.is_full());
702        heap.offer(GlobalDocId { segment: 0, row: 0 }, 9.0);
703        assert_eq!(heap.threshold(), 0.0, "a partial heap admits everything");
704        heap.offer(GlobalDocId { segment: 0, row: 1 }, 4.0);
705        assert!(heap.is_full());
706        assert!((heap.threshold() - 4.0).abs() < 1e-12);
707    }
708
709    #[test]
710    fn top_k_insertion_matches_a_full_sort_exactly() {
711        // The sort-per-offer heap P1.4 replaced, as an oracle. Ties matter
712        // most here, so scores are drawn from a deliberately small set: a
713        // tie-break that drifts by one position is a changed result and the
714        // equivalence property is right to reject it.
715        use rand::Rng as _;
716        let mut rng = crate::test_support::seeded_rng(
717            "fts::prune::tests::top_k_insertion_matches_a_full_sort_exactly",
718        );
719        for capacity in [0_usize, 1, 2, 10, 100] {
720            for _ in 0..64 {
721                let mut heap = TopK::new(capacity);
722                let mut oracle: Vec<(GlobalDocId, f64)> = Vec::new();
723                for _ in 0..250 {
724                    let doc = GlobalDocId {
725                        segment: rng.random_range(0..3_u32),
726                        row: rng.random_range(0..40_u32),
727                    };
728                    let score = f64::from(rng.random_range(0..5_u32));
729                    heap.offer(doc, score);
730                    oracle.push((doc, score));
731                    oracle.sort_by(TopK::rank);
732                    oracle.truncate(capacity);
733                    assert_eq!(
734                        heap.clone().into_hits(),
735                        oracle
736                            .iter()
737                            .map(|(doc, score)| ScoredDoc {
738                                doc: *doc,
739                                score: *score
740                            })
741                            .collect::<Vec<_>>(),
742                        "insertion diverged from a full sort at capacity {capacity}"
743                    );
744                    assert_eq!(
745                        heap.threshold().to_bits(),
746                        if oracle.len() < capacity {
747                            0.0_f64
748                        } else {
749                            oracle.last().map_or(0.0, |(_, score)| *score)
750                        }
751                        .to_bits()
752                    );
753                }
754            }
755        }
756    }
757
758    /// Builds a single-term index whose rows advance by `stride`.
759    fn strided_index(count: u32, stride: u32) -> LexicalIndex {
760        let analyzer = analyzer();
761        let mut segment = SegmentIndex::new();
762        let mut row = 0_u32;
763        for ordinal in 0..count {
764            // Rows between the term's occurrences carry a different term, so
765            // the posting list has real gaps rather than a dense run.
766            while row < ordinal.saturating_mul(stride) {
767                segment
768                    .push_document(&analyzer, &Document::with_text("filler"))
769                    .expect("indexable");
770                row = row.saturating_add(1);
771            }
772            segment
773                .push_document(&analyzer, &Document::with_text("alpha"))
774                .expect("indexable");
775            row = row.saturating_add(1);
776        }
777        let mut index = LexicalIndex::new();
778        index.push_segment(segment).expect("seals");
779        index
780    }
781
782    /// The rows a term occupies, as a linear oracle for `seek`.
783    fn rows_of(index: &LexicalIndex) -> Vec<u32> {
784        let segment = index.segments().first().expect("one segment");
785        let mut stream = crate::fts::sealed::TermStream::open(
786            segment,
787            b"alpha",
788            &crate::fts::search::FieldWeights::flat(&[DEFAULT_FIELD]),
789        )
790        .expect("stream");
791        let mut rows = Vec::new();
792        while let Some(row) = stream.current_row() {
793            rows.push(row);
794            stream.advance();
795        }
796        rows
797    }
798
799    /// Opens a stream over the one term of a `strided_index`.
800    fn stream_of(index: &LexicalIndex) -> crate::fts::sealed::TermStream<'_> {
801        crate::fts::sealed::TermStream::open(
802            index.segments().first().expect("one segment"),
803            b"alpha",
804            &crate::fts::search::FieldWeights::flat(&[DEFAULT_FIELD]),
805        )
806        .expect("stream")
807    }
808
809    #[test]
810    fn a_sealed_seek_lands_exactly_where_a_linear_walk_would() {
811        // Seeks are decisions, and decisions are pinned by the counter
812        // contracts, so a binary search over skip keys must land on the same
813        // row a walk would -- from every starting position, to every target.
814        for count in [1_u32, 7, 64, 65, 200] {
815            for stride in [1_u32, 3] {
816                let index = strided_index(count, stride);
817                let rows = rows_of(&index);
818                assert_eq!(rows.len(), usize::try_from(count).expect("small"));
819                let last = count.saturating_mul(stride).saturating_add(2);
820                for start in [0_usize, 1, 5, 63, 64, 199] {
821                    if start >= rows.len() {
822                        continue;
823                    }
824                    for target in 0..=last {
825                        let mut cursor = stream_of(&index);
826                        for _ in 0..start {
827                            cursor.advance();
828                        }
829                        cursor.seek(target);
830                        // The oracle: the first row at or after `target`,
831                        // never earlier than where the cursor already sat.
832                        let expected = rows.iter().skip(start).copied().find(|row| *row >= target);
833                        assert_eq!(
834                            cursor.current_row(),
835                            expected,
836                            "seek diverged (count={count}, stride={stride}, \
837                             start={start}, target={target})"
838                        );
839                    }
840                }
841            }
842        }
843    }
844
845    #[test]
846    fn a_sealed_seek_skips_whole_blocks_without_decoding_them() {
847        let index = strided_index(200, 1);
848        let mut stream = stream_of(&index);
849        let decoded = stream.blocks_decoded();
850        stream.seek(150);
851        assert_eq!(stream.current_row(), Some(150));
852        assert!(
853            stream.blocks_skipped() >= 1,
854            "seeking past a whole block must report it"
855        );
856        assert_eq!(
857            stream.blocks_decoded(),
858            decoded + 1,
859            "a seek decodes exactly the block it lands in"
860        );
861    }
862
863    #[test]
864    fn seeking_beyond_the_last_row_exhausts_the_cursor() {
865        let index = strided_index(10, 1);
866        let stats = index.corpus_stats().expect("stats");
867        let scorer = TermScorer::new(Df(10), &stats, Bm25Params::default());
868        let mut cursor = TermCursor {
869            stream: stream_of(&index),
870            df: Df(10),
871            scorer,
872            upper_bound: 1.0,
873        };
874        cursor.seek(9_999);
875        assert!(cursor.exhausted());
876        assert_eq!(cursor.current(), None);
877        assert_eq!(cursor.current_tf(), None);
878        // Every run is spent, so no block contributes and the bound
879        // collapses to zero rather than to a stale number.
880        assert!(cursor.current_block_max().abs() < 1e-12);
881    }
882
883    #[test]
884    fn the_sealed_block_bound_matches_the_in_memory_oracle() {
885        // `build_block_bounds` is the reference implementation of the same
886        // impact-pair formula. Reading a bound out of bytes and building one
887        // in memory must give the same number, or the equivalence property
888        // passes against a fixture and fails against a real segment.
889        let index = strided_index(200, 1);
890        let stats = index.corpus_stats().expect("stats");
891        let scorer = TermScorer::new(Df(200), &stats, Bm25Params::default());
892        let rows = rows_of(&index);
893        let entries: Vec<(u32, u32)> = rows.iter().map(|row| (*row, 1_u32)).collect();
894        let segment = index.segments().first().expect("one segment");
895        let lengths = segment.field_lengths(DEFAULT_FIELD).expect("lengths");
896        let oracle = build_block_bounds(&entries, lengths, PRUNE_BLOCK_SIZE, &scorer);
897
898        let mut stream = stream_of(&index);
899        for block in &oracle {
900            assert!(
901                (stream.block_bound(&scorer) - block.max_score).abs() < 1e-15,
902                "sealed and in-memory bounds diverged at block {:?}",
903                block.start
904            );
905            for _ in block.start..block.end {
906                stream.advance();
907            }
908        }
909        let overall = term_upper_bound(&oracle);
910        let fresh = stream_of(&index);
911        assert!((fresh.upper_bound(&scorer) - overall).abs() < 1e-15);
912    }
913
914    #[test]
915    fn maxscore_abandons_a_candidate_before_touching_the_cheap_cursor() {
916        // One rare, high-impact term and one common, low-impact term. The
917        // single document holding both fills the heap at k=1; after that,
918        // every rare-only candidate is provably unable to reach the
919        // threshold BEFORE the common cursor is asked to move, so the
920        // common list's later blocks must never be decoded. A probe order
921        // that seeks the cheap cursor first walks every common block and
922        // fails this bound.
923        let mut texts: Vec<String> = Vec::new();
924        // Row 0: the top document. The rare term dominates its score.
925        texts.push(format!("{} common", "rare ".repeat(30).trim()));
926        // Ten long rare-only documents interleaved every twenty common
927        // documents, so a seek-per-candidate traversal would land in every
928        // common block in turn.
929        for group in 0..10 {
930            for slot in 0..20 {
931                texts.push(format!("common c{group}x{slot}"));
932            }
933            let padding: String = (0..64).map(|word| format!("p{group}w{word} ")).collect();
934            texts.push(format!("rare {padding}"));
935        }
936        let index = index_of(&texts);
937        let query = TermQuery::flat(vec![b"rare".to_vec(), b"common".to_vec()], &[DEFAULT_FIELD]);
938        let pruned = search_pruned(
939            &index,
940            &query,
941            1,
942            Bm25Params::default(),
943            Strategy::BlockMaxMaxscore,
944        )
945        .expect("scores");
946        let oracle =
947            crate::fts::search::search(&index, &query, 1, Bm25Params::default()).expect("scores");
948        assert_eq!(pruned.hits, oracle.hits, "pruning must not change results");
949        assert!(
950            pruned.counters.blocks_decoded <= 2,
951            "abandonment must fire before the cheap cursor seeks; decoded {} blocks \
952             (one rare block plus the common list's first is the whole budget)",
953            pruned.counters.blocks_decoded
954        );
955    }
956
957    #[test]
958    fn the_metadata_block_bound_dominates_every_row_it_bounds() {
959        // `bound_for` is what MAXSCORE abandons candidates against, so it
960        // must dominate the exact contribution of any row the stream can
961        // still reach — from every cursor position, to every target. A
962        // bound below the truth is a dropped result, not a slow one.
963        let index = strided_index(200, 3);
964        let stats = index.corpus_stats().expect("stats");
965        let scorer = TermScorer::new(Df(200), &stats, Bm25Params::default());
966        let segment = index.segments().first().expect("one segment");
967        let lengths = segment.field_lengths(DEFAULT_FIELD).expect("lengths");
968        let last = rows_of(&index).last().copied().unwrap_or(0);
969        for start in [0_usize, 1, 63, 64, 150] {
970            for target in 0..=last.saturating_add(2) {
971                let mut stream = stream_of(&index);
972                for _ in 0..start {
973                    stream.advance();
974                }
975                let bound = stream.bound_for(target, &scorer);
976                let mut probe = stream.clone();
977                probe.seek(target);
978                if probe.current_row() == Some(target) {
979                    let tf = probe.current_tf().unwrap_or(0);
980                    let length = lengths
981                        .get(usize::try_from(target).expect("small"))
982                        .copied();
983                    let exact = scorer.score(
984                        crate::fts::bm25::Tf(tf),
985                        crate::fts::bm25::DocLen(length.unwrap_or(1)),
986                    );
987                    assert!(
988                        bound + 1e-12 >= exact,
989                        "bound {bound} fell below the exact score {exact} \
990                         (start={start}, target={target})"
991                    );
992                }
993            }
994        }
995    }
996
997    #[test]
998    fn the_exhaustive_strategy_delegates_to_the_oracle() {
999        let texts: Vec<String> = (0..5).map(|index| format!("alpha beta{index}")).collect();
1000        let index = index_of(&texts);
1001        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
1002        let pruned = search_pruned(
1003            &index,
1004            &query,
1005            3,
1006            Bm25Params::default(),
1007            Strategy::Exhaustive,
1008        )
1009        .expect("scores");
1010        let direct =
1011            crate::fts::search::search(&index, &query, 3, Bm25Params::default()).expect("scores");
1012        assert_eq!(pruned.hits, direct.hits);
1013    }
1014
1015    #[test]
1016    fn an_empty_index_is_a_typed_error_under_every_strategy() {
1017        let index = LexicalIndex::new();
1018        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
1019        for strategy in [
1020            Strategy::BlockMaxWand,
1021            Strategy::BlockMaxMaxscore,
1022            Strategy::Exhaustive,
1023        ] {
1024            assert!(search_pruned(&index, &query, 5, Bm25Params::default(), strategy).is_err());
1025        }
1026    }
1027}