use alloc::vec::Vec;
#[cfg(feature = "counters")]
use core::cell::Cell;
use plugmem_arena::{Arena, ArenaCfg, ShardMode, Slot, key};
use crate::error::Error;
use crate::id::FactId;
use crate::index::postings::PostingStore;
mod doclen_at {
use core::mem::size_of;
pub(super) const FACT: usize = 0;
pub(super) const KEY_LEN: usize = FACT + size_of::<u32>();
pub(super) const LEN: usize = KEY_LEN;
pub(super) const DISTINCT: usize = LEN + size_of::<u16>();
pub(super) const SIG: usize = DISTINCT + size_of::<u16>();
pub(super) const SIZE: usize = SIG + size_of::<u64>();
}
const SIG_BITS: u32 = 64;
const SIG_MULT: u64 = 0x9E37_79B9_7F4A_7C15;
pub(crate) fn sig_bit(term: u32) -> u64 {
let index = u64::from(term).wrapping_mul(SIG_MULT) >> (64 - SIG_BITS.trailing_zeros());
1u64 << index
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct DocLenSlot {
pub fact: FactId,
pub len: u16,
pub distinct: u16,
pub sig: u64,
}
impl DocLenSlot {
pub fn has_signature(&self) -> bool {
self.distinct != 0
}
pub fn overlap_bound(&self, terms: &[u32]) -> usize {
terms
.iter()
.filter(|&&term| self.sig & sig_bit(term) != 0)
.count()
}
}
impl Slot for DocLenSlot {
const SIZE: usize = doclen_at::SIZE;
const KEY_LEN: usize = doclen_at::KEY_LEN;
fn write(&self, out: &mut [u8]) {
key::write_u32(&mut out[doclen_at::FACT..], self.fact.0);
out[doclen_at::LEN..doclen_at::DISTINCT].copy_from_slice(&self.len.to_be_bytes());
out[doclen_at::DISTINCT..doclen_at::SIG].copy_from_slice(&self.distinct.to_be_bytes());
out[doclen_at::SIG..doclen_at::SIZE].copy_from_slice(&self.sig.to_be_bytes());
}
fn read(bytes: &[u8]) -> Self {
Self {
fact: FactId(key::read_u32(&bytes[doclen_at::FACT..])),
len: u16::from_be_bytes(
bytes[doclen_at::LEN..doclen_at::DISTINCT]
.try_into()
.unwrap(),
),
distinct: u16::from_be_bytes(
bytes[doclen_at::DISTINCT..doclen_at::SIG]
.try_into()
.unwrap(),
),
sig: u64::from_be_bytes(bytes[doclen_at::SIG..doclen_at::SIZE].try_into().unwrap()),
}
}
}
#[derive(Debug, Default)]
pub struct Bm25Scratch {
acc: Vec<(u32, f32)>,
merge: Vec<(u32, f32)>,
top: Vec<(f32, u32)>,
}
impl Bm25Scratch {
pub fn new() -> Self {
Self::default()
}
}
#[derive(Debug)]
pub struct Bm25Index<'a> {
postings: PostingStore<'a, true>,
doc_len: Arena<'a, DocLenSlot>,
total_docs: u64,
total_len: u64,
#[cfg(feature = "counters")]
decoded: Cell<u64>,
#[cfg(feature = "counters")]
scored: Cell<u64>,
#[cfg(feature = "counters")]
admitted: Cell<u64>,
unsummarized: bool,
doc_len_dense: Vec<u32>,
dense_limit: usize,
}
const DOC_LEN_ABSENT: u32 = u32::MAX;
impl<'a> Bm25Index<'a> {
pub fn new(shards: usize, max_bytes: usize) -> Result<Self, Error> {
Ok(Self {
postings: PostingStore::new(shards, max_bytes)?,
doc_len: Arena::new(
ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(max_bytes),
)?,
total_docs: 0,
total_len: 0,
#[cfg(feature = "counters")]
decoded: Cell::new(0),
#[cfg(feature = "counters")]
scored: Cell::new(0),
#[cfg(feature = "counters")]
admitted: Cell::new(0),
unsummarized: false,
doc_len_dense: Vec::new(),
dense_limit: usize::MAX,
})
}
pub fn index_doc(&mut self, fact: FactId, term_tfs: &[(u32, u8)]) -> Result<(), Error> {
let mut len = 0u32;
let mut sig = 0u64;
for &(term, tf) in term_tfs {
self.postings.push(term, fact, tf)?;
len += u32::from(tf);
sig |= sig_bit(term);
}
let doc = DocLenSlot {
fact,
len: u16::try_from(len).unwrap_or(u16::MAX),
distinct: u16::try_from(term_tfs.len()).unwrap_or(u16::MAX),
sig,
};
self.doc_len.insert(&doc)?;
self.note_dense(&doc);
self.total_docs += 1;
self.total_len += u64::from(len);
Ok(())
}
fn note_dense(&mut self, doc: &DocLenSlot) {
let at = doc.fact.0 as usize;
if at >= self.dense_capacity() {
self.dense_limit = self.dense_limit.min(at);
return;
}
if at >= self.doc_len_dense.len() {
self.doc_len_dense.resize(at + 1, DOC_LEN_ABSENT);
}
self.doc_len_dense[at] = u32::from(doc.len);
}
fn dense_capacity(&self) -> usize {
const SLACK: usize = 8;
let docs = self.total_docs.min(usize::MAX as u64) as usize;
docs.saturating_add(1).saturating_mul(SLACK)
}
fn rebuild_dense(&mut self) {
self.doc_len_dense.clear();
self.dense_limit = usize::MAX;
let docs: Vec<DocLenSlot> = self.doc_len.iter().collect();
for doc in docs {
self.note_dense(&doc);
}
}
pub fn doc(&self, fact: FactId) -> Option<DocLenSlot> {
self.doc_len.get(&fact.0.to_be_bytes())
}
pub(crate) fn needs_resummarize(&self) -> bool {
self.unsummarized
}
pub(crate) fn mark_unsummarized(&mut self) {
self.unsummarized = true;
}
pub(crate) fn compact_live(
&self,
shards: usize,
max_bytes: usize,
mut live: impl FnMut(FactId) -> bool,
) -> Result<Bm25Index<'static>, Error> {
let mut out = Bm25Index::new(shards, max_bytes)?;
let mut legacy = 0usize;
let mut max_fact = 0u32;
for doc in self.doc_len.iter() {
if live(doc.fact) {
out.doc_len.insert(&doc)?;
out.note_dense(&doc);
out.total_docs += 1;
out.total_len += u64::from(doc.len);
if !doc.has_signature() {
legacy += 1;
max_fact = max_fact.max(doc.fact.0);
}
}
}
let highest = max_fact as usize;
let mut rebuilt = if legacy > 0 && highest < out.dense_capacity() {
alloc::vec![(0u64, 0u16); highest + 1]
} else {
Vec::new()
};
for slot in self.postings.slots() {
for (fact, tf) in self.postings.entries(slot.key) {
if !live(fact) {
continue;
}
out.postings.push(slot.key, fact, tf)?;
if let Some(entry) = rebuilt.get_mut(fact.0 as usize) {
entry.0 |= sig_bit(slot.key);
entry.1 = entry.1.saturating_add(1);
}
}
}
if !rebuilt.is_empty() {
out.fill_missing_signatures(&rebuilt);
}
Ok(out)
}
fn fill_missing_signatures(&mut self, rebuilt: &[(u64, u16)]) {
let stale: Vec<DocLenSlot> = self
.doc_len
.iter()
.filter(|doc| !doc.has_signature())
.collect();
for mut doc in stale {
let Some(&(sig, distinct)) = rebuilt.get(doc.fact.0 as usize) else {
continue;
};
if distinct == 0 {
continue; }
doc.sig = sig;
doc.distinct = distinct;
let Some(payload) = self.doc_len.payload_mut(&doc.fact.0.to_be_bytes()) else {
continue;
};
let mut full = [0u8; DocLenSlot::SIZE];
doc.write(&mut full);
payload.copy_from_slice(&full[DocLenSlot::KEY_LEN..]);
}
}
pub fn df(&self, term: u32) -> u32 {
self.postings.count(term)
}
pub fn docs(&self) -> u64 {
self.total_docs
}
pub fn idf(&self, df: u32) -> f32 {
let n = self.total_docs as f32;
let df = df as f32;
libm::logf(1.0 + (n - df + 0.5) / (df + 0.5))
}
pub fn search(
&self,
(k1, b): (f32, f32),
terms: &[u32],
k: usize,
live: &mut dyn FnMut(FactId) -> bool,
scratch: &mut Bm25Scratch,
out: &mut Vec<(FactId, f32)>,
) {
out.clear();
if self.total_docs == 0 || k == 0 {
return;
}
let Bm25Scratch { acc, merge, top } = scratch;
acc.clear();
let avg_len = self.total_len as f32 / self.total_docs as f32;
#[cfg(feature = "counters")]
let (mut decoded, mut scored) = (0u64, 0u64);
for &term in terms {
let df = self.postings.count(term);
if df == 0 {
continue;
}
let idf = self.idf(df);
let mut ahead = 0usize;
merge.clear();
for (fact, tf) in self.postings.entries(term) {
#[cfg(feature = "counters")]
{
decoded += 1;
}
while let Some(&entry) = acc.get(ahead)
&& entry.0 < fact.0
{
merge.push(entry);
ahead += 1;
}
let carried = match acc.get(ahead) {
Some(&entry) if entry.0 == fact.0 => {
ahead += 1;
Some(entry.1)
}
_ => None,
};
let Some(len) = self.doc_len_of(fact) else {
if let Some(score) = carried {
merge.push((fact.0, score));
}
continue;
};
#[cfg(feature = "counters")]
{
scored += 1;
}
let tf = f32::from(tf);
let norm = tf * (k1 + 1.0) / (tf + k1 * (1.0 - b + b * f32::from(len) / avg_len));
merge.push((fact.0, carried.unwrap_or(0.0) + idf * norm));
}
merge.extend_from_slice(&acc[ahead.min(acc.len())..]);
core::mem::swap(acc, merge);
}
#[cfg(feature = "counters")]
{
self.decoded.set(self.decoded.get() + decoded);
self.scored.set(self.scored.get() + scored);
}
top.clear();
top.extend(acc.iter().map(|&(id, score)| (score, id)));
let order = |a: &(f32, u32), b: &(f32, u32)| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1));
#[cfg(feature = "counters")]
let mut admitted = 0u64;
let mut consume = |band: &[(f32, u32)], out: &mut Vec<(FactId, f32)>| {
for &(score, id) in band {
if out.len() == k {
return;
}
#[cfg(feature = "counters")]
{
admitted += 1;
}
if live(FactId(id)) {
out.push((FactId(id), score));
}
}
};
let band = k.min(top.len());
if band > 0 {
if band < top.len() {
top.select_nth_unstable_by(band - 1, order);
}
top[..band].sort_unstable_by(order);
consume(&top[..band], out);
}
if out.len() < k && band < top.len() {
top[band..].sort_unstable_by(order);
let (_, rest) = top.split_at(band);
consume(rest, out);
}
#[cfg(feature = "counters")]
self.admitted.set(self.admitted.get() + admitted);
}
fn doc_len_of(&self, fact: FactId) -> Option<u16> {
let at = fact.0 as usize;
if at < self.dense_limit
&& let Some(&len) = self.doc_len_dense.get(at)
{
return (len != DOC_LEN_ABSENT).then_some(len as u16);
}
self.doc_len.get(&fact.0.to_be_bytes()).map(|doc| doc.len)
}
pub fn pool_bytes(&self) -> usize {
self.postings.pool_bytes() + self.doc_len.pool_bytes()
}
pub fn total_len(&self) -> u64 {
self.total_len
}
pub(crate) fn postings(&self) -> &PostingStore<'a, true> {
&self.postings
}
pub(crate) fn doc_len_arena(&self) -> &Arena<'a, DocLenSlot> {
&self.doc_len
}
pub(crate) fn from_parts(
postings: PostingStore<'a, true>,
doc_len: Arena<'a, DocLenSlot>,
total_docs: u64,
total_len: u64,
) -> Self {
let mut index = Self {
postings,
doc_len,
total_docs,
total_len,
#[cfg(feature = "counters")]
decoded: Cell::new(0),
#[cfg(feature = "counters")]
scored: Cell::new(0),
#[cfg(feature = "counters")]
admitted: Cell::new(0),
unsummarized: false,
doc_len_dense: Vec::new(),
dense_limit: usize::MAX,
};
index.rebuild_dense();
index
}
#[cfg(feature = "counters")]
pub fn decoded(&self) -> u64 {
self.decoded.get()
}
#[cfg(feature = "counters")]
pub fn scored(&self) -> u64 {
self.scored.get()
}
#[cfg(feature = "counters")]
pub fn admitted(&self) -> u64 {
self.admitted.get()
}
#[cfg(feature = "counters")]
pub fn reset_query_counters(&self) {
self.decoded.set(0);
self.scored.set(0);
self.admitted.set(0);
}
}