#![forbid(unsafe_code)]
pub mod codec;
#[cfg(feature = "positional")]
pub mod positional;
use std::borrow::Borrow;
use std::collections::{HashMap, HashSet};
use std::hash::Hash;
pub type DocId = u32;
pub trait Weight: Copy + Default + std::fmt::Debug + 'static {
fn zero() -> Self;
fn accumulate(&mut self, other: Self);
fn to_f32(self) -> f32;
fn to_doc_len(self) -> u64;
}
impl Weight for u32 {
#[inline]
fn zero() -> Self {
0
}
#[inline]
fn accumulate(&mut self, other: Self) {
*self += other;
}
#[inline]
fn to_f32(self) -> f32 {
self as f32
}
#[inline]
fn to_doc_len(self) -> u64 {
self as u64
}
}
impl Weight for f32 {
#[inline]
fn zero() -> Self {
0.0
}
#[inline]
fn accumulate(&mut self, other: Self) {
*self += other;
}
#[inline]
fn to_f32(self) -> f32 {
self
}
#[inline]
fn to_doc_len(self) -> u64 {
1
}
}
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error("document already exists: {0}")]
DuplicateDocId(DocId),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CandidatePlan {
Candidates(Vec<DocId>),
ScanAll,
}
#[derive(Debug, Clone, Copy)]
pub struct PlannerConfig {
pub max_candidate_ratio: f32,
pub max_candidates: u32,
}
impl Default for PlannerConfig {
fn default() -> Self {
Self {
max_candidate_ratio: 0.6,
max_candidates: 200_000,
}
}
}
#[derive(Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(
feature = "serde",
serde(bound(
serialize = "Term: serde::Serialize, W: serde::Serialize",
deserialize = "Term: serde::Deserialize<'de> + Eq + std::hash::Hash, W: serde::Deserialize<'de>"
))
)]
struct Segment<Term, W: Weight = u32> {
doc_terms: HashMap<DocId, Vec<Term>>,
#[cfg_attr(feature = "serde", serde(skip))]
_w: std::marker::PhantomData<W>,
}
impl<Term, W: Weight> Default for Segment<Term, W> {
fn default() -> Self {
Self {
doc_terms: HashMap::new(),
_w: std::marker::PhantomData,
}
}
}
#[derive(Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(
feature = "serde",
serde(bound(
serialize = "Term: serde::Serialize, W: serde::Serialize",
deserialize = "Term: serde::Deserialize<'de> + Eq + std::hash::Hash, W: serde::Deserialize<'de>"
))
)]
pub struct PostingsIndex<Term = String, W: Weight = u32> {
segments: Vec<Segment<Term, W>>,
doc_segment: HashMap<DocId, usize>,
doc_len: HashMap<DocId, u32>,
df: HashMap<Term, u32>,
total_doc_len: u64,
global_postings: HashMap<Term, Vec<(DocId, W)>>,
}
impl<Term, W: Weight> Default for PostingsIndex<Term, W> {
fn default() -> Self {
Self {
segments: Vec::new(),
doc_segment: HashMap::new(),
doc_len: HashMap::new(),
df: HashMap::new(),
total_doc_len: 0,
global_postings: HashMap::new(),
}
}
}
impl<Term, W: Weight> PostingsIndex<Term, W>
where
Term: Clone + Eq + Hash + Ord,
{
pub fn new() -> Self {
Self::default()
}
pub fn num_docs(&self) -> u32 {
self.doc_len.len() as u32
}
pub fn avg_doc_len(&self) -> f32 {
let n = self.num_docs() as f32;
if n == 0.0 {
return 0.0;
}
(self.total_doc_len as f32) / n
}
pub fn document_ids(&self) -> impl Iterator<Item = DocId> + '_ {
self.doc_len.keys().copied()
}
pub fn document_len(&self, doc_id: DocId) -> u32 {
self.doc_len.get(&doc_id).copied().unwrap_or(0)
}
pub fn df<Q>(&self, term: &Q) -> u32
where
Term: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.df.get(term).copied().unwrap_or(0)
}
pub fn terms(&self) -> impl Iterator<Item = &Term> + '_ {
self.df.keys()
}
pub fn add_weighted_document(
&mut self,
doc_id: DocId,
weighted_terms: &[(Term, W)],
) -> Result<(), Error> {
if self.doc_segment.contains_key(&doc_id) {
return Err(Error::DuplicateDocId(doc_id));
}
let mut term_weights: HashMap<Term, W> = HashMap::new();
let mut doc_length: u64 = 0;
for (t, w) in weighted_terms {
term_weights
.entry(t.clone())
.and_modify(|existing| existing.accumulate(*w))
.or_insert(*w);
doc_length += w.to_doc_len();
}
self.index_document_inner(doc_id, term_weights, doc_length);
Ok(())
}
fn index_document_inner(
&mut self,
doc_id: DocId,
term_weights: HashMap<Term, W>,
doc_length: u64,
) {
let mut doc_terms: Vec<Term> = term_weights.keys().cloned().collect();
doc_terms.sort_unstable();
for (term, w) in &term_weights {
self.global_postings
.entry(term.clone())
.or_default()
.push((doc_id, *w));
}
for term in &doc_terms {
*self.df.entry(term.clone()).or_insert(0) += 1;
}
let mut seg = Segment::<Term, W>::default();
seg.doc_terms.insert(doc_id, doc_terms);
let seg_idx = self.segments.len();
self.segments.push(seg);
self.doc_segment.insert(doc_id, seg_idx);
self.doc_len.insert(doc_id, doc_length as u32);
self.total_doc_len += doc_length;
}
}
impl<Term> PostingsIndex<Term, u32>
where
Term: Clone + Eq + Hash + Ord,
{
pub fn add_document(&mut self, doc_id: DocId, terms: &[Term]) -> Result<(), Error> {
if self.doc_segment.contains_key(&doc_id) {
return Err(Error::DuplicateDocId(doc_id));
}
let mut term_counts: HashMap<Term, u32> = HashMap::new();
for t in terms {
*term_counts.entry(t.clone()).or_insert(0) += 1;
}
let doc_length = terms.len() as u64;
self.index_document_inner(doc_id, term_counts, doc_length);
Ok(())
}
}
impl<Term, W: Weight> PostingsIndex<Term, W>
where
Term: Clone + Eq + Hash + Ord,
{
pub fn delete_document(&mut self, doc_id: DocId) -> bool {
let seg_idx = match self.doc_segment.remove(&doc_id) {
Some(i) => i,
None => return false,
};
let doc_len = self.doc_len.remove(&doc_id).unwrap_or(0);
self.total_doc_len = self.total_doc_len.saturating_sub(doc_len as u64);
let seg = &self.segments[seg_idx];
if let Some(terms) = seg.doc_terms.get(&doc_id) {
for term in terms {
if let Some(df) = self.df.get_mut(term) {
*df = df.saturating_sub(1);
if *df == 0 {
self.df.remove(term);
}
}
}
}
true
}
pub fn term_frequency<Q>(&self, doc_id: DocId, term: &Q) -> W
where
Term: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
if !self.doc_segment.contains_key(&doc_id) {
return W::zero();
}
let postings = match self.global_postings.get(term) {
Some(p) => p,
None => return W::zero(),
};
match postings.binary_search_by_key(&doc_id, |(id, _)| *id) {
Ok(i) => postings[i].1,
Err(_) => W::zero(),
}
}
pub fn candidates<Q>(&self, query_terms: &[Q]) -> Vec<DocId>
where
Term: Borrow<Q>,
Q: Hash + Eq,
{
if query_terms.is_empty() {
return Vec::new();
}
let mut out: Vec<DocId> = Vec::new();
let mut seen: HashSet<DocId> = HashSet::new();
for term in query_terms {
for (doc_id, _) in self.postings_iter(term) {
if seen.insert(doc_id) {
out.push(doc_id);
}
}
}
out.sort_unstable();
out
}
pub fn candidates_all_terms<Q>(&self, query_terms: &[Q]) -> Vec<DocId>
where
Term: Borrow<Q>,
Q: Hash + Eq,
{
if query_terms.is_empty() {
return Vec::new();
}
let mut uniq: Vec<&Q> = Vec::new();
let mut seen: HashSet<&Q> = HashSet::new();
for t in query_terms {
if seen.insert(t) {
uniq.push(t);
}
}
if uniq.is_empty() {
return Vec::new();
}
uniq.sort_by_key(|t| self.df(*t));
if self.df(uniq[0]) == 0 {
return Vec::new();
}
let mut acc: Vec<DocId> = {
let mut v: Vec<DocId> = self.postings_iter(uniq[0]).map(|(id, _)| id).collect();
v.sort_unstable();
v
};
let mut buf: Vec<DocId> = Vec::new();
for &t in uniq.iter().skip(1) {
if self.df(t) == 0 {
return Vec::new();
}
buf.clear();
match self.global_postings.get(t) {
Some(list) => buf.extend(list.iter().map(|(id, _)| *id)),
None => return Vec::new(),
}
buf.sort_unstable();
acc = intersect_sorted(&acc, &buf);
if acc.is_empty() {
break;
}
}
acc
}
pub fn plan_candidates<Q>(&self, query_terms: &[Q], cfg: PlannerConfig) -> CandidatePlan
where
Term: Borrow<Q>,
Q: Hash + Eq,
{
if query_terms.is_empty() {
return CandidatePlan::Candidates(Vec::new());
}
let n = self.num_docs();
if n == 0 {
return CandidatePlan::Candidates(Vec::new());
}
let mut seen_terms: HashSet<&Q> = HashSet::new();
let mut df_sum: u64 = 0;
for t in query_terms {
if !seen_terms.insert(t) {
continue;
}
df_sum = df_sum.saturating_add(self.df(t) as u64);
if df_sum >= cfg.max_candidates as u64 {
return CandidatePlan::ScanAll;
}
}
let ratio = (df_sum as f32) / (n as f32);
if ratio > cfg.max_candidate_ratio {
return CandidatePlan::ScanAll;
}
CandidatePlan::Candidates(self.candidates(query_terms))
}
pub fn postings_iter<'a, Q>(&'a self, term: &'a Q) -> impl Iterator<Item = (DocId, W)> + 'a
where
Term: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.global_postings
.get(term)
.into_iter()
.flat_map(|v| v.iter().copied())
.filter(|(doc_id, _)| self.doc_segment.contains_key(doc_id))
}
#[cfg(feature = "persistence")]
pub fn save<D: durability::Directory + ?Sized>(
&self,
dir: &D,
path: &str,
) -> Result<(), Box<dyn std::error::Error>>
where
Term: serde::Serialize,
W: serde::Serialize,
{
let bytes = postcard::to_allocvec(self)?;
dir.atomic_write(path, &bytes)?;
Ok(())
}
#[cfg(feature = "persistence")]
pub fn save_durable<D: durability::Directory + ?Sized>(
&self,
dir: &D,
path: &str,
) -> Result<(), Box<dyn std::error::Error>>
where
Term: serde::Serialize,
W: serde::Serialize,
{
let bytes = postcard::to_allocvec(self)?;
dir.atomic_write_durable(path, &bytes)?;
Ok(())
}
#[cfg(feature = "persistence")]
pub fn load<D: durability::Directory + ?Sized>(
dir: &D,
path: &str,
) -> Result<Self, Box<dyn std::error::Error>>
where
for<'de> Term: serde::Deserialize<'de>,
for<'de> W: serde::Deserialize<'de>,
{
use std::io::Read;
let mut f = dir.open_file(path)?;
let mut bytes = Vec::new();
f.read_to_end(&mut bytes)?;
let idx: Self = postcard::from_bytes(&bytes)?;
Ok(idx)
}
}
fn intersect_sorted(a: &[DocId], b: &[DocId]) -> Vec<DocId> {
let mut out = Vec::new();
let mut i = 0usize;
let mut j = 0usize;
while i < a.len() && j < b.len() {
match a[i].cmp(&b[j]) {
std::cmp::Ordering::Equal => {
out.push(a[i]);
i += 1;
j += 1;
}
std::cmp::Ordering::Less => {
i = gallop_forward(a, i, b[j]);
}
std::cmp::Ordering::Greater => {
j = gallop_forward(b, j, a[i]);
}
}
}
out
}
#[inline]
fn gallop_forward(list: &[DocId], start: usize, target: DocId) -> usize {
if start >= list.len() || list[start] >= target {
return start;
}
let mut step = 1usize;
while start + step < list.len() && list[start + step] < target {
step <<= 1;
}
let lo = start + (step >> 1);
let hi = (start + step).min(list.len());
match list[lo..hi].binary_search(&target) {
Ok(k) => lo + k,
Err(k) => lo + k,
}
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::prelude::*;
#[test]
fn add_and_lookup_basic() {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
idx.add_document(
0,
&[
String::from("the"),
String::from("quick"),
String::from("quick"),
],
)
.unwrap();
assert_eq!(idx.num_docs(), 1);
assert_eq!(idx.document_len(0), 3);
assert_eq!(idx.df("quick"), 1);
assert_eq!(idx.term_frequency(0, "quick"), 2);
assert_eq!(idx.term_frequency(0, "missing"), 0);
}
#[test]
fn delete_updates_df() {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
idx.add_document(0, &[String::from("a"), String::from("b")])
.unwrap();
idx.add_document(1, &[String::from("b"), String::from("c")])
.unwrap();
assert_eq!(idx.df("b"), 2);
assert!(idx.delete_document(0));
assert_eq!(idx.df("b"), 1);
assert_eq!(idx.df("a"), 0);
assert_eq!(idx.term_frequency(0, "b"), 0);
assert_eq!(idx.term_frequency(1, "b"), 1);
}
#[test]
fn multilingual_terms_do_not_panic() {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
idx.add_document(
0,
&[
String::from("Müller"), String::from("東京"), String::from("مرحبا"), String::from("Москва"), String::from("cafe\u{0301}"), String::from("👨\u{200D}👩\u{200D}👧\u{200D}👦"), ],
)
.unwrap();
assert_eq!(idx.num_docs(), 1);
assert_eq!(idx.df("東京"), 1);
assert_eq!(idx.term_frequency(0, "مرحبا"), 1);
}
proptest::proptest! {
#[test]
fn df_is_never_negative_and_upper_bounded(
docs in proptest::collection::vec(
proptest::collection::vec("[a-z]{1,6}", 0..20),
0..50
)
) {
use proptest::prelude::*;
let mut idx: PostingsIndex<String> = PostingsIndex::new();
for (i, doc) in docs.iter().enumerate() {
let terms: Vec<String> = doc.to_vec();
idx.add_document(i as u32, &terms).unwrap();
}
let n = idx.num_docs();
for t in idx.terms() {
let df = idx.df(t);
prop_assert!(df <= n);
}
}
}
proptest! {
#[test]
fn candidates_have_no_false_negatives(
docs in prop::collection::vec(
prop::collection::vec("[a-z]{1,6}", 0..20),
0..30
),
query in prop::collection::vec("[a-z]{1,6}", 0..10),
) {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
for (i, terms) in docs.iter().enumerate() {
let terms: Vec<String> = terms.to_vec();
idx.add_document(i as DocId, &terms).unwrap();
}
let q_terms: Vec<String> = query.to_vec();
let cands = idx.candidates(&q_terms);
let cand_set: std::collections::HashSet<DocId> = cands.into_iter().collect();
for doc_id in idx.document_ids() {
let mut hits = false;
for t in &q_terms {
if idx.term_frequency(doc_id, t) > 0 {
hits = true;
break;
}
}
if hits {
prop_assert!(cand_set.contains(&doc_id));
}
}
}
}
#[test]
fn planner_can_bail_out() {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
for i in 0..100u32 {
idx.add_document(i, &["common".to_string(), format!("u{i}")])
.unwrap();
}
let cfg = PlannerConfig {
max_candidate_ratio: 0.2,
max_candidates: 10,
};
let plan = idx.plan_candidates(&["common".to_string()], cfg);
assert_eq!(plan, CandidatePlan::ScanAll);
}
#[test]
fn generic_term_type_u32_works() {
let mut idx: PostingsIndex<u32> = PostingsIndex::new();
idx.add_document(0, &[1, 2, 2, 3]).unwrap();
idx.add_document(1, &[2, 4]).unwrap();
assert_eq!(idx.df(&2u32), 2);
assert_eq!(idx.term_frequency(0, &2u32), 2);
assert_eq!(idx.term_frequency(1, &2u32), 1);
let plan = idx.plan_candidates(
&[2u32],
PlannerConfig {
max_candidate_ratio: 1.0,
max_candidates: 10_000,
},
);
match plan {
CandidatePlan::Candidates(cands) => {
assert!(cands.contains(&0));
assert!(cands.contains(&1));
}
CandidatePlan::ScanAll => panic!("unexpected bailout for tiny corpus"),
}
}
#[test]
fn candidates_all_terms_intersects() {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
idx.add_document(0, &["a".into(), "b".into(), "b".into()])
.unwrap();
idx.add_document(1, &["a".into(), "c".into()]).unwrap();
idx.add_document(2, &["b".into(), "c".into()]).unwrap();
assert_eq!(
idx.candidates_all_terms(&["a".to_string(), "b".to_string()]),
vec![0]
);
assert_eq!(
idx.candidates_all_terms(&["b".to_string(), "c".to_string()]),
vec![2]
);
assert!(idx
.candidates_all_terms(&["missing".to_string()])
.is_empty());
}
#[test]
fn candidates_are_sorted_and_unique() {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
idx.add_document(2, &["a".into(), "b".into()]).unwrap();
idx.add_document(1, &["a".into()]).unwrap();
idx.add_document(3, &["b".into()]).unwrap();
let c = idx.candidates(&["b".to_string(), "a".to_string()]);
assert_eq!(c, vec![1, 2, 3]);
}
proptest! {
#[test]
fn candidates_all_terms_have_no_false_negatives(
docs in prop::collection::vec(
prop::collection::vec("[a-z]{1,6}", 0..20),
0..30
),
query in prop::collection::vec("[a-z]{1,6}", 0..10),
) {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
for (i, terms) in docs.iter().enumerate() {
let terms: Vec<String> = terms.to_vec();
idx.add_document(i as DocId, &terms).unwrap();
}
let q_terms: Vec<String> = query.to_vec();
let cands = idx.candidates_all_terms(&q_terms);
let cand_set: std::collections::HashSet<DocId> = cands.into_iter().collect();
let mut uniq: std::collections::HashSet<&String> = std::collections::HashSet::new();
for t in &q_terms {
uniq.insert(t);
}
for doc_id in idx.document_ids() {
let mut ok = !uniq.is_empty();
for t in &uniq {
if idx.term_frequency(doc_id, t.as_str()) == 0 {
ok = false;
break;
}
}
if ok {
prop_assert!(cand_set.contains(&doc_id));
}
}
}
}
proptest! {
#[test]
fn plan_candidates_candidates_respects_thresholds(
docs in prop::collection::vec(
prop::collection::vec("[a-z]{1,6}", 0..30),
0..60
),
query in prop::collection::vec("[a-z]{1,6}", 0..12),
max_ratio in 0.05f32..1.0f32,
max_abs in 1u32..5000u32,
) {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
for (i, doc) in docs.iter().enumerate() {
let terms: Vec<String> = doc.to_vec();
idx.add_document(i as DocId, &terms).unwrap();
}
let q_terms: Vec<String> = query.to_vec();
let cfg = PlannerConfig { max_candidate_ratio: max_ratio, max_candidates: max_abs };
let plan = idx.plan_candidates(&q_terms, cfg);
if let CandidatePlan::Candidates(_cands) = plan {
let n = idx.num_docs();
if n == 0 || q_terms.is_empty() {
return Ok(());
}
let mut seen: std::collections::HashSet<&String> = std::collections::HashSet::new();
let mut df_sum: u64 = 0;
for t in &q_terms {
if !seen.insert(t) { continue; }
df_sum = df_sum.saturating_add(idx.df(t) as u64);
}
prop_assert!(df_sum < (cfg.max_candidates as u64));
prop_assert!((df_sum as f32) / (n as f32) <= cfg.max_candidate_ratio);
}
}
}
#[test]
fn float_weighted_index_basic() {
let mut idx: PostingsIndex<String, f32> = PostingsIndex::new();
idx.add_weighted_document(
0,
&[
(String::from("neural"), 0.42),
(String::from("network"), 0.87),
(String::from("deep"), 0.15),
],
)
.unwrap();
assert_eq!(idx.num_docs(), 1);
assert!((idx.term_frequency(0, "neural") - 0.42).abs() < 1e-6);
assert!((idx.term_frequency(0, "network") - 0.87).abs() < 1e-6);
assert!((idx.term_frequency(0, "missing") - 0.0).abs() < 1e-6);
}
#[test]
fn float_weighted_candidates() {
let mut idx: PostingsIndex<String, f32> = PostingsIndex::new();
idx.add_weighted_document(0, &[(String::from("cat"), 0.9), (String::from("dog"), 0.3)])
.unwrap();
idx.add_weighted_document(
1,
&[(String::from("dog"), 0.8), (String::from("fish"), 0.5)],
)
.unwrap();
let cands = idx.candidates(&[String::from("dog")]);
assert_eq!(cands.len(), 2);
let cands = idx.candidates(&[String::from("cat")]);
assert_eq!(cands, vec![0]);
assert_eq!(idx.df("dog"), 2);
assert_eq!(idx.df("cat"), 1);
}
#[test]
fn float_weighted_delete() {
let mut idx: PostingsIndex<String, f32> = PostingsIndex::new();
idx.add_weighted_document(0, &[(String::from("a"), 0.5)])
.unwrap();
idx.add_weighted_document(1, &[(String::from("a"), 0.8)])
.unwrap();
assert_eq!(idx.df("a"), 2);
idx.delete_document(0);
assert_eq!(idx.df("a"), 1);
assert_eq!(idx.num_docs(), 1);
}
#[test]
fn float_weighted_accumulates_duplicates() {
let mut idx: PostingsIndex<String, f32> = PostingsIndex::new();
idx.add_weighted_document(
0,
&[(String::from("term"), 0.3), (String::from("term"), 0.4)],
)
.unwrap();
assert!((idx.term_frequency(0, "term") - 0.7).abs() < 1e-6);
}
#[test]
fn classic_u32_still_works_unchanged() {
let mut idx: PostingsIndex<String> = PostingsIndex::new();
idx.add_document(0, &[String::from("hello"), String::from("hello")])
.unwrap();
assert_eq!(idx.term_frequency(0, "hello"), 2);
assert_eq!(idx.document_len(0), 2); }
}