Skip to main content

zeppelin_embed/fts/
index.rs

1//! The lexical index: per-segment postings and store-wide statistics.
2//!
3//! # Global statistics are the whole point
4//!
5//! [`LexicalIndex`](crate::fts::index::LexicalIndex) owns every segment, and
6//! every statistic the scorer needs
7//! — `N`, `avgdl`, and per-term `df` — is derived across all of them. A
8//! segment cannot answer a scoring question by itself and is not asked to.
9//!
10//! Segment-local IDF is a documented real-world failure (Milvus Lite) and
11//! the prime suspect for the prior engine's ~30% nDCG@10 gap
12//! (`research/02a:256`). The structure here makes it unconstructible: the
13//! only public entry points that produce statistics hang off the store-wide
14//! index, and `prop_engine_bm25_equals_model` randomizes seal boundaries so
15//! any per-segment statistic diverges from the brute-force model on the
16//! first multi-segment case.
17//!
18//! # Field handling (BM25F-lite)
19//!
20//! A document is a map from field id to text. Postings are keyed by
21//! `(term, field)`, so the flat scorer sums a term's contribution across
22//! fields with weight one — reproducing the "flat" BEIR column — and the
23//! multifield scorer applies per-field weights to term frequencies and
24//! lengths. Both run the same code path with a different weight table.
25//!
26//! # No truncation, ever
27//!
28//! Task 13's guardrail is absolute: documents are never truncated before
29//! indexing, and no candidate cap exists outside task 14's bounds. An
30//! oversized field is a typed rejection at the tokenizer boundary, not a
31//! silent cut.
32
33use std::collections::{BTreeMap, HashMap};
34use std::sync::Arc;
35
36use super::bm25::{Bm25Error, CorpusStats};
37use super::postings::{Posting, PostingList, PostingsError};
38use super::sealed::SealedSegment;
39use super::tokenizer::{Analyzer, Token};
40use crate::lifecycle::StoreError;
41use crate::lifecycle::hash_table::allocation_bytes;
42use crate::lifecycle::stats::{AccountedCounter, Accounting, AllocationComponent};
43use crate::meta::DocBitmap;
44
45/// A field identifier within a document schema.
46#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
47pub struct FieldId(pub u16);
48
49/// The default field, used when a caller indexes plain text.
50pub const DEFAULT_FIELD: FieldId = FieldId(0);
51
52/// The key a posting list hangs off: term first, so prefix scans work.
53#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
54pub struct TermKey {
55    /// The analyzed term.
56    pub term: Vec<u8>,
57    /// The field the term occurred in.
58    pub field: FieldId,
59}
60
61/// An index construction failure.
62#[derive(Clone, Debug, Eq, PartialEq)]
63pub enum IndexError {
64    /// The weighted cursor scheduler referenced an absent cursor slot.
65    CursorSlotMissing {
66        /// Invalid scheduler slot.
67        slot: usize,
68    },
69    /// Exact live-frequency memoization refused inconsistent state or a
70    /// poisoned lock. No cache failure is interpreted as a zero frequency.
71    LiveFrequencyCache {
72        /// Failed cache invariant or synchronization boundary.
73        reason: &'static str,
74    },
75    /// A posting list rejected an append.
76    Postings(PostingsError),
77    /// The corpus had no documents or no tokens.
78    Stats(Bm25Error),
79    /// Documents were added out of row order.
80    RowsNotAscending {
81        /// The offending row.
82        row: u32,
83    },
84    /// A live-row counter named a row outside its lexical segment.
85    LiveRowOutOfRange {
86        /// Lexical segment ordinal.
87        segment: usize,
88        /// Invalid row identifier.
89        row: u32,
90        /// Dense row count of the segment.
91        row_count: u32,
92    },
93    /// A sealed segment's dense length counter omitted a valid live row.
94    LiveLengthMissing {
95        /// Lexical segment ordinal.
96        segment: usize,
97        /// Live row without a length counter.
98        row: u32,
99    },
100    /// A replacement named a row the segment does not hold.
101    ReplacedRowOutOfRange {
102        /// Invalid row identifier.
103        row: u32,
104        /// Dense row count of the segment.
105        row_count: u32,
106    },
107    /// A replaced row's previous document names a term the row has no
108    /// posting for, so the index was not built from that document.
109    ReplacedPostingMissing {
110        /// The replaced row.
111        row: u32,
112    },
113}
114
115impl std::fmt::Display for IndexError {
116    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        match self {
118            Self::CursorSlotMissing { slot } => {
119                write!(formatter, "weighted cursor slot {slot} is absent")
120            }
121            Self::LiveFrequencyCache { reason } => {
122                write!(formatter, "live document frequency cache failed: {reason}")
123            }
124            Self::Postings(error) => write!(formatter, "postings rejected: {error}"),
125            Self::Stats(error) => write!(formatter, "corpus statistics unavailable: {error}"),
126            Self::RowsNotAscending { row } => {
127                write!(formatter, "documents must arrive in row order; got {row}")
128            }
129            Self::LiveRowOutOfRange {
130                segment,
131                row,
132                row_count,
133            } => write!(
134                formatter,
135                "live lexical row {row} is outside segment {segment} row count {row_count}"
136            ),
137            Self::LiveLengthMissing { segment, row } => write!(
138                formatter,
139                "live lexical row {row} has no length counter in segment {segment}"
140            ),
141            Self::ReplacedRowOutOfRange { row, row_count } => write!(
142                formatter,
143                "replaced lexical row {row} is outside the segment's {row_count} rows"
144            ),
145            Self::ReplacedPostingMissing { row } => write!(
146                formatter,
147                "replaced lexical row {row} has no posting for a term of its previous document"
148            ),
149        }
150    }
151}
152
153impl std::error::Error for IndexError {}
154
155impl From<PostingsError> for IndexError {
156    fn from(error: PostingsError) -> Self {
157        Self::Postings(error)
158    }
159}
160
161impl From<Bm25Error> for IndexError {
162    fn from(error: Bm25Error) -> Self {
163        Self::Stats(error)
164    }
165}
166
167/// One document's text, by field.
168#[derive(Clone, Debug, Default, Eq, PartialEq)]
169pub struct Document {
170    fields: BTreeMap<FieldId, String>,
171}
172
173impl Document {
174    /// Creates an empty document.
175    #[must_use]
176    pub fn new() -> Self {
177        Self::default()
178    }
179
180    /// Creates a single-field document in [`DEFAULT_FIELD`].
181    #[must_use]
182    pub fn with_text(text: &str) -> Self {
183        let mut document = Self::new();
184        document.set(DEFAULT_FIELD, text);
185        document
186    }
187
188    /// Sets one field's text, replacing any previous value.
189    pub fn set(&mut self, field: FieldId, text: &str) -> &mut Self {
190        self.fields.insert(field, text.to_owned());
191        self
192    }
193
194    /// Returns the fields, ascending by id.
195    pub fn fields(&self) -> impl Iterator<Item = (FieldId, &str)> {
196        self.fields
197            .iter()
198            .map(|(field, text)| (*field, text.as_str()))
199    }
200}
201
202/// One field's analyzed token counts, dense by row.
203///
204/// `lengths` is exactly `row_count` long at all times: a field absent from
205/// a document still occupies its row, carrying zero. That invariant is what
206/// lets the scorer borrow the array wholesale instead of rebuilding it.
207#[derive(Clone, Debug)]
208struct FieldLengths {
209    field: FieldId,
210    lengths: Vec<u32>,
211}
212
213/// Hashes term keys with xxh3, the one permitted hash family.
214#[derive(Clone, Debug, Default)]
215struct Xxh3State;
216
217impl std::hash::BuildHasher for Xxh3State {
218    type Hasher = xxhash_rust::xxh3::Xxh3;
219
220    fn build_hasher(&self) -> Self::Hasher {
221        xxhash_rust::xxh3::Xxh3::new()
222    }
223}
224
225#[cfg(test)]
226thread_local! {
227    static TOKENS_INDEXED: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
228}
229
230/// Tokens this thread has folded into or removed from any segment index.
231///
232/// A deterministic work counter: a gate compares deltas around one write to
233/// prove the write's lexical cost does not depend on the rows beside it.
234#[cfg(test)]
235pub(crate) fn tokens_indexed() -> u64 {
236    TOKENS_INDEXED.with(std::cell::Cell::get)
237}
238
239#[cfg_attr(not(test), allow(unused_variables))]
240fn count_tokens_indexed(tokens: usize) {
241    #[cfg(test)]
242    TOKENS_INDEXED.with(|count| {
243        count.set(
244            count
245                .get()
246                .saturating_add(u64::try_from(tokens).unwrap_or(u64::MAX)),
247        );
248    });
249}
250
251/// One segment's postings and lengths, addressed by dense row id.
252///
253/// # Why lengths are dense arrays and not a map
254///
255/// Rows are dense and ascending by the task 07 invariant, so a
256/// `BTreeMap<(row, field), u32>` spends a pointer-chasing descent — about
257/// seventeen levels at 171,000 documents — on what an array answers with one
258/// index. The scorer performs that lookup once per scored posting, so the
259/// descent was among the largest per-posting costs in the engine.
260///
261/// # Why the term table is a hash map and not a tree
262///
263/// Accumulation looks a term up once per (document, term, field): tens of
264/// millions of probes on a real corpus, against a table with millions of
265/// entries. A B-tree descent is ~20 pointer-chasing key compares per
266/// probe; an xxh3 probe is one hash and usually one compare. The sorted
267/// view the seal needs is materialized ONCE, in [`Self::postings`],
268/// instead of being maintained on every insert.
269#[derive(Debug, Default)]
270pub struct SegmentIndex {
271    /// Term key to slot in `lists`.
272    terms: HashMap<TermKey, usize, Xxh3State>,
273    // Removals can lower HashMap::capacity without freeing any buckets.
274    term_capacity: usize,
275    /// Posting lists, in first-seen order; `terms` holds the slots.
276    lists: Vec<PostingList>,
277    /// Per-field analyzed token counts, ascending by field id.
278    lengths: Vec<FieldLengths>,
279    row_count: u32,
280    /// Slots in `lists` whose term a replacement dropped; new terms reuse
281    /// them before growing `lists`. Each holds an empty list no key names.
282    free_slots: Vec<usize>,
283    /// Heap owned by term keys and postings; updated only for touched terms.
284    term_bytes: usize,
285    posting_bytes: usize,
286}
287
288// Vec cloning can shrink capacities, so a copy measures its own allocations.
289impl Clone for SegmentIndex {
290    fn clone(&self) -> Self {
291        self.clone_with_terms(self.terms.clone())
292    }
293}
294
295impl SegmentIndex {
296    fn clone_with_terms(&self, terms: HashMap<TermKey, usize, Xxh3State>) -> Self {
297        let mut copy = Self {
298            terms,
299            term_capacity: self.term_capacity,
300            lists: self.lists.clone(),
301            lengths: self.lengths.clone(),
302            row_count: self.row_count,
303            free_slots: self.free_slots.clone(),
304            term_bytes: 0,
305            posting_bytes: 0,
306        };
307        copy.term_bytes = copy
308            .terms
309            .keys()
310            .map(|key| key.term.capacity())
311            .fold(0, usize::saturating_add);
312        copy.posting_bytes = copy
313            .lists
314            .iter()
315            .map(PostingList::resident_bytes)
316            .fold(0, usize::saturating_add);
317        copy
318    }
319}
320
321/// Where [`SegmentIndex::accumulate`] places a row's postings.
322#[derive(Clone, Copy)]
323enum Placement {
324    /// The row is the next dense row: append, refusing any lower row.
325    Append,
326    /// The row already exists: insert at its document-id position.
327    Insert,
328}
329
330impl SegmentIndex {
331    pub(crate) fn clone_accounted(
332        &self,
333        accounting: &Arc<Accounting>,
334    ) -> Result<(Self, AccountedCounter), StoreError> {
335        let mut counter = AccountedCounter::new(accounting, AllocationComponent::Active)?;
336        counter.set(self.resident_bytes())?;
337        let mut terms = HashMap::with_hasher(Xxh3State);
338        terms
339            .try_reserve(self.term_capacity)
340            .map_err(|_| StoreError::AllocationFailed {
341                needed: allocation_bytes::<(TermKey, usize)>(self.term_capacity)
342                    .unwrap_or(usize::MAX) as u64,
343                component: "lexical terms map",
344            })?;
345        terms.extend(self.terms.iter().map(|(key, slot)| (key.clone(), *slot)));
346        let copy = self.clone_with_terms(terms);
347        counter.set(copy.resident_bytes())?;
348        Ok((copy, counter))
349    }
350
351    pub(crate) fn reserve_terms(
352        &mut self,
353        capacity: usize,
354        counter: &mut AccountedCounter,
355    ) -> Result<(), StoreError> {
356        // Below half-full, hashbrown only clears deleted control tags in place.
357        let allocation_capacity = if capacity <= self.term_capacity / 2 {
358            0
359        } else {
360            capacity.max(
361                self.term_capacity
362                    .checked_add(1)
363                    .ok_or(StoreError::ActiveRowOverflow)?,
364            )
365        };
366        let needed = allocation_bytes::<(TermKey, usize)>(allocation_capacity)
367            .ok_or(StoreError::ActiveRowOverflow)?;
368        // Keep the old table charged until try_reserve releases it.
369        counter.set(
370            self.resident_bytes()
371                .checked_add(needed)
372                .ok_or(StoreError::ActiveRowOverflow)?,
373        )?;
374        self.terms
375            .try_reserve(capacity.saturating_sub(self.terms.len()))
376            .map_err(|_| StoreError::AllocationFailed {
377                needed: needed as u64,
378                component: "lexical terms map",
379            })?;
380        self.term_capacity = self.term_capacity.max(self.terms.capacity());
381        counter.set(self.resident_bytes())
382    }
383
384    /// Creates an empty segment index.
385    #[must_use]
386    pub fn new() -> Self {
387        Self::default()
388    }
389
390    /// Returns the number of rows.
391    #[must_use]
392    pub const fn row_count(&self) -> u32 {
393        self.row_count
394    }
395
396    /// Returns true when the segment holds no rows.
397    #[must_use]
398    pub const fn is_empty(&self) -> bool {
399        self.row_count == 0
400    }
401
402    /// Returns the posting lists, ordered by term then field.
403    ///
404    /// The sorted view is built here, once per call, rather than being
405    /// maintained on every insert; seal is its one hot caller.
406    pub fn postings(&self) -> impl Iterator<Item = (&TermKey, &PostingList)> {
407        let mut entries: Vec<(&TermKey, usize)> =
408            self.terms.iter().map(|(key, slot)| (key, *slot)).collect();
409        entries.sort_unstable_by(|left, right| left.0.cmp(right.0));
410        entries
411            .into_iter()
412            .filter_map(|(key, slot)| self.lists.get(slot).map(|list| (key, list)))
413    }
414
415    /// Returns one term's posting list within one field.
416    #[must_use]
417    pub fn posting_list(&self, term: &[u8], field: FieldId) -> Option<&PostingList> {
418        let slot = *self.terms.get(&TermKey {
419            term: term.to_vec(),
420            field,
421        })?;
422        self.lists.get(slot)
423    }
424
425    /// Returns the analyzed token count of one row's field.
426    #[must_use]
427    pub fn field_length(&self, row: u32, field: FieldId) -> u32 {
428        self.field_lengths(field)
429            .and_then(|lengths| {
430                usize::try_from(row)
431                    .ok()
432                    .and_then(|slot| lengths.get(slot).copied())
433            })
434            .unwrap_or(0)
435    }
436
437    /// Returns the fields this segment recorded lengths for, ascending.
438    pub fn fields(&self) -> impl Iterator<Item = FieldId> + '_ {
439        self.lengths.iter().map(|entry| entry.field)
440    }
441
442    /// Returns one field's dense length array, if the field was ever set.
443    ///
444    /// The slice is exactly [`Self::row_count`] long, so a caller scoring a
445    /// single unweighted field can borrow it rather than rebuild it. A
446    /// linear scan over the field table is right here: a schema has a
447    /// handful of fields, and a map would trade an index for a descent.
448    #[must_use]
449    pub fn field_lengths(&self, field: FieldId) -> Option<&[u32]> {
450        self.lengths
451            .iter()
452            .find(|entry| entry.field == field)
453            .map(|entry| entry.lengths.as_slice())
454    }
455
456    /// Returns the analyzed token count of one row across every field.
457    #[must_use]
458    pub fn document_length(&self, row: u32) -> u32 {
459        let Ok(slot) = usize::try_from(row) else {
460            return 0;
461        };
462        self.lengths
463            .iter()
464            .filter_map(|entry| entry.lengths.get(slot).copied())
465            .fold(0_u32, u32::saturating_add)
466    }
467
468    /// Returns the segment's total analyzed token count.
469    #[must_use]
470    pub fn total_tokens(&self) -> u64 {
471        self.lengths
472            .iter()
473            .flat_map(|entry| entry.lengths.iter())
474            .map(|length| u64::from(*length))
475            .sum()
476    }
477
478    pub(crate) fn resident_bytes(&self) -> usize {
479        let term_slots =
480            allocation_bytes::<(TermKey, usize)>(self.term_capacity).unwrap_or(usize::MAX);
481        let term_bytes = self.term_bytes;
482        let list_bytes = self
483            .lists
484            .capacity()
485            .saturating_mul(std::mem::size_of::<PostingList>())
486            .saturating_add(self.posting_bytes);
487        let free_slot_bytes = self
488            .free_slots
489            .capacity()
490            .saturating_mul(std::mem::size_of::<usize>());
491        let length_bytes = self
492            .lengths
493            .capacity()
494            .saturating_mul(std::mem::size_of::<FieldLengths>())
495            .saturating_add(
496                self.lengths
497                    .iter()
498                    .map(|entry| entry.lengths.capacity().saturating_mul(4))
499                    .fold(0_usize, usize::saturating_add),
500            );
501        term_slots
502            .saturating_add(term_bytes)
503            .saturating_add(list_bytes)
504            .saturating_add(length_bytes)
505            .saturating_add(free_slot_bytes)
506    }
507
508    /// Independent full scan used to prove the incremental charge stays exact.
509    #[cfg(test)]
510    pub(crate) fn resident_bytes_rescan(&self) -> usize {
511        let term_slots =
512            allocation_bytes::<(TermKey, usize)>(self.term_capacity).unwrap_or(usize::MAX);
513        let term_bytes = self
514            .terms
515            .keys()
516            .map(|key| key.term.capacity())
517            .fold(0_usize, usize::saturating_add);
518        let list_bytes = self
519            .lists
520            .capacity()
521            .saturating_mul(std::mem::size_of::<PostingList>())
522            .saturating_add(
523                self.lists
524                    .iter()
525                    .map(PostingList::resident_bytes)
526                    .fold(0_usize, usize::saturating_add),
527            );
528        let free_slot_bytes = self
529            .free_slots
530            .capacity()
531            .saturating_mul(std::mem::size_of::<usize>());
532        let length_bytes = self
533            .lengths
534            .capacity()
535            .saturating_mul(std::mem::size_of::<FieldLengths>())
536            .saturating_add(
537                self.lengths
538                    .iter()
539                    .map(|entry| entry.lengths.capacity().saturating_mul(4))
540                    .fold(0_usize, usize::saturating_add),
541            );
542        term_slots
543            .saturating_add(term_bytes)
544            .saturating_add(list_bytes)
545            .saturating_add(length_bytes)
546            .saturating_add(free_slot_bytes)
547    }
548
549    /// Records one row's analyzed length for one field.
550    ///
551    /// Pads the field's array with zeros for any earlier row that did not
552    /// carry the field, which is what keeps every array dense.
553    fn set_field_length(&mut self, row: u32, field: FieldId, length: u32) {
554        let Ok(slot) = usize::try_from(row) else {
555            return;
556        };
557        let position = match self.lengths.iter().position(|entry| entry.field == field) {
558            Some(position) => position,
559            None => {
560                // Fields stay ascending so the table has a stable order.
561                let position = self
562                    .lengths
563                    .iter()
564                    .position(|entry| entry.field > field)
565                    .unwrap_or(self.lengths.len());
566                self.lengths.insert(
567                    position,
568                    FieldLengths {
569                        field,
570                        lengths: Vec::new(),
571                    },
572                );
573                position
574            }
575        };
576        let Some(entry) = self.lengths.get_mut(position) else {
577            return;
578        };
579        if let Some(existing) = entry.lengths.get_mut(slot) {
580            *existing = length;
581        } else {
582            entry.lengths.resize(slot, 0);
583            entry.lengths.push(length);
584        }
585    }
586
587    /// Extends every field array to cover `rows` rows.
588    fn pad_lengths_to(&mut self, rows: usize) {
589        for entry in &mut self.lengths {
590            entry.lengths.resize(rows, 0);
591        }
592    }
593
594    /// Appends one analyzed document as the next dense row.
595    ///
596    /// # Errors
597    ///
598    /// Returns [`IndexError::Postings`] when the analyzer produces a token
599    /// stream a posting list refuses, which would be a tokenizer contract
600    /// violation rather than bad input.
601    pub fn push_document(
602        &mut self,
603        analyzer: &Analyzer,
604        document: &Document,
605    ) -> Result<u32, IndexError> {
606        self.push_document_controlled(analyzer, document, &mut |_, _| Ok(()))
607    }
608
609    pub(crate) fn push_document_controlled<E: From<IndexError>>(
610        &mut self,
611        analyzer: &Analyzer,
612        document: &Document,
613        reserve: &mut impl FnMut(&mut Self, usize) -> Result<(), E>,
614    ) -> Result<u32, E> {
615        let row = self.row_count;
616        for (field, text) in document.fields() {
617            let tokens = analyzer.analyze(text);
618            self.accumulate(row, field, &tokens, Placement::Append, reserve)?;
619        }
620        self.finish_row();
621        Ok(row)
622    }
623
624    /// Replaces one existing row's document in place.
625    ///
626    /// `previous` must be the document the row was indexed from. The row's
627    /// postings for every term of `previous` are removed, a term left with
628    /// no posting is dropped, the row's lengths are reset, and `replacement`
629    /// is folded in at the same row. The result answers exactly as a
630    /// segment rebuilt with the row's document swapped, and the cost is
631    /// proportional to the two documents, not to the segment.
632    ///
633    /// # Errors
634    ///
635    /// Returns [`IndexError::ReplacedRowOutOfRange`] for a row the segment
636    /// does not hold, [`IndexError::ReplacedPostingMissing`] when `previous`
637    /// is not what the row was indexed from, and [`IndexError::Postings`]
638    /// as [`Self::push_document`]. On an error the segment is partly
639    /// updated; the caller discards it.
640    pub fn replace_document(
641        &mut self,
642        analyzer: &Analyzer,
643        row: u32,
644        previous: &Document,
645        replacement: &Document,
646    ) -> Result<(), IndexError> {
647        self.replace_document_controlled(analyzer, row, previous, replacement, &mut |_, _| Ok(()))
648    }
649
650    pub(crate) fn replace_document_controlled<E: From<IndexError>>(
651        &mut self,
652        analyzer: &Analyzer,
653        row: u32,
654        previous: &Document,
655        replacement: &Document,
656        reserve: &mut impl FnMut(&mut Self, usize) -> Result<(), E>,
657    ) -> Result<(), E> {
658        let out_of_range = IndexError::ReplacedRowOutOfRange {
659            row,
660            row_count: self.row_count,
661        };
662        if row >= self.row_count {
663            return Err(out_of_range.into());
664        }
665        let slot = usize::try_from(row).map_err(|_| out_of_range)?;
666        for (field, text) in previous.fields() {
667            let tokens = analyzer.analyze(text);
668            count_tokens_indexed(tokens.len());
669            let mut terms: Vec<&str> = tokens.iter().map(|token| token.term.as_str()).collect();
670            terms.sort_unstable();
671            terms.dedup();
672            let mut probe = TermKey {
673                term: Vec::new(),
674                field,
675            };
676            for term in terms {
677                probe.term.clear();
678                probe.term.extend_from_slice(term.as_bytes());
679                let missing = IndexError::ReplacedPostingMissing { row };
680                let list_slot = self.terms.get(&probe).copied().ok_or(missing.clone())?;
681                let list = self.lists.get_mut(list_slot).ok_or(missing.clone())?;
682                let posting_slot = list
683                    .postings()
684                    .binary_search_by_key(&row, |posting| posting.docid)
685                    .map_err(|_| missing.clone())?;
686                let position_bytes = list
687                    .postings()
688                    .get(posting_slot)
689                    .ok_or(missing.clone())?
690                    .positions
691                    .capacity()
692                    .saturating_mul(4);
693                crate::fts::postings::count_accounting_visits(1);
694                if !list.remove(row) {
695                    return Err(missing.into());
696                }
697                self.posting_bytes = self.posting_bytes.saturating_sub(position_bytes);
698                if list.is_empty() {
699                    self.posting_bytes = self.posting_bytes.saturating_sub(list.allocation_bytes());
700                    *list = PostingList::new();
701                    if let Some((key, _)) = self.terms.remove_entry(&probe) {
702                        self.term_bytes = self.term_bytes.saturating_sub(key.term.capacity());
703                    }
704                    self.free_slots.push(list_slot);
705                }
706            }
707        }
708        for entry in &mut self.lengths {
709            if let Some(length) = entry.lengths.get_mut(slot) {
710                *length = 0;
711            }
712        }
713        for (field, text) in replacement.fields() {
714            let tokens = analyzer.analyze(text);
715            self.accumulate(row, field, &tokens, Placement::Insert, reserve)?;
716        }
717        // A field new to the segment must still cover every row.
718        self.pad_lengths_to(usize::try_from(self.row_count).unwrap_or(usize::MAX));
719        Ok(())
720    }
721
722    /// Appends many documents, analyzing them in parallel.
723    ///
724    /// The result is identical to pushing the same documents one at a time
725    /// in slice order: analysis is independent per document and runs on
726    /// explicit scoped threads, while accumulation happens strictly in row
727    /// order on the calling thread — so nothing about the outcome depends
728    /// on scheduling, and the single-writer invariant is untouched.
729    ///
730    /// `threads` is the caller's policy; one or fewer runs sequentially.
731    /// Analysis is batched so the in-flight token working set stays
732    /// proportional to the batch, never to the corpus.
733    ///
734    /// # Errors
735    ///
736    /// As [`Self::push_document`]. On an error, documents before the
737    /// failing one remain indexed, exactly as a sequential loop would
738    /// leave them.
739    pub fn push_documents(
740        &mut self,
741        analyzer: &Analyzer,
742        documents: &[Document],
743        threads: usize,
744    ) -> Result<(), IndexError> {
745        let workers = threads.max(1).min(documents.len().max(1));
746        if workers <= 1 {
747            for document in documents {
748                self.push_document(analyzer, document)?;
749            }
750            return Ok(());
751        }
752        let batch_size = workers.saturating_mul(32).max(1);
753        for batch in documents.chunks(batch_size) {
754            let mut analyzed: Vec<Vec<(FieldId, Vec<Token>)>> = Vec::new();
755            analyzed.resize_with(batch.len(), Vec::new);
756            let stride = batch.len().div_ceil(workers);
757            std::thread::scope(|scope| {
758                let mut rest = analyzed.as_mut_slice();
759                let mut offset = 0_usize;
760                while !rest.is_empty() {
761                    let take = stride.min(rest.len());
762                    let (mine, tail) = rest.split_at_mut(take);
763                    rest = tail;
764                    let end = offset.saturating_add(take);
765                    let docs = batch.get(offset..end).unwrap_or(&[]);
766                    offset = end;
767                    scope.spawn(move || {
768                        for (slot, document) in mine.iter_mut().zip(docs) {
769                            *slot = document
770                                .fields()
771                                .map(|(field, text)| (field, analyzer.analyze(text)))
772                                .collect();
773                        }
774                    });
775                }
776            });
777            for fields in &analyzed {
778                let row = self.row_count;
779                for (field, tokens) in fields {
780                    self.accumulate(row, *field, tokens, Placement::Append, &mut |_, _| {
781                        Ok::<(), IndexError>(())
782                    })?;
783                }
784                self.finish_row();
785            }
786        }
787        Ok(())
788    }
789
790    /// Folds one field's analyzed tokens into the postings for one row.
791    fn accumulate<E: From<IndexError>>(
792        &mut self,
793        row: u32,
794        field: FieldId,
795        tokens: &[Token],
796        placement: Placement,
797        reserve: &mut impl FnMut(&mut Self, usize) -> Result<(), E>,
798    ) -> Result<(), E> {
799        count_tokens_indexed(tokens.len());
800        // Position count is the analyzed length: stacked variants share
801        // a position and must not inflate the document length, or avgdl
802        // stops matching the unit being scored.
803        let length = tokens
804            .iter()
805            .map(|token| token.position)
806            .max()
807            .map_or(0, |highest| highest.saturating_add(1));
808        self.set_field_length(row, field, length);
809
810        // Sort-based grouping over borrowed terms replaces the per-document
811        // `BTreeMap<Vec<u8>, Vec<u32>>`: no map to build and tear down per
812        // document, no allocation per (document, term). `str` compares
813        // byte-wise, so groups arrive in exactly the order the map's keys
814        // did, and a group's positions arrive already sorted.
815        let mut scratch: Vec<(&str, u32)> = tokens
816            .iter()
817            .map(|token| (token.term.as_str(), token.position))
818            .collect();
819        scratch.sort_unstable();
820
821        // One reusable probe key per field: lookups borrow it, and only a
822        // term new to the segment pays an owned copy.
823        let mut probe = TermKey {
824            term: Vec::new(),
825            field,
826        };
827        let mut start = 0_usize;
828        while let Some((term, _)) = scratch.get(start).copied() {
829            let mut end = start.saturating_add(1);
830            while scratch.get(end).is_some_and(|(next, _)| *next == term) {
831                end = end.saturating_add(1);
832            }
833            let group = scratch.get(start..end).unwrap_or(&[]);
834            let mut positions: Vec<u32> = Vec::with_capacity(group.len());
835            for (_, position) in group {
836                if positions.last() != Some(position) {
837                    positions.push(*position);
838                }
839            }
840            let tf = u32::try_from(positions.len()).unwrap_or(u32::MAX);
841            probe.term.clear();
842            probe.term.extend_from_slice(term.as_bytes());
843            let posting = Posting {
844                docid: row,
845                tf,
846                positions,
847            };
848            let position_bytes = posting.positions.capacity().saturating_mul(4);
849            crate::fts::postings::count_accounting_visits(1);
850            match self.terms.get(&probe).copied() {
851                Some(slot) => {
852                    // The map only ever hands out slots minted below, so
853                    // the lookup cannot miss; the `if let` is the panic-free
854                    // spelling of that impossibility.
855                    if let Some(list) = self.lists.get_mut(slot) {
856                        let previous = list.allocation_bytes();
857                        match placement {
858                            Placement::Append => list.push(posting).map_err(IndexError::from)?,
859                            Placement::Insert => list.insert(posting).map_err(IndexError::from)?,
860                        }
861                        self.posting_bytes = self
862                            .posting_bytes
863                            .saturating_sub(previous)
864                            .saturating_add(list.allocation_bytes())
865                            .saturating_add(position_bytes);
866                    }
867                }
868                None => {
869                    let grows = self.terms.len() == self.terms.capacity();
870                    if grows {
871                        reserve(self, self.terms.len().saturating_add(1))?;
872                    }
873                    let mut list = PostingList::new();
874                    list.push(posting).map_err(IndexError::from)?;
875                    self.posting_bytes = self
876                        .posting_bytes
877                        .saturating_add(list.allocation_bytes())
878                        .saturating_add(position_bytes);
879                    let slot = if let Some(slot) = self.free_slots.pop()
880                        && let Some(free) = self.lists.get_mut(slot)
881                    {
882                        *free = list;
883                        slot
884                    } else {
885                        let slot = self.lists.len();
886                        self.lists.push(list);
887                        slot
888                    };
889                    let key = TermKey {
890                        term: probe.term.clone(),
891                        field,
892                    };
893                    self.term_bytes = self.term_bytes.saturating_add(key.term.capacity());
894                    self.terms.insert(key, slot);
895                    if grows {
896                        self.term_capacity = self.term_capacity.max(self.terms.capacity());
897                    }
898                }
899            }
900            start = end;
901        }
902        Ok(())
903    }
904
905    /// Seals the row: every field array is padded to cover it.
906    fn finish_row(&mut self) {
907        self.row_count = self.row_count.saturating_add(1);
908        self.pad_lengths_to(usize::try_from(self.row_count).unwrap_or(usize::MAX));
909    }
910}
911
912/// The store-wide lexical index across every segment.
913///
914/// # Segments here are SEALED
915///
916/// [`push_segment`](Self::push_segment) encodes the active segment into the
917/// persisted posting layout and keeps only that. There is therefore exactly
918/// one scoring path in the engine — over bytes — rather than an in-memory
919/// one that ships and a persisted one that is only ever tested.
920#[derive(Clone, Debug, Default)]
921pub struct LexicalIndex {
922    segments: Vec<Arc<SealedSegment>>,
923    statistics: Vec<Arc<LiveSegmentStatistics>>,
924    has_frequency_cache: bool,
925}
926
927#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
928struct LiveSegmentCounters {
929    documents: u64,
930    tokens: u64,
931}
932
933/// Validated live membership and counters, independent of source ordinal.
934#[derive(Clone, Debug)]
935pub(crate) struct LiveSegmentStatistics {
936    pub(crate) rows: DocBitmap,
937    counters: LiveSegmentCounters,
938    pub(crate) frequency_cache: Option<Arc<super::live_df::LiveFrequencyCache>>,
939}
940
941impl LiveSegmentStatistics {
942    pub(crate) fn build_controlled<E: From<IndexError>>(
943        ordinal: usize,
944        segment: &SealedSegment,
945        rows: &DocBitmap,
946        work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
947    ) -> Result<Self, E> {
948        work.check_now()?;
949        let value = Self {
950            counters: live_segment_counters_controlled(ordinal, segment, rows, work)?,
951            rows: rows.clone(),
952            frequency_cache: None,
953        };
954        work.check_now()?;
955        Ok(value)
956    }
957
958    pub(crate) fn build(
959        ordinal: usize,
960        segment: &SealedSegment,
961        rows: &DocBitmap,
962    ) -> Result<Self, IndexError> {
963        Self::build_controlled(
964            ordinal,
965            segment,
966            rows,
967            &mut super::control::WorkCheck::new(|| Ok(())),
968        )
969    }
970
971    pub(crate) fn allocation_bytes(rows: &DocBitmap) -> Option<usize> {
972        std::mem::size_of::<Self>()
973            .checked_add(2 * std::mem::size_of::<usize>())?
974            .checked_add(rows.resident_bytes()?)
975    }
976}
977
978impl LexicalIndex {
979    /// Creates an empty index.
980    #[must_use]
981    pub fn new() -> Self {
982        Self::default()
983    }
984
985    /// Seals one active segment and appends it.
986    ///
987    /// # Errors
988    ///
989    /// Returns [`IndexError::Postings`] when the encoder rejects the
990    /// segment's geometry or refuses to read back what it just wrote. Both
991    /// are broken-writer conditions, and they surface here rather than
992    /// becoming a decode fault inside a later query.
993    pub fn push_segment(&mut self, segment: SegmentIndex) -> Result<(), IndexError> {
994        self.push_sealed(SealedSegment::seal(&segment)?);
995        Ok(())
996    }
997
998    /// Appends an already-sealed segment.
999    pub fn push_sealed(&mut self, segment: SealedSegment) {
1000        self.push_shared(Arc::new(segment));
1001    }
1002
1003    fn push_shared(&mut self, segment: Arc<SealedSegment>) {
1004        self.statistics.push(Arc::new(LiveSegmentStatistics {
1005            rows: DocBitmap::full(segment.row_count()),
1006            frequency_cache: None,
1007            counters: LiveSegmentCounters {
1008                documents: u64::from(segment.row_count()),
1009                tokens: segment.total_tokens(),
1010            },
1011        }));
1012        self.segments.push(segment);
1013    }
1014
1015    /// Appends an already-sealed segment with exact live-row counters.
1016    ///
1017    /// The row set is the segment tombstone state's live bitmap, not a query
1018    /// filter. Capturing it here keeps BM25 corpus statistics independent of
1019    /// later metadata allow-lists while avoiding any posting-list rewrite.
1020    ///
1021    /// # Errors
1022    ///
1023    /// Returns [`IndexError::LiveRowOutOfRange`] when the tombstone state
1024    /// names a row outside the sealed segment, or
1025    /// [`IndexError::LiveLengthMissing`] when a sealed length counter violates
1026    /// its dense-row invariant.
1027    pub fn push_sealed_with_live_rows(
1028        &mut self,
1029        segment: SealedSegment,
1030        live_rows: &DocBitmap,
1031    ) -> Result<(), IndexError> {
1032        self.push_shared_with_live_rows(Arc::new(segment), live_rows)
1033    }
1034
1035    pub(crate) fn push_shared_with_live_rows(
1036        &mut self,
1037        segment: Arc<SealedSegment>,
1038        live_rows: &DocBitmap,
1039    ) -> Result<(), IndexError> {
1040        let statistics = Arc::new(LiveSegmentStatistics::build(
1041            self.segments.len(),
1042            &segment,
1043            live_rows,
1044        )?);
1045        self.push_shared_with_statistics(segment, statistics);
1046        Ok(())
1047    }
1048
1049    /// Only a validated contribution can enter this path; callers cannot supply
1050    /// independent counters or change the membership behind the shared value.
1051    pub(crate) fn push_shared_with_statistics(
1052        &mut self,
1053        segment: Arc<SealedSegment>,
1054        statistics: Arc<LiveSegmentStatistics>,
1055    ) {
1056        self.has_frequency_cache |= statistics.frequency_cache.is_some();
1057        self.segments.push(segment);
1058        self.statistics.push(statistics);
1059    }
1060
1061    /// Returns the segments in seal order.
1062    #[must_use]
1063    pub fn segments(&self) -> &[Arc<SealedSegment>] {
1064        &self.segments
1065    }
1066
1067    /// Bytes for the index spines. Cached contributions own and account their
1068    /// shared membership/counters separately for the full Arc lifetime.
1069    pub(crate) fn resident_spine_bytes(&self) -> Option<usize> {
1070        self.segments
1071            .capacity()
1072            .checked_mul(std::mem::size_of::<Arc<SealedSegment>>())?
1073            .checked_add(
1074                self.statistics
1075                    .capacity()
1076                    .checked_mul(std::mem::size_of::<Arc<LiveSegmentStatistics>>())?,
1077            )
1078    }
1079
1080    /// Iterates the analyzed vocabulary of every segment.
1081    ///
1082    /// Terms may repeat across segments; callers that expand a structured
1083    /// query must impose their own deterministic de-duplication.
1084    pub fn terms(&self) -> impl Iterator<Item = &[u8]> {
1085        self.segments.iter().flat_map(|segment| segment.terms())
1086    }
1087
1088    pub(crate) fn vocabulary_terms(&self) -> impl Iterator<Item = (&[u8], FieldId)> + Clone {
1089        self.segments
1090            .iter()
1091            .flat_map(|segment| segment.vocabulary_terms())
1092    }
1093
1094    /// Returns the store-wide document count.
1095    #[must_use]
1096    pub fn document_count(&self) -> u64 {
1097        self.statistics
1098            .iter()
1099            .map(|statistics| statistics.counters.documents)
1100            .fold(0_u64, u64::saturating_add)
1101    }
1102
1103    /// Returns the store-wide analyzed token count.
1104    #[must_use]
1105    pub fn total_tokens(&self) -> u64 {
1106        self.statistics
1107            .iter()
1108            .map(|statistics| statistics.counters.tokens)
1109            .fold(0_u64, u64::saturating_add)
1110    }
1111
1112    /// Returns the store-wide corpus statistics.
1113    ///
1114    /// This is the ONLY source of `N` and `avgdl` for scoring. There is
1115    /// deliberately no per-segment equivalent.
1116    ///
1117    /// # Errors
1118    ///
1119    /// Returns [`IndexError::Stats`] when the index holds no documents or
1120    /// no tokens, because `avgdl` is undefined then.
1121    pub fn corpus_stats(&self) -> Result<CorpusStats, IndexError> {
1122        #[cfg(any(test, feature = "test-seams"))]
1123        super::preparation_observer::corpus_statistics();
1124        Ok(CorpusStats::new(
1125            self.document_count(),
1126            self.total_tokens(),
1127        )?)
1128    }
1129
1130    /// Returns the store-wide document frequency of one term.
1131    ///
1132    /// Summed across every segment and every requested field. A document
1133    /// containing the term in two fields counts once, which is what makes
1134    /// this a *document* frequency rather than a posting count. Segments
1135    /// without tombstones answer from their sealed dictionary.
1136    /// Segments with tombstones decode document identifiers and count only
1137    /// rows in the same live bitmap used for `N` and `avgdl`.
1138    #[must_use]
1139    pub fn document_frequency(&self, term: &[u8], fields: &[FieldId]) -> u32 {
1140        #[cfg(any(test, feature = "test-seams"))]
1141        super::preparation_observer::frequency();
1142        self.segments
1143            .iter()
1144            .zip(&self.statistics)
1145            .map(|(segment, statistics)| {
1146                let live_rows = &statistics.rows;
1147                if live_rows.cardinality() == u64::from(segment.row_count()) {
1148                    segment.document_frequency(term, fields)
1149                } else {
1150                    segment.live_document_frequency(term, fields, live_rows)
1151                }
1152            })
1153            .fold(0_u32, u32::saturating_add)
1154    }
1155
1156    /// Scoring can report cache synchronization errors without changing the
1157    /// public infallible, uncached statistics inspection method. Callers pass
1158    /// the canonical field list from FieldWeights.
1159    #[inline]
1160    pub(crate) fn prepared_document_frequency(
1161        &self,
1162        term: &[u8],
1163        fields: &[FieldId],
1164    ) -> Result<u32, IndexError> {
1165        // Preserve the original dictionary/statistics loop when no contribution
1166        // owns a cache. Cache synchronization and fallible accumulation belong
1167        // only to snapshots that can actually use them.
1168        if !self.has_frequency_cache {
1169            return Ok(self.document_frequency(term, fields));
1170        }
1171        self.cached_document_frequency(term, fields)
1172    }
1173
1174    pub(crate) fn prepared_document_frequency_controlled<E>(
1175        &self,
1176        term: &[u8],
1177        fields: &[FieldId],
1178        mut check: impl FnMut() -> Result<(), E>,
1179    ) -> Result<u32, super::search::ControlledSearchError<E>> {
1180        use super::search::ControlledSearchError;
1181        #[cfg(any(test, feature = "test-seams"))]
1182        super::preparation_observer::frequency();
1183        check().map_err(ControlledSearchError::Control)?;
1184        let mut total = 0_u32;
1185        for (segment, statistics) in self.segments.iter().zip(&self.statistics) {
1186            let live = &statistics.rows;
1187            let frequency = if live.cardinality() == u64::from(segment.row_count()) {
1188                segment
1189                    .document_frequency_controlled(term, fields, &mut check)
1190                    .map_err(ControlledSearchError::Control)?
1191            } else if let Some(cache) = &statistics.frequency_cache {
1192                cache.frequency_controlled(segment, live, term, fields, &mut check)?
1193            } else {
1194                segment
1195                    .live_document_frequency_controlled(term, fields, live, &mut check)
1196                    .map_err(ControlledSearchError::Control)?
1197            };
1198            total = total.saturating_add(frequency);
1199        }
1200        check().map_err(ControlledSearchError::Control)?;
1201        Ok(total)
1202    }
1203
1204    fn cached_document_frequency(
1205        &self,
1206        term: &[u8],
1207        fields: &[FieldId],
1208    ) -> Result<u32, IndexError> {
1209        #[cfg(any(test, feature = "test-seams"))]
1210        super::preparation_observer::frequency();
1211        self.segments.iter().zip(&self.statistics).try_fold(
1212            0_u32,
1213            |total, (segment, statistics)| {
1214                let live = &statistics.rows;
1215                let frequency = if live.cardinality() == u64::from(segment.row_count()) {
1216                    segment.document_frequency(term, fields)
1217                } else if let Some(cache) = &statistics.frequency_cache {
1218                    cache.frequency(segment, live, term, fields)?
1219                } else {
1220                    segment.live_document_frequency(term, fields, live)
1221                };
1222                Ok(total.saturating_add(frequency))
1223            },
1224        )
1225    }
1226}
1227
1228fn live_segment_counters_controlled<E: From<IndexError>>(
1229    ordinal: usize,
1230    segment: &SealedSegment,
1231    live_rows: &DocBitmap,
1232    work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1233) -> Result<LiveSegmentCounters, E> {
1234    work.check_now()?;
1235    if let Some(row) = live_rows.first_at_or_after(segment.row_count()) {
1236        return Err(IndexError::LiveRowOutOfRange {
1237            segment: ordinal,
1238            row,
1239            row_count: segment.row_count(),
1240        }
1241        .into());
1242    }
1243    let tokens =
1244        live_rows.iter().try_fold(0_u64, |total, row| {
1245            work.step()?;
1246            #[cfg(any(test, feature = "test-seams"))]
1247            super::preparation_observer::live_statistics_row();
1248            let slot = usize::try_from(row).map_err(|_| IndexError::LiveLengthMissing {
1249                segment: ordinal,
1250                row,
1251            })?;
1252            let length = segment.total_lengths().get(slot).copied().ok_or(
1253                IndexError::LiveLengthMissing {
1254                    segment: ordinal,
1255                    row,
1256                },
1257            )?;
1258            Ok::<u64, E>(total.saturating_add(u64::from(length)))
1259        })?;
1260    work.check_now()?;
1261    Ok(LiveSegmentCounters {
1262        documents: live_rows.cardinality(),
1263        tokens,
1264    })
1265}
1266
1267#[cfg(test)]
1268#[allow(
1269    clippy::expect_used,
1270    clippy::indexing_slicing,
1271    clippy::panic,
1272    clippy::unwrap_used
1273)]
1274mod tests {
1275
1276    #[test]
1277    fn terms_table_charge_includes_empty_buckets_and_controls() {
1278        let mut index = SegmentIndex::new();
1279        index
1280            .push_document(&analyzer(), &Document::with_text("alpha"))
1281            .expect("index");
1282        let group = if cfg!(all(
1283            any(target_arch = "x86", target_arch = "x86_64"),
1284            target_feature = "sse2",
1285            not(miri)
1286        )) {
1287            16
1288        } else {
1289            8
1290        };
1291        let other = index.term_bytes
1292            + index.posting_bytes
1293            + index.lists.capacity() * std::mem::size_of::<PostingList>()
1294            + index.lengths.capacity() * std::mem::size_of::<FieldLengths>()
1295            + index
1296                .lengths
1297                .iter()
1298                .map(|entry| entry.lengths.capacity() * 4)
1299                .sum::<usize>();
1300        assert_eq!(
1301            index.resident_bytes() - other,
1302            4 * std::mem::size_of::<(TermKey, usize)>() + 4 + group
1303        );
1304    }
1305
1306    #[test]
1307    fn parallel_analysis_is_identical_to_the_sequential_loop() {
1308        // push_documents claims scheduling independence; this holds it to
1309        // the sequential oracle exactly -- every posting, every position,
1310        // every length array -- across thread counts and a document count
1311        // chosen to leave a ragged final batch.
1312        let analyzer = analyzer();
1313        let documents: Vec<Document> = (0..97)
1314            .map(|index| {
1315                let mut document = Document::new();
1316                document.set(
1317                    FieldId(0),
1318                    &format!("alpha beta{} gamma-{index} 401k", index % 7),
1319                );
1320                document.set(FieldId(1), &format!("delta epsilon{} the quick", index % 5));
1321                document
1322            })
1323            .collect();
1324        let mut sequential = SegmentIndex::new();
1325        for document in &documents {
1326            sequential
1327                .push_document(&analyzer, document)
1328                .expect("indexable");
1329        }
1330        for threads in [1_usize, 3, 8] {
1331            let mut parallel = SegmentIndex::new();
1332            parallel
1333                .push_documents(&analyzer, &documents, threads)
1334                .expect("indexable");
1335            assert_eq!(parallel.row_count(), sequential.row_count());
1336            let ours: Vec<_> = parallel.postings().collect();
1337            let oracle: Vec<_> = sequential.postings().collect();
1338            assert_eq!(
1339                ours.len(),
1340                oracle.len(),
1341                "term count diverged at {threads} threads"
1342            );
1343            for ((key, list), (oracle_key, oracle_list)) in ours.iter().zip(&oracle) {
1344                assert_eq!(key, oracle_key);
1345                assert_eq!(
1346                    list.postings(),
1347                    oracle_list.postings(),
1348                    "postings diverged for {:?} at {threads} threads",
1349                    String::from_utf8_lossy(&key.term)
1350                );
1351            }
1352            for field in [FieldId(0), FieldId(1)] {
1353                assert_eq!(
1354                    parallel.field_lengths(field),
1355                    sequential.field_lengths(field),
1356                    "lengths diverged at {threads} threads"
1357                );
1358            }
1359        }
1360    }
1361
1362    use super::*;
1363    use crate::fts::tokenizer::{Profile, TokenizerConfig};
1364
1365    #[test]
1366    fn astra_18_live_statistics_cancel_inside_row_walk() {
1367        #[derive(Debug)]
1368        enum Error {
1369            Index(IndexError),
1370            Cancelled,
1371        }
1372        impl From<IndexError> for Error {
1373            fn from(error: IndexError) -> Self {
1374                Self::Index(error)
1375            }
1376        }
1377        let segment = SealedSegment::seal(&segment_of(&analyzer(), &vec!["alpha beta"; 4_097]))
1378            .expect("sealed row fixture");
1379        let live = DocBitmap::from_ids(0..4_096);
1380        super::super::preparation_observer::begin();
1381        let mut work = super::super::control::WorkCheck::new(|| {
1382            if super::super::preparation_observer::live_statistics_rows() >= 64 {
1383                Err(Error::Cancelled)
1384            } else {
1385                Ok(())
1386            }
1387        });
1388        let result = LiveSegmentStatistics::build_controlled(7, &segment, &live, &mut work);
1389        let visited = super::super::preparation_observer::live_statistics_rows();
1390        super::super::preparation_observer::take();
1391        println!("live statistics rows={visited}");
1392        assert!(matches!(result, Err(Error::Cancelled)));
1393        assert!((64..=128).contains(&visited));
1394        let mut clean = super::super::control::WorkCheck::new(|| Ok::<(), Error>(()));
1395        let value = LiveSegmentStatistics::build_controlled(7, &segment, &live, &mut clean)
1396            .expect("fresh control");
1397        assert_eq!(value.counters.documents, 4_096);
1398        assert_eq!(value.counters.tokens, 8_192);
1399        assert_eq!(value.rows, live);
1400        assert!(value.frequency_cache.is_none());
1401        assert!(matches!(
1402            LiveSegmentStatistics::build_controlled(
1403                7,
1404                &segment,
1405                &DocBitmap::from_ids([4_097]),
1406                &mut clean
1407            ),
1408            Err(Error::Index(IndexError::LiveRowOutOfRange {
1409                segment: 7,
1410                row: 4_097,
1411                row_count: 4_097
1412            }))
1413        ));
1414    }
1415
1416    fn analyzer() -> Analyzer {
1417        Analyzer::new(TokenizerConfig::text_default()).expect("valid config")
1418    }
1419
1420    #[test]
1421    fn replace_document_refuses_a_row_or_previous_text_the_index_does_not_hold() {
1422        let analyzer = analyzer();
1423        let mut index = SegmentIndex::new();
1424        index
1425            .push_document(&analyzer, &Document::with_text("alpha beta"))
1426            .expect("indexable");
1427        let replacement = Document::with_text("gamma");
1428
1429        assert_eq!(
1430            index.replace_document(&analyzer, 1, &Document::with_text("alpha"), &replacement),
1431            Err(IndexError::ReplacedRowOutOfRange {
1432                row: 1,
1433                row_count: 1
1434            })
1435        );
1436        assert_eq!(
1437            index.replace_document(&analyzer, 0, &Document::with_text("delta"), &replacement),
1438            Err(IndexError::ReplacedPostingMissing { row: 0 })
1439        );
1440        // A term the index holds, but not for this row.
1441        index
1442            .push_document(&analyzer, &Document::with_text("epsilon"))
1443            .expect("indexable");
1444        assert_eq!(
1445            index.replace_document(&analyzer, 0, &Document::with_text("epsilon"), &replacement),
1446            Err(IndexError::ReplacedPostingMissing { row: 0 })
1447        );
1448        assert!(
1449            IndexError::ReplacedPostingMissing { row: 0 }
1450                .to_string()
1451                .contains("row 0")
1452        );
1453        assert!(
1454            IndexError::ReplacedRowOutOfRange {
1455                row: 1,
1456                row_count: 1
1457            }
1458            .to_string()
1459            .contains("1 rows")
1460        );
1461    }
1462
1463    #[test]
1464    fn replace_document_reuses_a_dropped_term_slot_and_pads_a_new_field() {
1465        let analyzer = analyzer();
1466        let mut index = SegmentIndex::new();
1467        for text in ["alpha", "beta", "alpha delta"] {
1468            index
1469                .push_document(&analyzer, &Document::with_text(text))
1470                .expect("indexable");
1471        }
1472        let mut replacement = Document::with_text("alpha");
1473        replacement.set(FieldId(1), "gamma");
1474        index
1475            .replace_document(&analyzer, 1, &Document::with_text("beta"), &replacement)
1476            .expect("replace row 1");
1477        // "gamma" took the slot "beta" freed; no list was appended.
1478        assert_eq!(index.lists.len(), 3);
1479        assert!(index.posting_list(b"beta", DEFAULT_FIELD).is_none());
1480        assert!(index.free_slots.is_empty());
1481        assert_eq!(
1482            index.posting_list(b"alpha", DEFAULT_FIELD).map(|list| list
1483                .postings()
1484                .iter()
1485                .map(|p| p.docid)
1486                .collect::<Vec<_>>()),
1487            Some(vec![0, 1, 2])
1488        );
1489        assert_eq!(index.field_lengths(FieldId(1)), Some(&[0, 1, 0][..]));
1490        index
1491            .replace_document(
1492                &analyzer,
1493                2,
1494                &Document::with_text("alpha delta"),
1495                &Document::new(),
1496            )
1497            .expect("replace row 2");
1498        assert!(index.posting_list(b"delta", DEFAULT_FIELD).is_none());
1499        assert_eq!(index.free_slots.len(), 1);
1500        assert_eq!(index.total_tokens(), 3);
1501    }
1502
1503    #[test]
1504    fn posting_list_insert_refuses_a_duplicate_row() {
1505        let mut list = PostingList::new();
1506        let posting = |docid| Posting {
1507            docid,
1508            tf: 1,
1509            positions: vec![0],
1510        };
1511        list.push(posting(0)).expect("append");
1512        list.push(posting(4)).expect("append");
1513        list.insert(posting(2)).expect("insert between");
1514        assert_eq!(
1515            list.insert(posting(2)),
1516            Err(PostingsError::DocidsNotAscending { docid: 2 })
1517        );
1518        assert_eq!(
1519            list.insert(Posting {
1520                docid: 3,
1521                tf: 0,
1522                positions: Vec::new()
1523            }),
1524            Err(PostingsError::ZeroTermFrequency)
1525        );
1526        assert_eq!(
1527            list.insert(Posting {
1528                docid: 3,
1529                tf: 2,
1530                positions: vec![1, 1]
1531            }),
1532            Err(PostingsError::PositionsNotAscending { position: 1 })
1533        );
1534        assert!(list.remove(2));
1535        assert!(!list.remove(2));
1536        assert_eq!(
1537            list.postings().iter().map(|p| p.docid).collect::<Vec<_>>(),
1538            vec![0, 4]
1539        );
1540    }
1541
1542    fn code_analyzer() -> Analyzer {
1543        Analyzer::new(Profile::Code.config()).expect("valid config")
1544    }
1545
1546    fn segment_of(analyzer: &Analyzer, texts: &[&str]) -> SegmentIndex {
1547        let mut segment = SegmentIndex::new();
1548        for text in texts {
1549            segment
1550                .push_document(analyzer, &Document::with_text(text))
1551                .expect("indexable");
1552        }
1553        segment
1554    }
1555
1556    #[test]
1557    fn rows_are_dense_and_ascending() {
1558        let analyzer = analyzer();
1559        let mut segment = SegmentIndex::new();
1560        for (expected, text) in ["alpha", "beta", "gamma"].into_iter().enumerate() {
1561            let row = segment
1562                .push_document(&analyzer, &Document::with_text(text))
1563                .expect("indexable");
1564            assert_eq!(row, u32::try_from(expected).expect("small"));
1565        }
1566        assert_eq!(segment.row_count(), 3);
1567    }
1568
1569    #[test]
1570    fn document_frequency_is_summed_across_segments_not_per_segment() {
1571        let analyzer = code_analyzer();
1572        let mut index = LexicalIndex::new();
1573        index
1574            .push_segment(segment_of(&analyzer, &["engine", "engine"]))
1575            .expect("seals");
1576        index
1577            .push_segment(segment_of(&analyzer, &["engine", "other"]))
1578            .expect("seals");
1579        assert_eq!(index.document_count(), 4);
1580        // Three of the four documents contain "engine", across two segments.
1581        assert_eq!(index.document_frequency(b"engine", &[DEFAULT_FIELD]), 3);
1582    }
1583
1584    #[test]
1585    fn a_term_in_two_fields_of_one_document_counts_once() {
1586        let analyzer = code_analyzer();
1587        let mut segment = SegmentIndex::new();
1588        let mut document = Document::new();
1589        document.set(FieldId(0), "engine");
1590        document.set(FieldId(1), "engine");
1591        segment
1592            .push_document(&analyzer, &document)
1593            .expect("indexable");
1594        let mut index = LexicalIndex::new();
1595        index.push_segment(segment).expect("seals");
1596        assert_eq!(
1597            index.document_frequency(b"engine", &[FieldId(0), FieldId(1)]),
1598            1
1599        );
1600    }
1601
1602    #[test]
1603    fn corpus_stats_span_every_segment() {
1604        let analyzer = code_analyzer();
1605        let mut index = LexicalIndex::new();
1606        index
1607            .push_segment(segment_of(&analyzer, &["a b c"]))
1608            .expect("seals");
1609        index
1610            .push_segment(segment_of(&analyzer, &["d e f g h"]))
1611            .expect("seals");
1612        let stats = index.corpus_stats().expect("two documents");
1613        assert_eq!(stats.document_count(), 2);
1614        assert_eq!(stats.total_tokens(), 8);
1615        assert!((stats.average_document_length() - 4.0).abs() < 1e-12);
1616    }
1617
1618    #[test]
1619    fn live_row_counters_and_document_frequency_exclude_tombstones() {
1620        let analyzer = code_analyzer();
1621        let sealed = SealedSegment::seal(&segment_of(
1622            &analyzer,
1623            &["engine padding padding", "engine"],
1624        ))
1625        .expect("seal lexical segment");
1626        let live_rows = DocBitmap::from_ids([1]);
1627        let mut index = LexicalIndex::new();
1628        index
1629            .push_sealed_with_live_rows(sealed, &live_rows)
1630            .expect("valid live rows");
1631
1632        let stats = index.corpus_stats().expect("one live document");
1633        assert_eq!(stats.document_count(), 1);
1634        assert_eq!(stats.total_tokens(), 1);
1635        assert_eq!(
1636            index.document_frequency(b"engine", &[DEFAULT_FIELD]),
1637            1,
1638            "document frequency must use the same live rows as corpus stats"
1639        );
1640    }
1641
1642    #[test]
1643    fn an_empty_index_refuses_to_produce_statistics() {
1644        let index = LexicalIndex::new();
1645        assert!(index.corpus_stats().is_err());
1646    }
1647
1648    #[test]
1649    fn document_length_counts_positions_not_tokens() {
1650        // Stacked variants share a position, so an identifier that emits
1651        // three tokens at one position must not count as three tokens of
1652        // length, or avgdl stops matching the unit being scored.
1653        let analyzer = code_analyzer();
1654        let mut segment = SegmentIndex::new();
1655        segment
1656            .push_document(&analyzer, &Document::with_text("put_if_match"))
1657            .expect("indexable");
1658        // Positions 0, 1, 2 for put / if / match.
1659        assert_eq!(segment.document_length(0), 3);
1660    }
1661
1662    #[test]
1663    fn field_lengths_are_tracked_separately_and_summed() {
1664        let analyzer = code_analyzer();
1665        let mut segment = SegmentIndex::new();
1666        let mut document = Document::new();
1667        document.set(FieldId(0), "one two");
1668        document.set(FieldId(1), "three four five");
1669        segment
1670            .push_document(&analyzer, &document)
1671            .expect("indexable");
1672        assert_eq!(segment.field_length(0, FieldId(0)), 2);
1673        assert_eq!(segment.field_length(0, FieldId(1)), 3);
1674        assert_eq!(segment.document_length(0), 5);
1675        assert_eq!(segment.total_tokens(), 5);
1676        assert_eq!(segment.field_length(0, FieldId(9)), 0);
1677    }
1678
1679    #[test]
1680    fn posting_lists_carry_positions() {
1681        let analyzer = code_analyzer();
1682        let segment = segment_of(&analyzer, &["alpha beta alpha"]);
1683        let list = segment
1684            .posting_list(b"alpha", DEFAULT_FIELD)
1685            .expect("alpha present");
1686        let posting = list.postings().first().expect("one posting");
1687        assert_eq!(posting.tf, 2);
1688        assert_eq!(posting.positions, vec![0, 2]);
1689    }
1690
1691    #[test]
1692    fn an_absent_term_has_no_posting_list_and_zero_frequency() {
1693        let analyzer = code_analyzer();
1694        let segment = segment_of(&analyzer, &["alpha"]);
1695        assert!(segment.posting_list(b"omega", DEFAULT_FIELD).is_none());
1696        let mut index = LexicalIndex::new();
1697        index.push_segment(segment).expect("seals");
1698        assert_eq!(index.document_frequency(b"omega", &[DEFAULT_FIELD]), 0);
1699    }
1700
1701    #[test]
1702    fn errors_render_useful_text() {
1703        assert!(
1704            IndexError::RowsNotAscending { row: 3 }
1705                .to_string()
1706                .contains("row order")
1707        );
1708        assert!(
1709            IndexError::Stats(Bm25Error::EmptyCorpus)
1710                .to_string()
1711                .contains("statistics")
1712        );
1713        assert!(
1714            IndexError::Postings(PostingsError::ZeroTermFrequency)
1715                .to_string()
1716                .contains("postings")
1717        );
1718    }
1719}