Skip to main content

rank_bm25/
lib.rs

1//! BM25Okapi, BM25L and BM25Plus with scores identical to Python's
2//! [rank_bm25](https://github.com/dorianbrown/rank_bm25) 0.2.2, down to the last bit.
3//!
4//! rank_bm25 computes each term's contribution for every document with NumPy. This crate keeps
5//! an inverted index, visits only the documents that contain a term where the other documents'
6//! contribution is exactly zero, and performs the remaining floating-point operations in the
7//! same order as NumPy, so every score equals rank_bm25's.
8//!
9//! ```
10//! use rank_bm25::{Builder, Params};
11//!
12//! let mut builder = Builder::new(Params::okapi());
13//! builder.add_document(["hello", "there", "good", "man"])?;
14//! builder.add_document(["it", "is", "quite", "windy", "in", "london"])?;
15//! builder.add_document(["how", "is", "the", "weather", "today"])?;
16//! let index = builder.build()?;
17//! let scores = index.get_scores(&["windy", "london"]);
18//! assert!(scores[1] > 0.0 && scores[0] == 0.0);
19//! # Ok::<(), rank_bm25::Error>(())
20//! ```
21
22use std::fmt;
23
24use rustc_hash::FxHashMap;
25
26/// The BM25 variant and its own parameter.
27#[derive(Clone, Copy, Debug, PartialEq)]
28pub enum Variant {
29    /// `BM25Okapi`: idf floored at `epsilon * average idf` for terms in more than half of the
30    /// documents.
31    Okapi { epsilon: f64 },
32    /// `BM25L`.
33    L { delta: f64 },
34    /// `BM25Plus`.
35    Plus { delta: f64 },
36}
37
38/// `k1`, `b` and the variant, with rank_bm25's defaults in the constructors.
39#[derive(Clone, Copy, Debug, PartialEq)]
40pub struct Params {
41    pub k1: f64,
42    pub b: f64,
43    pub variant: Variant,
44}
45
46impl Params {
47    pub fn okapi() -> Self {
48        Params {
49            k1: 1.5,
50            b: 0.75,
51            variant: Variant::Okapi { epsilon: 0.25 },
52        }
53    }
54
55    pub fn l() -> Self {
56        Params {
57            k1: 1.5,
58            b: 0.75,
59            variant: Variant::L { delta: 0.5 },
60        }
61    }
62
63    pub fn plus() -> Self {
64        Params {
65            k1: 1.5,
66            b: 0.75,
67            variant: Variant::Plus { delta: 1.0 },
68        }
69    }
70}
71
72/// rank_bm25 divides by the number of documents, and BM25Okapi by the number of distinct terms.
73#[derive(Clone, Copy, Debug, PartialEq, Eq)]
74pub enum Error {
75    /// An empty corpus, or for BM25Okapi a corpus without any term (Python's `ZeroDivisionError`).
76    ZeroDivision,
77    /// A serialised index that cannot be read.
78    InvalidData,
79    /// More than `u32::MAX` documents, distinct terms or occurrences of a term in a document.
80    TooLarge,
81}
82
83impl fmt::Display for Error {
84    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
85        match self {
86            Error::ZeroDivision => f.write_str("division by zero"),
87            Error::InvalidData => f.write_str("invalid serialised index"),
88            Error::TooLarge => f.write_str("corpus too large for 32-bit identifiers"),
89        }
90    }
91}
92
93impl std::error::Error for Error {}
94
95/// Documents and frequencies of every term, stored contiguously term after term.
96#[derive(Clone, Debug, Default)]
97struct Postings {
98    off: Vec<usize>,
99    docs: Vec<u32>,
100    tfs: Vec<u32>,
101}
102
103impl Postings {
104    #[inline]
105    fn get(&self, term: u32) -> (&[u32], &[u32]) {
106        let (a, b) = (self.off[term as usize], self.off[term as usize + 1]);
107        (&self.docs[a..b], &self.tfs[a..b])
108    }
109
110    fn terms(&self) -> usize {
111        self.off.len() - 1
112    }
113
114    fn df(&self, term: usize) -> usize {
115        self.off[term + 1] - self.off[term]
116    }
117}
118
119/// Collects documents, then builds an [`Index`].
120#[derive(Debug)]
121pub struct Builder {
122    params: Params,
123    vocab: FxHashMap<Box<str>, u32>,
124    terms: Vec<Box<str>>,
125    df: Vec<u32>,
126    doc_len: Vec<u64>,
127    forward: Vec<u32>,
128    forward_tf: Vec<u32>,
129    forward_off: Vec<usize>,
130    counts: Vec<u32>,
131    touched: Vec<u32>,
132    current_len: u64,
133    /// Vocabulary size when the current document started.
134    terms_at_start: usize,
135}
136
137impl Builder {
138    pub fn new(params: Params) -> Self {
139        Builder {
140            params,
141            vocab: FxHashMap::default(),
142            terms: Vec::new(),
143            df: Vec::new(),
144            doc_len: Vec::new(),
145            forward: Vec::new(),
146            forward_tf: Vec::new(),
147            forward_off: vec![0],
148            counts: Vec::new(),
149            touched: Vec::new(),
150            current_len: 0,
151            terms_at_start: 0,
152        }
153    }
154
155    /// Adds one token to the current document.
156    #[inline]
157    pub fn push_token(&mut self, token: &str) -> Result<(), Error> {
158        let id = match self.vocab.get(token) {
159            Some(&id) => id,
160            None => {
161                let id = u32::try_from(self.terms.len()).map_err(|_| Error::TooLarge)?;
162                self.vocab.insert(token.into(), id);
163                self.terms.push(token.into());
164                self.df.push(0);
165                self.counts.push(0);
166                id
167            }
168        };
169        let count = &mut self.counts[id as usize];
170        if *count == 0 {
171            self.touched.push(id);
172        }
173        *count = count.checked_add(1).ok_or(Error::TooLarge)?;
174        self.current_len += 1;
175        Ok(())
176    }
177
178    /// Ends the current document; the next token starts a new one.
179    pub fn end_document(&mut self) -> Result<(), Error> {
180        if self.doc_len.len() >= u32::MAX as usize {
181            return Err(Error::TooLarge);
182        }
183        for &t in &self.touched {
184            let count = std::mem::take(&mut self.counts[t as usize]);
185            self.df[t as usize] += 1;
186            self.forward.push(t);
187            self.forward_tf.push(count);
188        }
189        self.touched.clear();
190        self.doc_len.push(self.current_len);
191        self.current_len = 0;
192        self.forward_off.push(self.forward.len());
193        self.terms_at_start = self.terms.len();
194        Ok(())
195    }
196
197    /// Discards the tokens of the current, unfinished document, including the terms it
198    /// introduced.
199    pub fn abandon_document(&mut self) {
200        for &t in &self.touched {
201            self.counts[t as usize] = 0;
202        }
203        self.touched.clear();
204        self.current_len = 0;
205        for term in self.terms.drain(self.terms_at_start..) {
206            self.vocab.remove(&term);
207        }
208        self.df.truncate(self.terms_at_start);
209        self.counts.truncate(self.terms_at_start);
210    }
211
212    pub fn add_document<'a, I: IntoIterator<Item = &'a str>>(
213        &mut self,
214        tokens: I,
215    ) -> Result<(), Error> {
216        for token in tokens {
217            if let Err(e) = self.push_token(token) {
218                self.abandon_document();
219                return Err(e);
220            }
221        }
222        self.end_document()
223    }
224
225    pub fn build(mut self) -> Result<Index, Error> {
226        self.abandon_document();
227        let mut off = Vec::with_capacity(self.df.len() + 1);
228        off.push(0usize);
229        for &df in &self.df {
230            off.push(off.last().unwrap() + df as usize);
231        }
232        let total = *off.last().unwrap();
233        let mut docs = vec![0u32; total];
234        let mut tfs = vec![0u32; total];
235        let mut cursor: Vec<usize> = off[..off.len() - 1].to_vec();
236        for d in 0..self.doc_len.len() {
237            let (a, b) = (self.forward_off[d], self.forward_off[d + 1]);
238            for (&t, &tf) in self.forward[a..b].iter().zip(&self.forward_tf[a..b]) {
239                let c = &mut cursor[t as usize];
240                docs[*c] = d as u32;
241                tfs[*c] = tf;
242                *c += 1;
243            }
244        }
245        drop(self.forward_tf);
246        let postings = Postings { off, docs, tfs };
247        Index::new(
248            self.params,
249            self.vocab,
250            self.terms,
251            postings,
252            self.doc_len,
253            self.forward,
254            self.forward_off,
255        )
256    }
257}
258
259/// A BM25 index over a fixed corpus.
260#[derive(Clone, Debug)]
261pub struct Index {
262    params: Params,
263    avgdl: f64,
264    vocab: FxHashMap<Box<str>, u32>,
265    terms: Vec<Box<str>>,
266    idf: Vec<f64>,
267    average_idf: Option<f64>,
268    postings: Postings,
269    doc_len: Vec<u64>,
270    forward: Vec<u32>,
271    forward_off: Vec<usize>,
272    /// Okapi and Plus: `k1 * ((1 - b) + (b * dl) / avgdl)`; L: `(1 - b) + (b * dl) / avgdl`.
273    norm: Vec<f64>,
274    kp1: f64,
275    /// Documents without a term contribute exactly zero, so only its postings need visiting.
276    sparse: bool,
277    /// Every contribution is finite whatever the query, so scores need no NaN check.
278    always_finite: bool,
279    max_tf: u32,
280}
281
282impl Index {
283    fn new(
284        params: Params,
285        vocab: FxHashMap<Box<str>, u32>,
286        terms: Vec<Box<str>>,
287        postings: Postings,
288        doc_len: Vec<u64>,
289        forward: Vec<u32>,
290        forward_off: Vec<usize>,
291    ) -> Result<Self, Error> {
292        let n = doc_len.len();
293        if n == 0 {
294            return Err(Error::ZeroDivision);
295        }
296        let num_doc = doc_len
297            .iter()
298            .try_fold(0u64, |a, &l| a.checked_add(l))
299            .ok_or(Error::TooLarge)?;
300        let avgdl = num_doc as f64 / n as f64;
301        let nterms = postings.terms();
302        let mut idf = Vec::with_capacity(nterms);
303        let mut average_idf = None;
304        match params.variant {
305            Variant::Okapi { epsilon } => {
306                if nterms == 0 {
307                    return Err(Error::ZeroDivision);
308                }
309                let mut idf_sum = 0.0f64;
310                let mut negative = Vec::new();
311                for t in 0..nterms {
312                    let freq = postings.df(t) as u64;
313                    let v = ((n as u64 - freq) as f64 + 0.5).ln() - (freq as f64 + 0.5).ln();
314                    idf.push(v);
315                    idf_sum += v;
316                    if v < 0.0 {
317                        negative.push(t);
318                    }
319                }
320                let average = idf_sum / nterms as f64;
321                let eps = epsilon * average;
322                for t in negative {
323                    idf[t] = eps;
324                }
325                average_idf = Some(average);
326            }
327            Variant::L { .. } => {
328                for t in 0..nterms {
329                    idf.push(((n + 1) as f64).ln() - (postings.df(t) as f64 + 0.5).ln());
330                }
331            }
332            Variant::Plus { .. } => {
333                for t in 0..nterms {
334                    idf.push(((n + 1) as f64 / postings.df(t) as f64).ln());
335                }
336            }
337        }
338        let Params { k1, b, variant } = params;
339        let norm: Vec<f64> = doc_len
340            .iter()
341            .map(|&dl| {
342                let l = (1.0 - b) + (b * dl as f64) / avgdl;
343                match variant {
344                    Variant::L { .. } => l,
345                    _ => k1 * l,
346                }
347            })
348            .collect();
349        let kp1 = k1 + 1.0;
350        let norm_ok = norm.iter().all(|&x| x != 0.0 && !x.is_nan());
351        let sparse = norm_ok
352            && kp1.is_finite()
353            && match variant {
354                Variant::Okapi { .. } => true,
355                Variant::L { delta } => {
356                    let denominator = (k1 + 0.0) + delta;
357                    delta.is_finite() && denominator != 0.0 && !denominator.is_nan()
358                }
359                Variant::Plus { .. } => false,
360            };
361        let moderate = |x: f64| x.is_finite() && x.abs() <= 1e100;
362        let extra = match variant {
363            Variant::Okapi { epsilon } => epsilon,
364            Variant::L { delta } | Variant::Plus { delta } => delta,
365        };
366        let always_finite = moderate(k1)
367            && moderate(b)
368            && moderate(extra)
369            && norm.iter().all(|&x| x > 0.0 && x.is_finite())
370            && idf.iter().all(|x| x.is_finite())
371            && match variant {
372                Variant::L { delta } => k1 >= 0.0 && delta >= 0.0 && k1 + delta > 0.0,
373                _ => true,
374            };
375        Ok(Index {
376            always_finite,
377            params,
378            avgdl,
379            vocab,
380            terms,
381            idf,
382            average_idf,
383            max_tf: postings.tfs.iter().copied().max().unwrap_or(0),
384            postings,
385            doc_len,
386            forward,
387            forward_off,
388            norm,
389            kp1,
390            sparse,
391        })
392    }
393
394    pub fn params(&self) -> Params {
395        self.params
396    }
397
398    /// Number of documents.
399    pub fn corpus_size(&self) -> usize {
400        self.doc_len.len()
401    }
402
403    pub fn avgdl(&self) -> f64 {
404        self.avgdl
405    }
406
407    /// The average idf before flooring, for BM25Okapi.
408    pub fn average_idf(&self) -> Option<f64> {
409        self.average_idf
410    }
411
412    /// True when no query can produce a NaN or infinite score.
413    pub fn always_finite(&self) -> bool {
414        self.always_finite
415    }
416
417    /// The largest frequency of a term in a document.
418    pub fn max_tf(&self) -> u32 {
419        self.max_tf
420    }
421
422    pub fn doc_len(&self) -> &[u64] {
423        &self.doc_len
424    }
425
426    /// Terms with their idf, in the order rank_bm25 stores them (first occurrence in the corpus).
427    pub fn idf(&self) -> impl Iterator<Item = (&str, f64)> {
428        self.terms
429            .iter()
430            .map(|t| &**t)
431            .zip(self.idf.iter().copied())
432    }
433
434    /// The idf of a term returned by [`Index::term_id`].
435    pub fn term_idf(&self, term: u32) -> f64 {
436        self.idf[term as usize]
437    }
438
439    pub fn term_id(&self, term: &str) -> Option<u32> {
440        self.vocab.get(term).copied()
441    }
442
443    /// The terms of document `doc` with their frequencies, in order of first occurrence.
444    pub fn document(&self, doc: usize) -> impl Iterator<Item = (&str, u32)> + '_ {
445        self.forward[self.forward_off[doc]..self.forward_off[doc + 1]]
446            .iter()
447            .map(move |&t| (&*self.terms[t as usize], self.tf(t, doc as u32)))
448    }
449
450    /// Frequency of term `term` in document `doc`.
451    pub fn tf(&self, term: u32, doc: u32) -> u32 {
452        let (docs, tfs) = self.postings.get(term);
453        match docs.binary_search(&doc) {
454            Ok(i) => tfs[i],
455            Err(_) => 0,
456        }
457    }
458
459    #[inline]
460    fn contribution(&self, idf: f64, f: f64, d: usize) -> f64 {
461        let Params { k1, variant, .. } = self.params;
462        match variant {
463            Variant::Okapi { .. } => idf * ((f * self.kp1) / (f + self.norm[d])),
464            Variant::L { delta } => {
465                let ctd = f / self.norm[d];
466                (((idf * f) * self.kp1) * (ctd + delta)) / ((k1 + ctd) + delta)
467            }
468            Variant::Plus { delta } => idf * (delta + ((f * self.kp1) / (self.norm[d] + f))),
469        }
470    }
471
472    /// Adds the contributions of `term` (None for a term not in the corpus) to `score`.
473    pub fn add_term(&self, term: Option<u32>, score: &mut [f64]) {
474        let (idf, docs, tfs): (f64, &[u32], &[u32]) = match term {
475            Some(t) => {
476                let (docs, tfs) = self.postings.get(t);
477                (self.idf[t as usize], docs, tfs)
478            }
479            None => (0.0, &[], &[]),
480        };
481        if self.sparse && idf.is_finite() {
482            for (&d, &tf) in docs.iter().zip(tfs) {
483                score[d as usize] += self.contribution(idf, tf as f64, d as usize);
484            }
485            return;
486        }
487        let mut next = docs.iter().zip(tfs).peekable();
488        for (d, s) in score.iter_mut().enumerate() {
489            let f = match next.peek() {
490                Some(&(&pd, &tf)) if pd as usize == d => {
491                    next.next();
492                    tf as f64
493                }
494                _ => 0.0,
495            };
496            *s += self.contribution(idf, f, d);
497        }
498    }
499
500    /// `get_scores` for query terms already looked up with [`Index::term_id`].
501    pub fn scores_into(&self, query: &[Option<u32>], score: &mut [f64]) {
502        score.fill(0.0);
503        for &term in query {
504            self.add_term(term, score);
505        }
506    }
507
508    /// The score of every document, as `get_scores` in rank_bm25.
509    pub fn get_scores<S: AsRef<str>>(&self, query: &[S]) -> Vec<f64> {
510        let ids: Vec<Option<u32>> = query.iter().map(|q| self.term_id(q.as_ref())).collect();
511        let mut score = vec![0.0; self.corpus_size()];
512        self.scores_into(&ids, &mut score);
513        score
514    }
515
516    /// Serialises the index.
517    pub fn to_bytes(&self) -> Vec<u8> {
518        let mut w = Writer(Vec::new());
519        w.0.extend_from_slice(b"BM25R1");
520        let (kind, extra) = match self.params.variant {
521            Variant::Okapi { epsilon } => (0u8, epsilon),
522            Variant::L { delta } => (1, delta),
523            Variant::Plus { delta } => (2, delta),
524        };
525        w.0.push(kind);
526        w.f64(self.params.k1);
527        w.f64(self.params.b);
528        w.f64(extra);
529        w.u64(self.doc_len.len() as u64);
530        for &dl in &self.doc_len {
531            w.u64(dl);
532        }
533        w.u64(self.terms.len() as u64);
534        for (t, term) in self.terms.iter().enumerate() {
535            w.u64(term.len() as u64);
536            w.0.extend_from_slice(term.as_bytes());
537            let (docs, tfs) = self.postings.get(t as u32);
538            w.u64(docs.len() as u64);
539            for (&d, &tf) in docs.iter().zip(tfs) {
540                w.u32(d);
541                w.u32(tf);
542            }
543        }
544        w.u64(self.forward.len() as u64);
545        for &t in &self.forward {
546            w.u32(t);
547        }
548        for &o in &self.forward_off {
549            w.u64(o as u64);
550        }
551        w.0
552    }
553
554    /// Restores an index saved with [`Index::to_bytes`].
555    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
556        let mut r = Reader { buf: bytes, pos: 0 };
557        if r.take(6)? != b"BM25R1" {
558            return Err(Error::InvalidData);
559        }
560        let kind = r.take(1)?[0];
561        let k1 = r.f64()?;
562        let b = r.f64()?;
563        let extra = r.f64()?;
564        let variant = match kind {
565            0 => Variant::Okapi { epsilon: extra },
566            1 => Variant::L { delta: extra },
567            2 => Variant::Plus { delta: extra },
568            _ => return Err(Error::InvalidData),
569        };
570        let n = r.len()?;
571        let mut doc_len = Vec::with_capacity(n);
572        for _ in 0..n {
573            doc_len.push(r.u64()?);
574        }
575        let nterms = r.len()?;
576        let mut vocab = FxHashMap::default();
577        let mut terms = Vec::with_capacity(nterms);
578        let mut postings = Postings {
579            off: vec![0],
580            ..Default::default()
581        };
582        for t in 0..nterms {
583            let len = r.len()?;
584            let term: Box<str> = std::str::from_utf8(r.take(len)?)
585                .map_err(|_| Error::InvalidData)?
586                .into();
587            if vocab.insert(term.clone(), t as u32).is_some() {
588                return Err(Error::InvalidData);
589            }
590            terms.push(term);
591            let plen = r.len()?;
592            let mut last = None;
593            for _ in 0..plen {
594                let d = r.u32()?;
595                if d as usize >= n || last.is_some_and(|l| d <= l) {
596                    return Err(Error::InvalidData);
597                }
598                last = Some(d);
599                postings.docs.push(d);
600                postings.tfs.push(r.u32()?);
601            }
602            postings.off.push(postings.docs.len());
603        }
604        let nf = r.len()?;
605        let mut forward = Vec::with_capacity(nf);
606        for _ in 0..nf {
607            let t = r.u32()?;
608            if t as usize >= nterms {
609                return Err(Error::InvalidData);
610            }
611            forward.push(t);
612        }
613        let mut forward_off = Vec::with_capacity(n + 1);
614        for _ in 0..=n {
615            let o = r.u64()? as usize;
616            if o > nf || forward_off.last().is_some_and(|&p| o < p) {
617                return Err(Error::InvalidData);
618            }
619            forward_off.push(o);
620        }
621        if r.pos != bytes.len()
622            || forward_off.first() != Some(&0)
623            || forward_off.last() != Some(&forward.len())
624            || postings.docs.len() != forward.len()
625            || postings.tfs.contains(&0)
626            || (0..nterms).any(|t| postings.df(t) == 0)
627        {
628            return Err(Error::InvalidData);
629        }
630        // Every document's terms must match the postings and add up to its length.
631        let mut seen = vec![u32::MAX; nterms];
632        for d in 0..n {
633            let mut len = 0u64;
634            for &t in &forward[forward_off[d]..forward_off[d + 1]] {
635                if seen[t as usize] == d as u32 {
636                    return Err(Error::InvalidData);
637                }
638                seen[t as usize] = d as u32;
639                let (docs, tfs) = postings.get(t);
640                let i = docs
641                    .binary_search(&(d as u32))
642                    .map_err(|_| Error::InvalidData)?;
643                len += tfs[i] as u64;
644            }
645            if len != doc_len[d] {
646                return Err(Error::InvalidData);
647            }
648        }
649        Index::new(
650            Params { k1, b, variant },
651            vocab,
652            terms,
653            postings,
654            doc_len,
655            forward,
656            forward_off,
657        )
658    }
659}
660
661struct Writer(Vec<u8>);
662
663impl Writer {
664    fn u32(&mut self, v: u32) {
665        self.0.extend_from_slice(&v.to_le_bytes());
666    }
667    fn u64(&mut self, v: u64) {
668        self.0.extend_from_slice(&v.to_le_bytes());
669    }
670    fn f64(&mut self, v: f64) {
671        self.u64(v.to_bits());
672    }
673}
674
675struct Reader<'a> {
676    buf: &'a [u8],
677    pos: usize,
678}
679
680impl<'a> Reader<'a> {
681    fn take(&mut self, n: usize) -> Result<&'a [u8], Error> {
682        let end = self.pos.checked_add(n).ok_or(Error::InvalidData)?;
683        let s = self.buf.get(self.pos..end).ok_or(Error::InvalidData)?;
684        self.pos = end;
685        Ok(s)
686    }
687    fn u32(&mut self) -> Result<u32, Error> {
688        Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
689    }
690    fn u64(&mut self) -> Result<u64, Error> {
691        Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
692    }
693    fn f64(&mut self) -> Result<f64, Error> {
694        Ok(f64::from_bits(self.u64()?))
695    }
696    fn len(&mut self) -> Result<usize, Error> {
697        let v = self.u64()? as usize;
698        if v > self.buf.len() {
699            return Err(Error::InvalidData);
700        }
701        Ok(v)
702    }
703}
704
705#[cfg(test)]
706mod tests {
707    use super::*;
708
709    fn index(params: Params) -> Index {
710        let mut b = Builder::new(params);
711        b.add_document(["hello", "there", "good", "man"]).unwrap();
712        b.add_document(["it", "is", "quite", "windy", "in", "london"])
713            .unwrap();
714        b.add_document(["how", "is", "the", "weather", "today"])
715            .unwrap();
716        b.build().unwrap()
717    }
718
719    // Values from rank_bm25 0.2.2.
720    #[test]
721    fn matches_rank_bm25() {
722        let q = ["windy", "london"];
723        assert_eq!(
724            index(Params::okapi()).get_scores(&q),
725            [0.0, 0.9372947225064051, 0.0]
726        );
727        assert_eq!(
728            index(Params::l()).get_scores(&q),
729            [0.0, 2.340615262868892, 0.0]
730        );
731        assert_eq!(
732            index(Params::plus()).get_scores(&q),
733            [2.772588722239781, 5.316248100441416, 2.772588722239781]
734        );
735    }
736
737    #[test]
738    fn document_keeps_first_occurrence_order() {
739        let mut b = Builder::new(Params::okapi());
740        b.add_document(["b", "a", "b", "c"]).unwrap();
741        let idx = b.build().unwrap();
742        assert_eq!(
743            idx.document(0).collect::<Vec<_>>(),
744            [("b", 2), ("a", 1), ("c", 1)]
745        );
746    }
747
748    #[test]
749    fn empty_corpus_is_an_error() {
750        assert_eq!(
751            Builder::new(Params::okapi()).build().unwrap_err(),
752            Error::ZeroDivision
753        );
754        let mut b = Builder::new(Params::okapi());
755        b.add_document(Vec::<&str>::new()).unwrap();
756        assert_eq!(b.build().unwrap_err(), Error::ZeroDivision);
757        let mut b = Builder::new(Params::l());
758        b.add_document(Vec::<&str>::new()).unwrap();
759        assert!(b.build().is_ok());
760    }
761
762    #[test]
763    fn abandoned_terms_leave_no_trace() {
764        let mut b = Builder::new(Params::okapi());
765        b.add_document(["a"]).unwrap();
766        b.push_token("discarded").unwrap();
767        b.abandon_document();
768        let with_abandon = b.build().unwrap();
769        let mut c = Builder::new(Params::okapi());
770        c.add_document(["a"]).unwrap();
771        let plain = c.build().unwrap();
772        assert_eq!(with_abandon.idf().count(), 1);
773        assert_eq!(with_abandon.get_scores(&["a"]), plain.get_scores(&["a"]));
774    }
775
776    #[test]
777    fn inconsistent_bytes_are_rejected() {
778        let idx = index(Params::l());
779        let bytes = idx.to_bytes();
780        let mut zeroed = bytes.clone();
781        let n = zeroed.len();
782        zeroed[n - 16..].fill(0);
783        assert!(Index::from_bytes(&zeroed).is_err());
784        for i in 0..bytes.len() {
785            let mut b = bytes.clone();
786            b[i] ^= 0xff;
787            let _ = Index::from_bytes(&b);
788        }
789    }
790
791    #[test]
792    fn terms_without_postings_are_rejected() {
793        let mut w = Writer(Vec::new());
794        w.0.extend_from_slice(b"BM25R1");
795        w.0.push(0);
796        for v in [1.5, 0.75, 0.25] {
797            w.f64(v);
798        }
799        w.u64(1);
800        w.u64(1);
801        w.u64(2);
802        for (term, postings) in [("a", 1u64), ("ghost", 0)] {
803            w.u64(term.len() as u64);
804            w.0.extend_from_slice(term.as_bytes());
805            w.u64(postings);
806            for _ in 0..postings {
807                w.u32(0);
808                w.u32(1);
809            }
810        }
811        w.u64(1);
812        w.u32(0);
813        w.u64(0);
814        w.u64(1);
815        assert_eq!(Index::from_bytes(&w.0).unwrap_err(), Error::InvalidData);
816    }
817
818    #[test]
819    fn index_is_send_and_sync() {
820        fn check<T: Send + Sync>() {}
821        check::<Index>();
822    }
823
824    #[test]
825    fn serialisation_round_trips() {
826        for params in [Params::okapi(), Params::l(), Params::plus()] {
827            let idx = index(params);
828            let back = Index::from_bytes(&idx.to_bytes()).unwrap();
829            let q = ["is", "windy", "nothing"];
830            assert_eq!(back.get_scores(&q), idx.get_scores(&q));
831        }
832        assert!(Index::from_bytes(b"BM25R1").is_err());
833    }
834}