Skip to main content

plugmem_core/index/
bm25.rs

1//! The lexical index: classic BM25 over delta-encoded postings
2//!
3//! Scoring is the standard formula with the Robertson idf:
4//!
5//! ```text
6//! idf(t)      = ln(1 + (N - df + 0.5) / (df + 0.5))
7//! tf_norm(d)  = tf · (k1 + 1) / (tf + k1 · (1 - b + b · len(d) / avg_len))
8//! score(d, q) = Σ_t idf(t) · tf_norm(d, t)
9//! ```
10//!
11//! A query decodes every query term's postings — O(Σ df), and there is no
12//! WAND-style pruning to avoid it. That is affordable because a decode is a
13//! few nanoseconds of varint over a contiguous chunk chain; what is not
14//! affordable is a *random* lookup per posting, and the scan is built to have
15//! none:
16//!
17//! - document lengths come from a flat array indexed by fact id, not from the
18//!   stored arena (see `Bm25Index::doc_len_dense`);
19//! - partial scores accumulate by merging sorted runs, because the postings
20//!   are already sorted by fact id, so no map is probed;
21//! - the caller's `live` predicate — a fact-record lookup on the engine side —
22//!   is asked only about documents in contention for the top `k`, since a
23//!   filter can remove entries from a ranking but never reorder it.
24//!
25//! The `decoded`, `scored` and `admitted` counters gate exactly this split in
26//! CI, and `bm25_probe_work_is_bounded` pins the last one at `k`.
27//!
28//! Deletions never touch the postings: tombstoned facts are filtered per
29//! candidate by the caller's `live` predicate and fall out physically
30//! when `maintain` rebuilds the index.
31
32use alloc::vec::Vec;
33
34#[cfg(feature = "counters")]
35use core::cell::Cell;
36
37use plugmem_arena::{Arena, ArenaCfg, ShardMode, Slot, key};
38
39use crate::error::Error;
40use crate::id::FactId;
41use crate::index::postings::PostingStore;
42
43/// Byte layout of [`DocLenSlot`]. Every offset is the previous field's offset
44/// plus its width, so a field cannot be moved by editing one number.
45mod doclen_at {
46    use core::mem::size_of;
47
48    pub(super) const FACT: usize = 0;
49    pub(super) const KEY_LEN: usize = FACT + size_of::<u32>();
50    pub(super) const LEN: usize = KEY_LEN;
51    pub(super) const DISTINCT: usize = LEN + size_of::<u16>();
52    pub(super) const SIG: usize = DISTINCT + size_of::<u16>();
53    pub(super) const SIZE: usize = SIG + size_of::<u64>();
54}
55
56/// Width of a [`DocLenSlot::sig`] signature in bits.
57const SIG_BITS: u32 = 64;
58/// Fibonacci hashing multiplier — the same constant the arena shards with, so
59/// term ids scatter over the signature's bits without a second hash family.
60const SIG_MULT: u64 = 0x9E37_79B9_7F4A_7C15;
61
62/// The signature bit a term claims. Distinct terms may collide; that is what
63/// makes [`DocLenSlot::sig`] an over-approximation and never an
64/// under-approximation of a document's term set.
65pub(crate) fn sig_bit(term: u32) -> u64 {
66    // The top `log2(SIG_BITS)` bits of the multiplied word, so the index is in
67    // range by construction.
68    let index = u64::from(term).wrapping_mul(SIG_MULT) >> (64 - SIG_BITS.trailing_zeros());
69    1u64 << index
70}
71
72/// Per-document record: `[fact 4 | len u16 | distinct u16 | sig u64]`,
73/// Uniform arena.
74///
75/// Beyond the length BM25 scores with, the slot carries a summary of the
76/// document's *term set*: how many distinct terms it has, and one bit per
77/// term hashed into a 64-bit word. That summary is what lets the write path
78/// bound the term-set overlap of two facts without re-reading and
79/// re-tokenizing their texts (see `Memory::find_similar`). It is written when
80/// the document is indexed, where the term set is already in hand, so it
81/// costs nothing to produce.
82#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
84pub struct DocLenSlot {
85    /// The document (fact) id — the key.
86    pub fact: FactId,
87    /// Token count of the document, saturated at `u16::MAX`.
88    pub len: u16,
89    /// Number of distinct terms, saturated at `u16::MAX`. Zero means
90    /// "unknown": a document indexed before the signature existed, read
91    /// through the legacy migration.
92    pub distinct: u16,
93    /// Union of `sig_bit` over the document's distinct terms. A term absent
94    /// from this word is definitely absent from the document; a term present
95    /// may still be absent (bits collide). Zero alongside `distinct == 0`
96    /// means "unknown".
97    pub sig: u64,
98}
99
100impl DocLenSlot {
101    /// Whether the term-set summary is present. A legacy document carries
102    /// none, and callers must fall back to reading its text.
103    pub fn has_signature(&self) -> bool {
104        self.distinct != 0
105    }
106
107    /// An upper bound on how many of `terms` this document also holds.
108    ///
109    /// Exact in the direction that matters: a term whose bit is clear cannot
110    /// be in the document, so the true intersection is never larger than the
111    /// count returned here. Callers use it to rule overlap *out*.
112    pub fn overlap_bound(&self, terms: &[u32]) -> usize {
113        terms
114            .iter()
115            .filter(|&&term| self.sig & sig_bit(term) != 0)
116            .count()
117    }
118}
119
120impl Slot for DocLenSlot {
121    const SIZE: usize = doclen_at::SIZE;
122    const KEY_LEN: usize = doclen_at::KEY_LEN;
123
124    fn write(&self, out: &mut [u8]) {
125        key::write_u32(&mut out[doclen_at::FACT..], self.fact.0);
126        out[doclen_at::LEN..doclen_at::DISTINCT].copy_from_slice(&self.len.to_be_bytes());
127        out[doclen_at::DISTINCT..doclen_at::SIG].copy_from_slice(&self.distinct.to_be_bytes());
128        out[doclen_at::SIG..doclen_at::SIZE].copy_from_slice(&self.sig.to_be_bytes());
129    }
130
131    fn read(bytes: &[u8]) -> Self {
132        Self {
133            fact: FactId(key::read_u32(&bytes[doclen_at::FACT..])),
134            len: u16::from_be_bytes(
135                bytes[doclen_at::LEN..doclen_at::DISTINCT]
136                    .try_into()
137                    .unwrap(),
138            ),
139            distinct: u16::from_be_bytes(
140                bytes[doclen_at::DISTINCT..doclen_at::SIG]
141                    .try_into()
142                    .unwrap(),
143            ),
144            sig: u64::from_be_bytes(bytes[doclen_at::SIG..doclen_at::SIZE].try_into().unwrap()),
145        }
146    }
147}
148
149/// Reusable query scratch: score accumulator and top-k selection buffer. One
150/// per concurrent reader; after warm-up a query allocates nothing (the
151/// zero-alloc recall invariant).
152#[derive(Debug, Default)]
153pub struct Bm25Scratch {
154    /// Partial scores as `(fact, score)` **sorted by fact id**.
155    ///
156    /// A map would be the obvious shape, and was the original one, but the
157    /// postings are already sorted by fact id: accumulating a term is then a
158    /// linear merge of two sorted runs instead of one hash probe per posting.
159    /// The difference is not the hashing — it is that a map big enough to hold
160    /// a frequent term's postings misses cache on essentially every probe,
161    /// while a merge walks three arrays forwards.
162    acc: Vec<(u32, f32)>,
163    /// Merge target, swapped with `acc` after each term past the first.
164    merge: Vec<(u32, f32)>,
165    /// Selection buffer for the top-k extraction.
166    top: Vec<(f32, u32)>,
167}
168
169impl Bm25Scratch {
170    /// Empty scratch buffers.
171    pub fn new() -> Self {
172        Self::default()
173    }
174}
175
176/// The BM25 index: postings with term frequencies plus per-document
177/// lengths and corpus statistics.
178#[derive(Debug)]
179pub struct Bm25Index<'a> {
180    postings: PostingStore<'a, true>,
181    doc_len: Arena<'a, DocLenSlot>,
182    total_docs: u64,
183    total_len: u64,
184    /// Posting entries decoded by queries (feature `counters`) — the
185    /// deterministic cost metric of the lexical source.
186    #[cfg(feature = "counters")]
187    decoded: Cell<u64>,
188    /// Documents whose BM25 contribution was actually evaluated — a document
189    /// length fetched and `tf_norm` computed (feature `counters`).
190    ///
191    /// Decoding a posting entry is a few nanoseconds of varint; *scoring* it
192    /// costs a document-length lookup, which is a random probe into an arena
193    /// that grows with the corpus. The two counters therefore measure
194    /// different things, and this is the one that dominates a large query.
195    #[cfg(feature = "counters")]
196    scored: Cell<u64>,
197    /// Calls to the caller's `live` predicate (feature `counters`).
198    ///
199    /// Every call is a fact-record lookup on the engine side — the second
200    /// random probe per candidate. A candidate that cannot reach the top `k`
201    /// should never cost one.
202    #[cfg(feature = "counters")]
203    admitted: Cell<u64>,
204    /// Some documents arrived without a term-set summary — this index came out
205    /// of a pre-signature image. Derived at load, never persisted: the next
206    /// compaction fills the summaries in and clears it.
207    unsummarized: bool,
208    /// Document length by fact id, [`DOC_LEN_ABSENT`] where the id names no
209    /// indexed document.
210    ///
211    /// Scoring needs the length of every document a query term's postings
212    /// name, which is one lookup per posting entry — the single hottest read
213    /// in the engine, and a random probe into an arena that deepens as the
214    /// corpus grows. Fact ids are dense and monotone, so a flat array answers
215    /// it by indexing, at four bytes per id.
216    ///
217    /// Runtime-only, like the arena's own page directory: rebuilt from
218    /// `doc_len` on load, never written to the snapshot. `doc_len` remains the
219    /// stored form and the only place the term-set summary lives.
220    doc_len_dense: Vec<u32>,
221    /// Exclusive upper bound of the fact ids `Bm25Index::doc_len_dense` may
222    /// be trusted for. `usize::MAX` while it has covered every document it was
223    /// offered; lowered to the first id it declined.
224    dense_limit: usize,
225}
226
227/// `Bm25Index::doc_len_dense` entry for a fact id with no indexed document.
228/// Lengths saturate at `u16::MAX`, so this cannot collide with a real one.
229const DOC_LEN_ABSENT: u32 = u32::MAX;
230
231impl<'a> Bm25Index<'a> {
232    /// Creates an empty index; `shards` per the engine config
233    /// (`shards_postings`), `max_bytes` bounds each underlying pool.
234    pub fn new(shards: usize, max_bytes: usize) -> Result<Self, Error> {
235        Ok(Self {
236            postings: PostingStore::new(shards, max_bytes)?,
237            doc_len: Arena::new(
238                ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(max_bytes),
239            )?,
240            total_docs: 0,
241            total_len: 0,
242            #[cfg(feature = "counters")]
243            decoded: Cell::new(0),
244            #[cfg(feature = "counters")]
245            scored: Cell::new(0),
246            #[cfg(feature = "counters")]
247            admitted: Cell::new(0),
248            unsummarized: false,
249            doc_len_dense: Vec::new(),
250            dense_limit: usize::MAX,
251        })
252    }
253
254    /// Indexes one document given its `(term, tf)` pairs (the caller
255    /// tokenizes and counts; pairs may arrive in any order, terms must be
256    /// unique). Documents must arrive in ascending fact-id order.
257    ///
258    /// # Errors
259    ///
260    /// [`Error::Arena`] when a pool hits its byte ceiling; the index may
261    /// then hold the document partially — the engine treats that as fatal
262    /// for the whole operation (journal replay rebuilds consistently).
263    pub fn index_doc(&mut self, fact: FactId, term_tfs: &[(u32, u8)]) -> Result<(), Error> {
264        let mut len = 0u32;
265        let mut sig = 0u64;
266        for &(term, tf) in term_tfs {
267            self.postings.push(term, fact, tf)?;
268            len += u32::from(tf);
269            sig |= sig_bit(term);
270        }
271        let doc = DocLenSlot {
272            fact,
273            len: u16::try_from(len).unwrap_or(u16::MAX),
274            // The caller's pairs are already unique per term, so their count
275            // is the distinct-term count.
276            distinct: u16::try_from(term_tfs.len()).unwrap_or(u16::MAX),
277            sig,
278        };
279        self.doc_len.insert(&doc)?;
280        self.note_dense(&doc);
281        self.total_docs += 1;
282        self.total_len += u64::from(len);
283        Ok(())
284    }
285
286    /// Records a document's length in the flat index, growing it to reach the
287    /// id. Fact ids are dense and monotone, so the growth is amortized.
288    ///
289    /// Growth stops where the id space stops being dense — see
290    /// [`Bm25Index::dense_capacity`]. Nothing is lost when it does:
291    /// [`Bm25Index::doc_len_of`] reads the stored arena for an id the flat
292    /// index does not cover.
293    fn note_dense(&mut self, doc: &DocLenSlot) {
294        // A `u32` id always fits a `usize`, on wasm32 as on a 64-bit host; the
295        // bound that matters is the capacity below.
296        let at = doc.fact.0 as usize;
297        if at >= self.dense_capacity() {
298            // Declined. The array can never speak for this id, and documents
299            // do not arrive in id order (compaction walks a hashed arena), so
300            // a later id may extend the array right over this one — hence the
301            // watermark rather than a bare skip.
302            self.dense_limit = self.dense_limit.min(at);
303            return;
304        }
305        if at >= self.doc_len_dense.len() {
306            self.doc_len_dense.resize(at + 1, DOC_LEN_ABSENT);
307        }
308        self.doc_len_dense[at] = u32::from(doc.len);
309    }
310
311    /// Highest fact id the flat length index will grow to cover.
312    ///
313    /// The index is worth its four bytes per id only while ids are dense,
314    /// which they are by construction — they are handed out in order, and only
315    /// purged facts leave holes. An id far past the document count means the
316    /// space is sparse or the image is lying, and a flat array would then cost
317    /// more memory than the documents themselves; a snapshot is untrusted
318    /// input, and nothing range-checks the ids inside the stored records.
319    /// `usize` being 32 bits on wasm32 makes the same claim sharper there.
320    ///
321    /// This bounds memory, never correctness: an id past the cap is answered
322    /// from the arena.
323    /// Saturating throughout, which is also what makes `at + 1` safe wherever
324    /// `at < dense_capacity()` holds: the capacity never exceeds `usize::MAX`,
325    /// so an id below it is below `usize::MAX` too.
326    fn dense_capacity(&self) -> usize {
327        const SLACK: usize = 8;
328        let docs = self.total_docs.min(usize::MAX as u64) as usize;
329        docs.saturating_add(1).saturating_mul(SLACK)
330    }
331
332    /// Rebuilds the flat length index from the stored records (the load path).
333    fn rebuild_dense(&mut self) {
334        self.doc_len_dense.clear();
335        self.dense_limit = usize::MAX;
336        let docs: Vec<DocLenSlot> = self.doc_len.iter().collect();
337        for doc in docs {
338            self.note_dense(&doc);
339        }
340    }
341
342    /// The per-document record of `fact`, or `None` when the document is not
343    /// indexed. Carries the term-set summary the write path bounds overlap
344    /// with — see [`DocLenSlot`].
345    pub fn doc(&self, fact: FactId) -> Option<DocLenSlot> {
346        self.doc_len.get(&fact.0.to_be_bytes())
347    }
348
349    /// Whether this index holds documents with no term-set summary, which
350    /// compaction can fill in from the postings. True only after opening a
351    /// pre-signature image.
352    pub(crate) fn needs_resummarize(&self) -> bool {
353        self.unsummarized
354    }
355
356    /// Marks the index as holding unsummarized documents (the load path, after
357    /// a legacy migration).
358    pub(crate) fn mark_unsummarized(&mut self) {
359        self.unsummarized = true;
360    }
361
362    /// Builds a compacted BM25 index by filtering this index's existing
363    /// postings and document lengths through `live`.
364    ///
365    /// This is the ordinary-maintenance path: it preserves the exact term ids
366    /// and term frequencies already indexed, so compaction does not have to
367    /// read and tokenize every live document again. A tokenizer migration must
368    /// use the text reindex path instead.
369    pub(crate) fn compact_live(
370        &self,
371        shards: usize,
372        max_bytes: usize,
373        mut live: impl FnMut(FactId) -> bool,
374    ) -> Result<Bm25Index<'static>, Error> {
375        let mut out = Bm25Index::new(shards, max_bytes)?;
376        // Documents carried over from a pre-signature image, and the id range
377        // they span. Compaction is the cheapest place to fill their term-set
378        // summaries in: it already walks every posting, which is the transpose
379        // of what a signature needs, so no text is read and nothing is
380        // tokenized.
381        let mut legacy = 0usize;
382        let mut max_fact = 0u32;
383        for doc in self.doc_len.iter() {
384            if live(doc.fact) {
385                out.doc_len.insert(&doc)?;
386                out.note_dense(&doc);
387                out.total_docs += 1;
388                out.total_len += u64::from(doc.len);
389                if !doc.has_signature() {
390                    legacy += 1;
391                    max_fact = max_fact.max(doc.fact.0);
392                }
393            }
394        }
395        // `(sig, distinct)` per fact id, dense because the transpose visits
396        // facts in posting order rather than document order. Allocated only
397        // when there is something to fill, and dropped with this call.
398        //
399        // Sized like the flat length index and for the same reasons: `usize`
400        // is 32 bits on wasm32, and the ids come from an untrusted image, so
401        // the array is allocated only for an id space that is actually dense.
402        // Declining costs speed and nothing else — an unsummarized document
403        // keeps reading its text.
404        let highest = max_fact as usize;
405        let mut rebuilt = if legacy > 0 && highest < out.dense_capacity() {
406            alloc::vec![(0u64, 0u16); highest + 1]
407        } else {
408            Vec::new()
409        };
410        for slot in self.postings.slots() {
411            for (fact, tf) in self.postings.entries(slot.key) {
412                if !live(fact) {
413                    continue;
414                }
415                out.postings.push(slot.key, fact, tf)?;
416                if let Some(entry) = rebuilt.get_mut(fact.0 as usize) {
417                    entry.0 |= sig_bit(slot.key);
418                    entry.1 = entry.1.saturating_add(1);
419                }
420            }
421        }
422        if !rebuilt.is_empty() {
423            out.fill_missing_signatures(&rebuilt);
424        }
425        Ok(out)
426    }
427
428    /// Writes the recomputed term-set summaries of [`Self::compact_live`]
429    /// into the documents that arrived without one. Documents that already
430    /// carry a signature keep it: it was written from the exact term set the
431    /// indexer saw, and the transpose can only reproduce it.
432    ///
433    /// The compacted index never inherits [`Self::unsummarized`]. A document
434    /// still unsummarized after this pass has no indexed terms at all, so no
435    /// later pass could summarize it either — carrying the flag forward would
436    /// make every future `maintain` recompact for work that cannot be done.
437    fn fill_missing_signatures(&mut self, rebuilt: &[(u64, u16)]) {
438        let stale: Vec<DocLenSlot> = self
439            .doc_len
440            .iter()
441            .filter(|doc| !doc.has_signature())
442            .collect();
443        for mut doc in stale {
444            let Some(&(sig, distinct)) = rebuilt.get(doc.fact.0 as usize) else {
445                continue;
446            };
447            if distinct == 0 {
448                continue; // a document with no indexed terms has nothing to summarize
449            }
450            doc.sig = sig;
451            doc.distinct = distinct;
452            let Some(payload) = self.doc_len.payload_mut(&doc.fact.0.to_be_bytes()) else {
453                continue;
454            };
455            let mut full = [0u8; DocLenSlot::SIZE];
456            doc.write(&mut full);
457            payload.copy_from_slice(&full[DocLenSlot::KEY_LEN..]);
458        }
459    }
460
461    /// Document frequency of a term.
462    pub fn df(&self, term: u32) -> u32 {
463        self.postings.count(term)
464    }
465
466    /// Number of indexed documents.
467    pub fn docs(&self) -> u64 {
468        self.total_docs
469    }
470
471    /// Robertson idf for a term with document frequency `df` in this
472    /// corpus (monotonically decreasing in `df`, always positive).
473    pub fn idf(&self, df: u32) -> f32 {
474        let n = self.total_docs as f32;
475        let df = df as f32;
476        libm::logf(1.0 + (n - df + 0.5) / (df + 0.5))
477    }
478
479    /// Scores `terms` against the corpus and writes the top `k` live
480    /// documents into `out` (descending score, ties by ascending id).
481    /// `live` filters candidates (tombstones, as-of, tag allow-sets) —
482    /// filtered documents cost their posting decode but never rank.
483    ///
484    /// Duplicate query terms are the caller's choice: each occurrence
485    /// accumulates again (a term repeated in the query weighs more).
486    ///
487    /// Cost is O(Σ df) in *decodes*, which is the honest price of a lexical
488    /// scan, but the expensive part of a candidate is not its decode — it is
489    /// the two random lookups that used to follow it, one for the document
490    /// length and one for the `live` predicate. Neither is paid per candidate
491    /// any more: lengths come from a flat array, and `live` is asked only
492    /// about documents that are actually in contention for the top `k`.
493    pub fn search(
494        &self,
495        (k1, b): (f32, f32),
496        terms: &[u32],
497        k: usize,
498        live: &mut dyn FnMut(FactId) -> bool,
499        scratch: &mut Bm25Scratch,
500        out: &mut Vec<(FactId, f32)>,
501    ) {
502        out.clear();
503        if self.total_docs == 0 || k == 0 {
504            return;
505        }
506        let Bm25Scratch { acc, merge, top } = scratch;
507        acc.clear();
508        let avg_len = self.total_len as f32 / self.total_docs as f32;
509        #[cfg(feature = "counters")]
510        let (mut decoded, mut scored) = (0u64, 0u64);
511
512        for &term in terms {
513            let df = self.postings.count(term);
514            if df == 0 {
515                continue;
516            }
517            let idf = self.idf(df);
518            // A term's postings are already ascending by fact id and `acc`
519            // holds the same order, so accumulating is a merge of two sorted
520            // runs. The first term has nothing to merge against and fills
521            // `acc` directly.
522            let mut ahead = 0usize;
523            merge.clear();
524            for (fact, tf) in self.postings.entries(term) {
525                #[cfg(feature = "counters")]
526                {
527                    decoded += 1;
528                }
529                // Everything in `acc` below this posting keeps its score.
530                while let Some(&entry) = acc.get(ahead)
531                    && entry.0 < fact.0
532                {
533                    merge.push(entry);
534                    ahead += 1;
535                }
536                let carried = match acc.get(ahead) {
537                    Some(&entry) if entry.0 == fact.0 => {
538                        ahead += 1;
539                        Some(entry.1)
540                    }
541                    _ => None,
542                };
543                // A posting naming a document with no length record scores
544                // nothing — and must not create a candidate either.
545                let Some(len) = self.doc_len_of(fact) else {
546                    if let Some(score) = carried {
547                        merge.push((fact.0, score));
548                    }
549                    continue;
550                };
551                #[cfg(feature = "counters")]
552                {
553                    scored += 1;
554                }
555                let tf = f32::from(tf);
556                let norm = tf * (k1 + 1.0) / (tf + k1 * (1.0 - b + b * f32::from(len) / avg_len));
557                // Terms are summed in query order, the order the accumulating
558                // map used, so the float result is bit-for-bit the same.
559                merge.push((fact.0, carried.unwrap_or(0.0) + idf * norm));
560            }
561            merge.extend_from_slice(&acc[ahead.min(acc.len())..]);
562            core::mem::swap(acc, merge);
563        }
564        #[cfg(feature = "counters")]
565        {
566            self.decoded.set(self.decoded.get() + decoded);
567            self.scored.set(self.scored.get() + scored);
568        }
569
570        // Top-k. Ranking is by score alone, so `live` cannot change the order
571        // — only remove entries from it. Asking it about every candidate is
572        // therefore wasted work: rank first, then walk the ranking and ask
573        // only until `k` survivors are found. The result is the same set in
574        // the same order an exhaustive filter produces.
575        top.clear();
576        top.extend(acc.iter().map(|&(id, score)| (score, id)));
577        let order = |a: &(f32, u32), b: &(f32, u32)| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1));
578        #[cfg(feature = "counters")]
579        let mut admitted = 0u64;
580        let mut consume = |band: &[(f32, u32)], out: &mut Vec<(FactId, f32)>| {
581            for &(score, id) in band {
582                if out.len() == k {
583                    return;
584                }
585                #[cfg(feature = "counters")]
586                {
587                    admitted += 1;
588                }
589                if live(FactId(id)) {
590                    out.push((FactId(id), score));
591                }
592            }
593        };
594
595        // The usual case: partition the `k` highest scores to the front,
596        // order them, and take the survivors. Linear, and it asks `live`
597        // about `k` documents rather than every candidate.
598        let band = k.min(top.len());
599        if band > 0 {
600            if band < top.len() {
601                top.select_nth_unstable_by(band - 1, order);
602            }
603            top[..band].sort_unstable_by(order);
604            consume(&top[..band], out);
605        }
606        // The band was thinned by tombstones or a filter. Order what is left
607        // in one pass and continue down it — the same total cost the
608        // exhaustive filter used to pay on every query, now only on a query
609        // that needs it.
610        if out.len() < k && band < top.len() {
611            top[band..].sort_unstable_by(order);
612            let (_, rest) = top.split_at(band);
613            consume(rest, out);
614        }
615        #[cfg(feature = "counters")]
616        self.admitted.set(self.admitted.get() + admitted);
617    }
618
619    /// Length of the document `fact` names, or `None` when it names none.
620    ///
621    /// The flat index is authoritative below [`Bm25Index::dense_limit`],
622    /// because every stored document with an id there was written into it. At
623    /// or above that mark it may have holes it never filled, so the answer
624    /// comes from the stored arena — slower, and correct.
625    fn doc_len_of(&self, fact: FactId) -> Option<u16> {
626        let at = fact.0 as usize;
627        if at < self.dense_limit
628            && let Some(&len) = self.doc_len_dense.get(at)
629        {
630            return (len != DOC_LEN_ABSENT).then_some(len as u16);
631        }
632        self.doc_len.get(&fact.0.to_be_bytes()).map(|doc| doc.len)
633    }
634
635    /// Bytes held by the underlying pools.
636    pub fn pool_bytes(&self) -> usize {
637        self.postings.pool_bytes() + self.doc_len.pool_bytes()
638    }
639
640    /// Total token count across the corpus (persisted in the engine
641    /// state).
642    pub fn total_len(&self) -> u64 {
643        self.total_len
644    }
645
646    /// The underlying posting store (the persistence composer dumps it).
647    pub(crate) fn postings(&self) -> &PostingStore<'a, true> {
648        &self.postings
649    }
650
651    /// The per-document length arena (the persistence composer dumps it).
652    pub(crate) fn doc_len_arena(&self) -> &Arena<'a, DocLenSlot> {
653        &self.doc_len
654    }
655
656    /// Assembles an index from already-validated parts (the load path).
657    pub(crate) fn from_parts(
658        postings: PostingStore<'a, true>,
659        doc_len: Arena<'a, DocLenSlot>,
660        total_docs: u64,
661        total_len: u64,
662    ) -> Self {
663        let mut index = Self {
664            postings,
665            doc_len,
666            total_docs,
667            total_len,
668            #[cfg(feature = "counters")]
669            decoded: Cell::new(0),
670            #[cfg(feature = "counters")]
671            scored: Cell::new(0),
672            #[cfg(feature = "counters")]
673            admitted: Cell::new(0),
674            unsummarized: false,
675            doc_len_dense: Vec::new(),
676            dense_limit: usize::MAX,
677        };
678        // One sequential pass over the stored records; the flat index is
679        // derived state and is never part of the image.
680        index.rebuild_dense();
681        index
682    }
683
684    /// Posting entries decoded so far (feature `counters`).
685    #[cfg(feature = "counters")]
686    pub fn decoded(&self) -> u64 {
687        self.decoded.get()
688    }
689
690    /// Documents scored so far — document-length fetches (feature
691    /// `counters`). See the [`Bm25Index::scored`] field docs for why this is
692    /// tracked apart from [`Bm25Index::decoded`].
693    #[cfg(feature = "counters")]
694    pub fn scored(&self) -> u64 {
695        self.scored.get()
696    }
697
698    /// `live` predicate calls so far (feature `counters`).
699    #[cfg(feature = "counters")]
700    pub fn admitted(&self) -> u64 {
701        self.admitted.get()
702    }
703
704    /// Resets the query work counters (feature `counters`).
705    #[cfg(feature = "counters")]
706    pub fn reset_query_counters(&self) {
707        self.decoded.set(0);
708        self.scored.set(0);
709        self.admitted.set(0);
710    }
711}