use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use super::bm25::{Bm25Error, CorpusStats};
use super::postings::{Posting, PostingList, PostingsError};
use super::sealed::SealedSegment;
use super::tokenizer::{Analyzer, Token};
use crate::lifecycle::StoreError;
use crate::lifecycle::hash_table::allocation_bytes;
use crate::lifecycle::stats::{AccountedCounter, Accounting, AllocationComponent};
use crate::meta::DocBitmap;
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct FieldId(pub u16);
pub const DEFAULT_FIELD: FieldId = FieldId(0);
#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct TermKey {
pub term: Vec<u8>,
pub field: FieldId,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum IndexError {
CursorSlotMissing {
slot: usize,
},
LiveFrequencyCache {
reason: &'static str,
},
Postings(PostingsError),
Stats(Bm25Error),
RowsNotAscending {
row: u32,
},
LiveRowOutOfRange {
segment: usize,
row: u32,
row_count: u32,
},
LiveLengthMissing {
segment: usize,
row: u32,
},
ReplacedRowOutOfRange {
row: u32,
row_count: u32,
},
ReplacedPostingMissing {
row: u32,
},
}
impl std::fmt::Display for IndexError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::CursorSlotMissing { slot } => {
write!(formatter, "weighted cursor slot {slot} is absent")
}
Self::LiveFrequencyCache { reason } => {
write!(formatter, "live document frequency cache failed: {reason}")
}
Self::Postings(error) => write!(formatter, "postings rejected: {error}"),
Self::Stats(error) => write!(formatter, "corpus statistics unavailable: {error}"),
Self::RowsNotAscending { row } => {
write!(formatter, "documents must arrive in row order; got {row}")
}
Self::LiveRowOutOfRange {
segment,
row,
row_count,
} => write!(
formatter,
"live lexical row {row} is outside segment {segment} row count {row_count}"
),
Self::LiveLengthMissing { segment, row } => write!(
formatter,
"live lexical row {row} has no length counter in segment {segment}"
),
Self::ReplacedRowOutOfRange { row, row_count } => write!(
formatter,
"replaced lexical row {row} is outside the segment's {row_count} rows"
),
Self::ReplacedPostingMissing { row } => write!(
formatter,
"replaced lexical row {row} has no posting for a term of its previous document"
),
}
}
}
impl std::error::Error for IndexError {}
impl From<PostingsError> for IndexError {
fn from(error: PostingsError) -> Self {
Self::Postings(error)
}
}
impl From<Bm25Error> for IndexError {
fn from(error: Bm25Error) -> Self {
Self::Stats(error)
}
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Document {
fields: BTreeMap<FieldId, String>,
}
impl Document {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_text(text: &str) -> Self {
let mut document = Self::new();
document.set(DEFAULT_FIELD, text);
document
}
pub fn set(&mut self, field: FieldId, text: &str) -> &mut Self {
self.fields.insert(field, text.to_owned());
self
}
pub fn fields(&self) -> impl Iterator<Item = (FieldId, &str)> {
self.fields
.iter()
.map(|(field, text)| (*field, text.as_str()))
}
}
#[derive(Clone, Debug)]
struct FieldLengths {
field: FieldId,
lengths: Vec<u32>,
}
#[derive(Clone, Debug, Default)]
struct Xxh3State;
impl std::hash::BuildHasher for Xxh3State {
type Hasher = xxhash_rust::xxh3::Xxh3;
fn build_hasher(&self) -> Self::Hasher {
xxhash_rust::xxh3::Xxh3::new()
}
}
#[cfg(test)]
thread_local! {
static TOKENS_INDEXED: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
}
#[cfg(test)]
pub(crate) fn tokens_indexed() -> u64 {
TOKENS_INDEXED.with(std::cell::Cell::get)
}
#[cfg_attr(not(test), allow(unused_variables))]
fn count_tokens_indexed(tokens: usize) {
#[cfg(test)]
TOKENS_INDEXED.with(|count| {
count.set(
count
.get()
.saturating_add(u64::try_from(tokens).unwrap_or(u64::MAX)),
);
});
}
#[derive(Debug, Default)]
pub struct SegmentIndex {
terms: HashMap<TermKey, usize, Xxh3State>,
term_capacity: usize,
lists: Vec<PostingList>,
lengths: Vec<FieldLengths>,
row_count: u32,
free_slots: Vec<usize>,
term_bytes: usize,
posting_bytes: usize,
}
impl Clone for SegmentIndex {
fn clone(&self) -> Self {
self.clone_with_terms(self.terms.clone())
}
}
impl SegmentIndex {
fn clone_with_terms(&self, terms: HashMap<TermKey, usize, Xxh3State>) -> Self {
let mut copy = Self {
terms,
term_capacity: self.term_capacity,
lists: self.lists.clone(),
lengths: self.lengths.clone(),
row_count: self.row_count,
free_slots: self.free_slots.clone(),
term_bytes: 0,
posting_bytes: 0,
};
copy.term_bytes = copy
.terms
.keys()
.map(|key| key.term.capacity())
.fold(0, usize::saturating_add);
copy.posting_bytes = copy
.lists
.iter()
.map(PostingList::resident_bytes)
.fold(0, usize::saturating_add);
copy
}
}
#[derive(Clone, Copy)]
enum Placement {
Append,
Insert,
}
impl SegmentIndex {
pub(crate) fn clone_accounted(
&self,
accounting: &Arc<Accounting>,
) -> Result<(Self, AccountedCounter), StoreError> {
let mut counter = AccountedCounter::new(accounting, AllocationComponent::Active)?;
counter.set(self.resident_bytes())?;
let mut terms = HashMap::with_hasher(Xxh3State);
terms
.try_reserve(self.term_capacity)
.map_err(|_| StoreError::AllocationFailed {
needed: allocation_bytes::<(TermKey, usize)>(self.term_capacity)
.unwrap_or(usize::MAX) as u64,
component: "lexical terms map",
})?;
terms.extend(self.terms.iter().map(|(key, slot)| (key.clone(), *slot)));
let copy = self.clone_with_terms(terms);
counter.set(copy.resident_bytes())?;
Ok((copy, counter))
}
pub(crate) fn reserve_terms(
&mut self,
capacity: usize,
counter: &mut AccountedCounter,
) -> Result<(), StoreError> {
let allocation_capacity = if capacity <= self.term_capacity / 2 {
0
} else {
capacity.max(
self.term_capacity
.checked_add(1)
.ok_or(StoreError::ActiveRowOverflow)?,
)
};
let needed = allocation_bytes::<(TermKey, usize)>(allocation_capacity)
.ok_or(StoreError::ActiveRowOverflow)?;
counter.set(
self.resident_bytes()
.checked_add(needed)
.ok_or(StoreError::ActiveRowOverflow)?,
)?;
self.terms
.try_reserve(capacity.saturating_sub(self.terms.len()))
.map_err(|_| StoreError::AllocationFailed {
needed: needed as u64,
component: "lexical terms map",
})?;
self.term_capacity = self.term_capacity.max(self.terms.capacity());
counter.set(self.resident_bytes())
}
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub const fn row_count(&self) -> u32 {
self.row_count
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.row_count == 0
}
pub fn postings(&self) -> impl Iterator<Item = (&TermKey, &PostingList)> {
let mut entries: Vec<(&TermKey, usize)> =
self.terms.iter().map(|(key, slot)| (key, *slot)).collect();
entries.sort_unstable_by(|left, right| left.0.cmp(right.0));
entries
.into_iter()
.filter_map(|(key, slot)| self.lists.get(slot).map(|list| (key, list)))
}
#[must_use]
pub fn posting_list(&self, term: &[u8], field: FieldId) -> Option<&PostingList> {
let slot = *self.terms.get(&TermKey {
term: term.to_vec(),
field,
})?;
self.lists.get(slot)
}
#[must_use]
pub fn field_length(&self, row: u32, field: FieldId) -> u32 {
self.field_lengths(field)
.and_then(|lengths| {
usize::try_from(row)
.ok()
.and_then(|slot| lengths.get(slot).copied())
})
.unwrap_or(0)
}
pub fn fields(&self) -> impl Iterator<Item = FieldId> + '_ {
self.lengths.iter().map(|entry| entry.field)
}
#[must_use]
pub fn field_lengths(&self, field: FieldId) -> Option<&[u32]> {
self.lengths
.iter()
.find(|entry| entry.field == field)
.map(|entry| entry.lengths.as_slice())
}
#[must_use]
pub fn document_length(&self, row: u32) -> u32 {
let Ok(slot) = usize::try_from(row) else {
return 0;
};
self.lengths
.iter()
.filter_map(|entry| entry.lengths.get(slot).copied())
.fold(0_u32, u32::saturating_add)
}
#[must_use]
pub fn total_tokens(&self) -> u64 {
self.lengths
.iter()
.flat_map(|entry| entry.lengths.iter())
.map(|length| u64::from(*length))
.sum()
}
pub(crate) fn resident_bytes(&self) -> usize {
let term_slots =
allocation_bytes::<(TermKey, usize)>(self.term_capacity).unwrap_or(usize::MAX);
let term_bytes = self.term_bytes;
let list_bytes = self
.lists
.capacity()
.saturating_mul(std::mem::size_of::<PostingList>())
.saturating_add(self.posting_bytes);
let free_slot_bytes = self
.free_slots
.capacity()
.saturating_mul(std::mem::size_of::<usize>());
let length_bytes = self
.lengths
.capacity()
.saturating_mul(std::mem::size_of::<FieldLengths>())
.saturating_add(
self.lengths
.iter()
.map(|entry| entry.lengths.capacity().saturating_mul(4))
.fold(0_usize, usize::saturating_add),
);
term_slots
.saturating_add(term_bytes)
.saturating_add(list_bytes)
.saturating_add(length_bytes)
.saturating_add(free_slot_bytes)
}
#[cfg(test)]
pub(crate) fn resident_bytes_rescan(&self) -> usize {
let term_slots =
allocation_bytes::<(TermKey, usize)>(self.term_capacity).unwrap_or(usize::MAX);
let term_bytes = self
.terms
.keys()
.map(|key| key.term.capacity())
.fold(0_usize, usize::saturating_add);
let list_bytes = self
.lists
.capacity()
.saturating_mul(std::mem::size_of::<PostingList>())
.saturating_add(
self.lists
.iter()
.map(PostingList::resident_bytes)
.fold(0_usize, usize::saturating_add),
);
let free_slot_bytes = self
.free_slots
.capacity()
.saturating_mul(std::mem::size_of::<usize>());
let length_bytes = self
.lengths
.capacity()
.saturating_mul(std::mem::size_of::<FieldLengths>())
.saturating_add(
self.lengths
.iter()
.map(|entry| entry.lengths.capacity().saturating_mul(4))
.fold(0_usize, usize::saturating_add),
);
term_slots
.saturating_add(term_bytes)
.saturating_add(list_bytes)
.saturating_add(length_bytes)
.saturating_add(free_slot_bytes)
}
fn set_field_length(&mut self, row: u32, field: FieldId, length: u32) {
let Ok(slot) = usize::try_from(row) else {
return;
};
let position = match self.lengths.iter().position(|entry| entry.field == field) {
Some(position) => position,
None => {
let position = self
.lengths
.iter()
.position(|entry| entry.field > field)
.unwrap_or(self.lengths.len());
self.lengths.insert(
position,
FieldLengths {
field,
lengths: Vec::new(),
},
);
position
}
};
let Some(entry) = self.lengths.get_mut(position) else {
return;
};
if let Some(existing) = entry.lengths.get_mut(slot) {
*existing = length;
} else {
entry.lengths.resize(slot, 0);
entry.lengths.push(length);
}
}
fn pad_lengths_to(&mut self, rows: usize) {
for entry in &mut self.lengths {
entry.lengths.resize(rows, 0);
}
}
pub fn push_document(
&mut self,
analyzer: &Analyzer,
document: &Document,
) -> Result<u32, IndexError> {
self.push_document_controlled(analyzer, document, &mut |_, _| Ok(()))
}
pub(crate) fn push_document_controlled<E: From<IndexError>>(
&mut self,
analyzer: &Analyzer,
document: &Document,
reserve: &mut impl FnMut(&mut Self, usize) -> Result<(), E>,
) -> Result<u32, E> {
let row = self.row_count;
for (field, text) in document.fields() {
let tokens = analyzer.analyze(text);
self.accumulate(row, field, &tokens, Placement::Append, reserve)?;
}
self.finish_row();
Ok(row)
}
pub fn replace_document(
&mut self,
analyzer: &Analyzer,
row: u32,
previous: &Document,
replacement: &Document,
) -> Result<(), IndexError> {
self.replace_document_controlled(analyzer, row, previous, replacement, &mut |_, _| Ok(()))
}
pub(crate) fn replace_document_controlled<E: From<IndexError>>(
&mut self,
analyzer: &Analyzer,
row: u32,
previous: &Document,
replacement: &Document,
reserve: &mut impl FnMut(&mut Self, usize) -> Result<(), E>,
) -> Result<(), E> {
let out_of_range = IndexError::ReplacedRowOutOfRange {
row,
row_count: self.row_count,
};
if row >= self.row_count {
return Err(out_of_range.into());
}
let slot = usize::try_from(row).map_err(|_| out_of_range)?;
for (field, text) in previous.fields() {
let tokens = analyzer.analyze(text);
count_tokens_indexed(tokens.len());
let mut terms: Vec<&str> = tokens.iter().map(|token| token.term.as_str()).collect();
terms.sort_unstable();
terms.dedup();
let mut probe = TermKey {
term: Vec::new(),
field,
};
for term in terms {
probe.term.clear();
probe.term.extend_from_slice(term.as_bytes());
let missing = IndexError::ReplacedPostingMissing { row };
let list_slot = self.terms.get(&probe).copied().ok_or(missing.clone())?;
let list = self.lists.get_mut(list_slot).ok_or(missing.clone())?;
let posting_slot = list
.postings()
.binary_search_by_key(&row, |posting| posting.docid)
.map_err(|_| missing.clone())?;
let position_bytes = list
.postings()
.get(posting_slot)
.ok_or(missing.clone())?
.positions
.capacity()
.saturating_mul(4);
crate::fts::postings::count_accounting_visits(1);
if !list.remove(row) {
return Err(missing.into());
}
self.posting_bytes = self.posting_bytes.saturating_sub(position_bytes);
if list.is_empty() {
self.posting_bytes = self.posting_bytes.saturating_sub(list.allocation_bytes());
*list = PostingList::new();
if let Some((key, _)) = self.terms.remove_entry(&probe) {
self.term_bytes = self.term_bytes.saturating_sub(key.term.capacity());
}
self.free_slots.push(list_slot);
}
}
}
for entry in &mut self.lengths {
if let Some(length) = entry.lengths.get_mut(slot) {
*length = 0;
}
}
for (field, text) in replacement.fields() {
let tokens = analyzer.analyze(text);
self.accumulate(row, field, &tokens, Placement::Insert, reserve)?;
}
self.pad_lengths_to(usize::try_from(self.row_count).unwrap_or(usize::MAX));
Ok(())
}
pub fn push_documents(
&mut self,
analyzer: &Analyzer,
documents: &[Document],
threads: usize,
) -> Result<(), IndexError> {
let workers = threads.max(1).min(documents.len().max(1));
if workers <= 1 {
for document in documents {
self.push_document(analyzer, document)?;
}
return Ok(());
}
let batch_size = workers.saturating_mul(32).max(1);
for batch in documents.chunks(batch_size) {
let mut analyzed: Vec<Vec<(FieldId, Vec<Token>)>> = Vec::new();
analyzed.resize_with(batch.len(), Vec::new);
let stride = batch.len().div_ceil(workers);
std::thread::scope(|scope| {
let mut rest = analyzed.as_mut_slice();
let mut offset = 0_usize;
while !rest.is_empty() {
let take = stride.min(rest.len());
let (mine, tail) = rest.split_at_mut(take);
rest = tail;
let end = offset.saturating_add(take);
let docs = batch.get(offset..end).unwrap_or(&[]);
offset = end;
scope.spawn(move || {
for (slot, document) in mine.iter_mut().zip(docs) {
*slot = document
.fields()
.map(|(field, text)| (field, analyzer.analyze(text)))
.collect();
}
});
}
});
for fields in &analyzed {
let row = self.row_count;
for (field, tokens) in fields {
self.accumulate(row, *field, tokens, Placement::Append, &mut |_, _| {
Ok::<(), IndexError>(())
})?;
}
self.finish_row();
}
}
Ok(())
}
fn accumulate<E: From<IndexError>>(
&mut self,
row: u32,
field: FieldId,
tokens: &[Token],
placement: Placement,
reserve: &mut impl FnMut(&mut Self, usize) -> Result<(), E>,
) -> Result<(), E> {
count_tokens_indexed(tokens.len());
let length = tokens
.iter()
.map(|token| token.position)
.max()
.map_or(0, |highest| highest.saturating_add(1));
self.set_field_length(row, field, length);
let mut scratch: Vec<(&str, u32)> = tokens
.iter()
.map(|token| (token.term.as_str(), token.position))
.collect();
scratch.sort_unstable();
let mut probe = TermKey {
term: Vec::new(),
field,
};
let mut start = 0_usize;
while let Some((term, _)) = scratch.get(start).copied() {
let mut end = start.saturating_add(1);
while scratch.get(end).is_some_and(|(next, _)| *next == term) {
end = end.saturating_add(1);
}
let group = scratch.get(start..end).unwrap_or(&[]);
let mut positions: Vec<u32> = Vec::with_capacity(group.len());
for (_, position) in group {
if positions.last() != Some(position) {
positions.push(*position);
}
}
let tf = u32::try_from(positions.len()).unwrap_or(u32::MAX);
probe.term.clear();
probe.term.extend_from_slice(term.as_bytes());
let posting = Posting {
docid: row,
tf,
positions,
};
let position_bytes = posting.positions.capacity().saturating_mul(4);
crate::fts::postings::count_accounting_visits(1);
match self.terms.get(&probe).copied() {
Some(slot) => {
if let Some(list) = self.lists.get_mut(slot) {
let previous = list.allocation_bytes();
match placement {
Placement::Append => list.push(posting).map_err(IndexError::from)?,
Placement::Insert => list.insert(posting).map_err(IndexError::from)?,
}
self.posting_bytes = self
.posting_bytes
.saturating_sub(previous)
.saturating_add(list.allocation_bytes())
.saturating_add(position_bytes);
}
}
None => {
let grows = self.terms.len() == self.terms.capacity();
if grows {
reserve(self, self.terms.len().saturating_add(1))?;
}
let mut list = PostingList::new();
list.push(posting).map_err(IndexError::from)?;
self.posting_bytes = self
.posting_bytes
.saturating_add(list.allocation_bytes())
.saturating_add(position_bytes);
let slot = if let Some(slot) = self.free_slots.pop()
&& let Some(free) = self.lists.get_mut(slot)
{
*free = list;
slot
} else {
let slot = self.lists.len();
self.lists.push(list);
slot
};
let key = TermKey {
term: probe.term.clone(),
field,
};
self.term_bytes = self.term_bytes.saturating_add(key.term.capacity());
self.terms.insert(key, slot);
if grows {
self.term_capacity = self.term_capacity.max(self.terms.capacity());
}
}
}
start = end;
}
Ok(())
}
fn finish_row(&mut self) {
self.row_count = self.row_count.saturating_add(1);
self.pad_lengths_to(usize::try_from(self.row_count).unwrap_or(usize::MAX));
}
}
#[derive(Clone, Debug, Default)]
pub struct LexicalIndex {
segments: Vec<Arc<SealedSegment>>,
statistics: Vec<Arc<LiveSegmentStatistics>>,
has_frequency_cache: bool,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
struct LiveSegmentCounters {
documents: u64,
tokens: u64,
}
#[derive(Clone, Debug)]
pub(crate) struct LiveSegmentStatistics {
pub(crate) rows: DocBitmap,
counters: LiveSegmentCounters,
pub(crate) frequency_cache: Option<Arc<super::live_df::LiveFrequencyCache>>,
}
impl LiveSegmentStatistics {
pub(crate) fn build_controlled<E: From<IndexError>>(
ordinal: usize,
segment: &SealedSegment,
rows: &DocBitmap,
work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
) -> Result<Self, E> {
work.check_now()?;
let value = Self {
counters: live_segment_counters_controlled(ordinal, segment, rows, work)?,
rows: rows.clone(),
frequency_cache: None,
};
work.check_now()?;
Ok(value)
}
pub(crate) fn build(
ordinal: usize,
segment: &SealedSegment,
rows: &DocBitmap,
) -> Result<Self, IndexError> {
Self::build_controlled(
ordinal,
segment,
rows,
&mut super::control::WorkCheck::new(|| Ok(())),
)
}
pub(crate) fn allocation_bytes(rows: &DocBitmap) -> Option<usize> {
std::mem::size_of::<Self>()
.checked_add(2 * std::mem::size_of::<usize>())?
.checked_add(rows.resident_bytes()?)
}
}
impl LexicalIndex {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn push_segment(&mut self, segment: SegmentIndex) -> Result<(), IndexError> {
self.push_sealed(SealedSegment::seal(&segment)?);
Ok(())
}
pub fn push_sealed(&mut self, segment: SealedSegment) {
self.push_shared(Arc::new(segment));
}
fn push_shared(&mut self, segment: Arc<SealedSegment>) {
self.statistics.push(Arc::new(LiveSegmentStatistics {
rows: DocBitmap::full(segment.row_count()),
frequency_cache: None,
counters: LiveSegmentCounters {
documents: u64::from(segment.row_count()),
tokens: segment.total_tokens(),
},
}));
self.segments.push(segment);
}
pub fn push_sealed_with_live_rows(
&mut self,
segment: SealedSegment,
live_rows: &DocBitmap,
) -> Result<(), IndexError> {
self.push_shared_with_live_rows(Arc::new(segment), live_rows)
}
pub(crate) fn push_shared_with_live_rows(
&mut self,
segment: Arc<SealedSegment>,
live_rows: &DocBitmap,
) -> Result<(), IndexError> {
let statistics = Arc::new(LiveSegmentStatistics::build(
self.segments.len(),
&segment,
live_rows,
)?);
self.push_shared_with_statistics(segment, statistics);
Ok(())
}
pub(crate) fn push_shared_with_statistics(
&mut self,
segment: Arc<SealedSegment>,
statistics: Arc<LiveSegmentStatistics>,
) {
self.has_frequency_cache |= statistics.frequency_cache.is_some();
self.segments.push(segment);
self.statistics.push(statistics);
}
#[must_use]
pub fn segments(&self) -> &[Arc<SealedSegment>] {
&self.segments
}
pub(crate) fn resident_spine_bytes(&self) -> Option<usize> {
self.segments
.capacity()
.checked_mul(std::mem::size_of::<Arc<SealedSegment>>())?
.checked_add(
self.statistics
.capacity()
.checked_mul(std::mem::size_of::<Arc<LiveSegmentStatistics>>())?,
)
}
pub fn terms(&self) -> impl Iterator<Item = &[u8]> {
self.segments.iter().flat_map(|segment| segment.terms())
}
pub(crate) fn vocabulary_terms(&self) -> impl Iterator<Item = (&[u8], FieldId)> + Clone {
self.segments
.iter()
.flat_map(|segment| segment.vocabulary_terms())
}
#[must_use]
pub fn document_count(&self) -> u64 {
self.statistics
.iter()
.map(|statistics| statistics.counters.documents)
.fold(0_u64, u64::saturating_add)
}
#[must_use]
pub fn total_tokens(&self) -> u64 {
self.statistics
.iter()
.map(|statistics| statistics.counters.tokens)
.fold(0_u64, u64::saturating_add)
}
pub fn corpus_stats(&self) -> Result<CorpusStats, IndexError> {
#[cfg(any(test, feature = "test-seams"))]
super::preparation_observer::corpus_statistics();
Ok(CorpusStats::new(
self.document_count(),
self.total_tokens(),
)?)
}
#[must_use]
pub fn document_frequency(&self, term: &[u8], fields: &[FieldId]) -> u32 {
#[cfg(any(test, feature = "test-seams"))]
super::preparation_observer::frequency();
self.segments
.iter()
.zip(&self.statistics)
.map(|(segment, statistics)| {
let live_rows = &statistics.rows;
if live_rows.cardinality() == u64::from(segment.row_count()) {
segment.document_frequency(term, fields)
} else {
segment.live_document_frequency(term, fields, live_rows)
}
})
.fold(0_u32, u32::saturating_add)
}
#[inline]
pub(crate) fn prepared_document_frequency(
&self,
term: &[u8],
fields: &[FieldId],
) -> Result<u32, IndexError> {
if !self.has_frequency_cache {
return Ok(self.document_frequency(term, fields));
}
self.cached_document_frequency(term, fields)
}
pub(crate) fn prepared_document_frequency_controlled<E>(
&self,
term: &[u8],
fields: &[FieldId],
mut check: impl FnMut() -> Result<(), E>,
) -> Result<u32, super::search::ControlledSearchError<E>> {
use super::search::ControlledSearchError;
#[cfg(any(test, feature = "test-seams"))]
super::preparation_observer::frequency();
check().map_err(ControlledSearchError::Control)?;
let mut total = 0_u32;
for (segment, statistics) in self.segments.iter().zip(&self.statistics) {
let live = &statistics.rows;
let frequency = if live.cardinality() == u64::from(segment.row_count()) {
segment
.document_frequency_controlled(term, fields, &mut check)
.map_err(ControlledSearchError::Control)?
} else if let Some(cache) = &statistics.frequency_cache {
cache.frequency_controlled(segment, live, term, fields, &mut check)?
} else {
segment
.live_document_frequency_controlled(term, fields, live, &mut check)
.map_err(ControlledSearchError::Control)?
};
total = total.saturating_add(frequency);
}
check().map_err(ControlledSearchError::Control)?;
Ok(total)
}
fn cached_document_frequency(
&self,
term: &[u8],
fields: &[FieldId],
) -> Result<u32, IndexError> {
#[cfg(any(test, feature = "test-seams"))]
super::preparation_observer::frequency();
self.segments.iter().zip(&self.statistics).try_fold(
0_u32,
|total, (segment, statistics)| {
let live = &statistics.rows;
let frequency = if live.cardinality() == u64::from(segment.row_count()) {
segment.document_frequency(term, fields)
} else if let Some(cache) = &statistics.frequency_cache {
cache.frequency(segment, live, term, fields)?
} else {
segment.live_document_frequency(term, fields, live)
};
Ok(total.saturating_add(frequency))
},
)
}
}
fn live_segment_counters_controlled<E: From<IndexError>>(
ordinal: usize,
segment: &SealedSegment,
live_rows: &DocBitmap,
work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
) -> Result<LiveSegmentCounters, E> {
work.check_now()?;
if let Some(row) = live_rows.first_at_or_after(segment.row_count()) {
return Err(IndexError::LiveRowOutOfRange {
segment: ordinal,
row,
row_count: segment.row_count(),
}
.into());
}
let tokens =
live_rows.iter().try_fold(0_u64, |total, row| {
work.step()?;
#[cfg(any(test, feature = "test-seams"))]
super::preparation_observer::live_statistics_row();
let slot = usize::try_from(row).map_err(|_| IndexError::LiveLengthMissing {
segment: ordinal,
row,
})?;
let length = segment.total_lengths().get(slot).copied().ok_or(
IndexError::LiveLengthMissing {
segment: ordinal,
row,
},
)?;
Ok::<u64, E>(total.saturating_add(u64::from(length)))
})?;
work.check_now()?;
Ok(LiveSegmentCounters {
documents: live_rows.cardinality(),
tokens,
})
}
#[cfg(test)]
#[allow(
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic,
clippy::unwrap_used
)]
mod tests {
#[test]
fn terms_table_charge_includes_empty_buckets_and_controls() {
let mut index = SegmentIndex::new();
index
.push_document(&analyzer(), &Document::with_text("alpha"))
.expect("index");
let group = if cfg!(all(
any(target_arch = "x86", target_arch = "x86_64"),
target_feature = "sse2",
not(miri)
)) {
16
} else {
8
};
let other = index.term_bytes
+ index.posting_bytes
+ index.lists.capacity() * std::mem::size_of::<PostingList>()
+ index.lengths.capacity() * std::mem::size_of::<FieldLengths>()
+ index
.lengths
.iter()
.map(|entry| entry.lengths.capacity() * 4)
.sum::<usize>();
assert_eq!(
index.resident_bytes() - other,
4 * std::mem::size_of::<(TermKey, usize)>() + 4 + group
);
}
#[test]
fn parallel_analysis_is_identical_to_the_sequential_loop() {
let analyzer = analyzer();
let documents: Vec<Document> = (0..97)
.map(|index| {
let mut document = Document::new();
document.set(
FieldId(0),
&format!("alpha beta{} gamma-{index} 401k", index % 7),
);
document.set(FieldId(1), &format!("delta epsilon{} the quick", index % 5));
document
})
.collect();
let mut sequential = SegmentIndex::new();
for document in &documents {
sequential
.push_document(&analyzer, document)
.expect("indexable");
}
for threads in [1_usize, 3, 8] {
let mut parallel = SegmentIndex::new();
parallel
.push_documents(&analyzer, &documents, threads)
.expect("indexable");
assert_eq!(parallel.row_count(), sequential.row_count());
let ours: Vec<_> = parallel.postings().collect();
let oracle: Vec<_> = sequential.postings().collect();
assert_eq!(
ours.len(),
oracle.len(),
"term count diverged at {threads} threads"
);
for ((key, list), (oracle_key, oracle_list)) in ours.iter().zip(&oracle) {
assert_eq!(key, oracle_key);
assert_eq!(
list.postings(),
oracle_list.postings(),
"postings diverged for {:?} at {threads} threads",
String::from_utf8_lossy(&key.term)
);
}
for field in [FieldId(0), FieldId(1)] {
assert_eq!(
parallel.field_lengths(field),
sequential.field_lengths(field),
"lengths diverged at {threads} threads"
);
}
}
}
use super::*;
use crate::fts::tokenizer::{Profile, TokenizerConfig};
#[test]
fn astra_18_live_statistics_cancel_inside_row_walk() {
#[derive(Debug)]
enum Error {
Index(IndexError),
Cancelled,
}
impl From<IndexError> for Error {
fn from(error: IndexError) -> Self {
Self::Index(error)
}
}
let segment = SealedSegment::seal(&segment_of(&analyzer(), &vec!["alpha beta"; 4_097]))
.expect("sealed row fixture");
let live = DocBitmap::from_ids(0..4_096);
super::super::preparation_observer::begin();
let mut work = super::super::control::WorkCheck::new(|| {
if super::super::preparation_observer::live_statistics_rows() >= 64 {
Err(Error::Cancelled)
} else {
Ok(())
}
});
let result = LiveSegmentStatistics::build_controlled(7, &segment, &live, &mut work);
let visited = super::super::preparation_observer::live_statistics_rows();
super::super::preparation_observer::take();
println!("live statistics rows={visited}");
assert!(matches!(result, Err(Error::Cancelled)));
assert!((64..=128).contains(&visited));
let mut clean = super::super::control::WorkCheck::new(|| Ok::<(), Error>(()));
let value = LiveSegmentStatistics::build_controlled(7, &segment, &live, &mut clean)
.expect("fresh control");
assert_eq!(value.counters.documents, 4_096);
assert_eq!(value.counters.tokens, 8_192);
assert_eq!(value.rows, live);
assert!(value.frequency_cache.is_none());
assert!(matches!(
LiveSegmentStatistics::build_controlled(
7,
&segment,
&DocBitmap::from_ids([4_097]),
&mut clean
),
Err(Error::Index(IndexError::LiveRowOutOfRange {
segment: 7,
row: 4_097,
row_count: 4_097
}))
));
}
fn analyzer() -> Analyzer {
Analyzer::new(TokenizerConfig::text_default()).expect("valid config")
}
#[test]
fn replace_document_refuses_a_row_or_previous_text_the_index_does_not_hold() {
let analyzer = analyzer();
let mut index = SegmentIndex::new();
index
.push_document(&analyzer, &Document::with_text("alpha beta"))
.expect("indexable");
let replacement = Document::with_text("gamma");
assert_eq!(
index.replace_document(&analyzer, 1, &Document::with_text("alpha"), &replacement),
Err(IndexError::ReplacedRowOutOfRange {
row: 1,
row_count: 1
})
);
assert_eq!(
index.replace_document(&analyzer, 0, &Document::with_text("delta"), &replacement),
Err(IndexError::ReplacedPostingMissing { row: 0 })
);
index
.push_document(&analyzer, &Document::with_text("epsilon"))
.expect("indexable");
assert_eq!(
index.replace_document(&analyzer, 0, &Document::with_text("epsilon"), &replacement),
Err(IndexError::ReplacedPostingMissing { row: 0 })
);
assert!(
IndexError::ReplacedPostingMissing { row: 0 }
.to_string()
.contains("row 0")
);
assert!(
IndexError::ReplacedRowOutOfRange {
row: 1,
row_count: 1
}
.to_string()
.contains("1 rows")
);
}
#[test]
fn replace_document_reuses_a_dropped_term_slot_and_pads_a_new_field() {
let analyzer = analyzer();
let mut index = SegmentIndex::new();
for text in ["alpha", "beta", "alpha delta"] {
index
.push_document(&analyzer, &Document::with_text(text))
.expect("indexable");
}
let mut replacement = Document::with_text("alpha");
replacement.set(FieldId(1), "gamma");
index
.replace_document(&analyzer, 1, &Document::with_text("beta"), &replacement)
.expect("replace row 1");
assert_eq!(index.lists.len(), 3);
assert!(index.posting_list(b"beta", DEFAULT_FIELD).is_none());
assert!(index.free_slots.is_empty());
assert_eq!(
index.posting_list(b"alpha", DEFAULT_FIELD).map(|list| list
.postings()
.iter()
.map(|p| p.docid)
.collect::<Vec<_>>()),
Some(vec![0, 1, 2])
);
assert_eq!(index.field_lengths(FieldId(1)), Some(&[0, 1, 0][..]));
index
.replace_document(
&analyzer,
2,
&Document::with_text("alpha delta"),
&Document::new(),
)
.expect("replace row 2");
assert!(index.posting_list(b"delta", DEFAULT_FIELD).is_none());
assert_eq!(index.free_slots.len(), 1);
assert_eq!(index.total_tokens(), 3);
}
#[test]
fn posting_list_insert_refuses_a_duplicate_row() {
let mut list = PostingList::new();
let posting = |docid| Posting {
docid,
tf: 1,
positions: vec![0],
};
list.push(posting(0)).expect("append");
list.push(posting(4)).expect("append");
list.insert(posting(2)).expect("insert between");
assert_eq!(
list.insert(posting(2)),
Err(PostingsError::DocidsNotAscending { docid: 2 })
);
assert_eq!(
list.insert(Posting {
docid: 3,
tf: 0,
positions: Vec::new()
}),
Err(PostingsError::ZeroTermFrequency)
);
assert_eq!(
list.insert(Posting {
docid: 3,
tf: 2,
positions: vec![1, 1]
}),
Err(PostingsError::PositionsNotAscending { position: 1 })
);
assert!(list.remove(2));
assert!(!list.remove(2));
assert_eq!(
list.postings().iter().map(|p| p.docid).collect::<Vec<_>>(),
vec![0, 4]
);
}
fn code_analyzer() -> Analyzer {
Analyzer::new(Profile::Code.config()).expect("valid config")
}
fn segment_of(analyzer: &Analyzer, texts: &[&str]) -> SegmentIndex {
let mut segment = SegmentIndex::new();
for text in texts {
segment
.push_document(analyzer, &Document::with_text(text))
.expect("indexable");
}
segment
}
#[test]
fn rows_are_dense_and_ascending() {
let analyzer = analyzer();
let mut segment = SegmentIndex::new();
for (expected, text) in ["alpha", "beta", "gamma"].into_iter().enumerate() {
let row = segment
.push_document(&analyzer, &Document::with_text(text))
.expect("indexable");
assert_eq!(row, u32::try_from(expected).expect("small"));
}
assert_eq!(segment.row_count(), 3);
}
#[test]
fn document_frequency_is_summed_across_segments_not_per_segment() {
let analyzer = code_analyzer();
let mut index = LexicalIndex::new();
index
.push_segment(segment_of(&analyzer, &["engine", "engine"]))
.expect("seals");
index
.push_segment(segment_of(&analyzer, &["engine", "other"]))
.expect("seals");
assert_eq!(index.document_count(), 4);
assert_eq!(index.document_frequency(b"engine", &[DEFAULT_FIELD]), 3);
}
#[test]
fn a_term_in_two_fields_of_one_document_counts_once() {
let analyzer = code_analyzer();
let mut segment = SegmentIndex::new();
let mut document = Document::new();
document.set(FieldId(0), "engine");
document.set(FieldId(1), "engine");
segment
.push_document(&analyzer, &document)
.expect("indexable");
let mut index = LexicalIndex::new();
index.push_segment(segment).expect("seals");
assert_eq!(
index.document_frequency(b"engine", &[FieldId(0), FieldId(1)]),
1
);
}
#[test]
fn corpus_stats_span_every_segment() {
let analyzer = code_analyzer();
let mut index = LexicalIndex::new();
index
.push_segment(segment_of(&analyzer, &["a b c"]))
.expect("seals");
index
.push_segment(segment_of(&analyzer, &["d e f g h"]))
.expect("seals");
let stats = index.corpus_stats().expect("two documents");
assert_eq!(stats.document_count(), 2);
assert_eq!(stats.total_tokens(), 8);
assert!((stats.average_document_length() - 4.0).abs() < 1e-12);
}
#[test]
fn live_row_counters_and_document_frequency_exclude_tombstones() {
let analyzer = code_analyzer();
let sealed = SealedSegment::seal(&segment_of(
&analyzer,
&["engine padding padding", "engine"],
))
.expect("seal lexical segment");
let live_rows = DocBitmap::from_ids([1]);
let mut index = LexicalIndex::new();
index
.push_sealed_with_live_rows(sealed, &live_rows)
.expect("valid live rows");
let stats = index.corpus_stats().expect("one live document");
assert_eq!(stats.document_count(), 1);
assert_eq!(stats.total_tokens(), 1);
assert_eq!(
index.document_frequency(b"engine", &[DEFAULT_FIELD]),
1,
"document frequency must use the same live rows as corpus stats"
);
}
#[test]
fn an_empty_index_refuses_to_produce_statistics() {
let index = LexicalIndex::new();
assert!(index.corpus_stats().is_err());
}
#[test]
fn document_length_counts_positions_not_tokens() {
let analyzer = code_analyzer();
let mut segment = SegmentIndex::new();
segment
.push_document(&analyzer, &Document::with_text("put_if_match"))
.expect("indexable");
assert_eq!(segment.document_length(0), 3);
}
#[test]
fn field_lengths_are_tracked_separately_and_summed() {
let analyzer = code_analyzer();
let mut segment = SegmentIndex::new();
let mut document = Document::new();
document.set(FieldId(0), "one two");
document.set(FieldId(1), "three four five");
segment
.push_document(&analyzer, &document)
.expect("indexable");
assert_eq!(segment.field_length(0, FieldId(0)), 2);
assert_eq!(segment.field_length(0, FieldId(1)), 3);
assert_eq!(segment.document_length(0), 5);
assert_eq!(segment.total_tokens(), 5);
assert_eq!(segment.field_length(0, FieldId(9)), 0);
}
#[test]
fn posting_lists_carry_positions() {
let analyzer = code_analyzer();
let segment = segment_of(&analyzer, &["alpha beta alpha"]);
let list = segment
.posting_list(b"alpha", DEFAULT_FIELD)
.expect("alpha present");
let posting = list.postings().first().expect("one posting");
assert_eq!(posting.tf, 2);
assert_eq!(posting.positions, vec![0, 2]);
}
#[test]
fn an_absent_term_has_no_posting_list_and_zero_frequency() {
let analyzer = code_analyzer();
let segment = segment_of(&analyzer, &["alpha"]);
assert!(segment.posting_list(b"omega", DEFAULT_FIELD).is_none());
let mut index = LexicalIndex::new();
index.push_segment(segment).expect("seals");
assert_eq!(index.document_frequency(b"omega", &[DEFAULT_FIELD]), 0);
}
#[test]
fn errors_render_useful_text() {
assert!(
IndexError::RowsNotAscending { row: 3 }
.to_string()
.contains("row order")
);
assert!(
IndexError::Stats(Bm25Error::EmptyCorpus)
.to_string()
.contains("statistics")
);
assert!(
IndexError::Postings(PostingsError::ZeroTermFrequency)
.to_string()
.contains("postings")
);
}
}