use super::super::codec::u64_value;
use super::{
cluster_id, encode_occurrence_cluster_controlled, keys, other_error, BTreeMap, ClusterChanges,
DocId, DocumentFields, FieldName, FieldStats, KeyValueBatch, OccurrencePosting, OccurrenceRead,
StorageBackendResult, TokenTermKey,
};
use crate::inverted_index::{
visit_field_replacement, InvertedIndexChange, InvertedIndexChangeVisitor,
};
#[cfg(test)]
mod tests;
fn cluster_replacement_changes(
entries: &[OccurrencePosting],
changes: &BTreeMap<DocId, Option<OccurrencePosting>>,
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<bool> {
for (doc_id, replacement) in changes {
control.check()?;
let previous = entries
.binary_search_by_key(doc_id, |entry| entry.doc_id)
.ok()
.map(|position| &entries[position]);
match (previous, replacement) {
(None, None) => {}
(Some(previous), Some(replacement)) => {
if previous.doc_length != replacement.doc_length
|| previous.occurrences.len() != replacement.occurrences.len()
{
return Ok(true);
}
for (old, new) in previous
.occurrences
.chunks(128)
.zip(replacement.occurrences.chunks(128))
{
control.check()?;
if old != new {
return Ok(true);
}
}
}
_ => return Ok(true),
}
}
control.check()?;
Ok(false)
}
fn visit_replacement(
visit: &mut InvertedIndexChangeVisitor<'_>,
doc_id: DocId,
previous: &DocumentFields,
replacement: &DocumentFields,
) -> StorageBackendResult<()> {
for field in previous.keys().chain(
replacement
.keys()
.filter(|field| !previous.contains_key(*field)),
) {
let old = previous.get(field);
let new = replacement.get(field);
visit_field_replacement(
visit,
doc_id,
field,
old.map(|snapshot| snapshot.metadata.length),
new.map(|snapshot| snapshot.metadata.length),
)?;
for snapshot in old.into_iter().chain(new) {
for term in snapshot.terms.keys() {
visit(InvertedIndexChange::Posting {
doc_id,
field,
term,
})?;
}
}
}
Ok(())
}
fn merge_cluster_changes(
entries: Vec<OccurrencePosting>,
changes: BTreeMap<DocId, Option<OccurrencePosting>>,
) -> Vec<OccurrencePosting> {
let mut merged = Vec::with_capacity(entries.len().saturating_add(changes.len()));
let mut changes = changes.into_iter().peekable();
for entry in entries {
while changes.peek().is_some_and(|(id, _)| *id < entry.doc_id) {
merged.extend(changes.next().expect("peeked change exists").1);
}
if changes.peek().is_some_and(|(id, _)| *id == entry.doc_id) {
merged.extend(changes.next().expect("peeked change exists").1);
} else {
merged.push(entry);
}
}
merged.extend(changes.filter_map(|(_, entry)| entry));
merged
}
impl OccurrenceRead<'_> {
pub(super) fn put_cluster(
&self,
batch: &mut dyn KeyValueBatch,
field: &str,
term: &TokenTermKey,
cluster: u64,
entries: &[OccurrencePosting],
) -> StorageBackendResult<()> {
let score_key = keys::cluster_key(self.table, keys::SCORE, field, term, cluster)?;
let graph_key = keys::cluster_key(self.table, keys::POSITIONS, field, term, cluster)?;
if entries.is_empty() {
batch.replace_occurrence_record(&score_key, None)?;
batch.replace_occurrence_record(&graph_key, None)?;
} else {
let (score, graph) =
encode_occurrence_cluster_controlled(entries.iter(), self.store.control())?;
batch.replace_occurrence_record(&score_key, Some(&score))?;
batch.replace_occurrence_record(&graph_key, Some(&graph))?;
}
Ok(())
}
pub(super) fn put_encoded_cluster(
&self,
batch: &mut dyn KeyValueBatch,
cluster: &crate::inverted_index::StagedCluster<'_>,
) -> StorageBackendResult<()> {
let key = |kind| {
keys::cluster_key(
self.table,
kind,
cluster.field,
cluster.term,
cluster.cluster,
)
};
batch.replace_occurrence_record(&key(keys::SCORE)?, Some(cluster.score))?;
batch.replace_occurrence_record(&key(keys::POSITIONS)?, Some(cluster.positions))
}
pub(super) fn put_document(
&self,
batch: &mut dyn KeyValueBatch,
doc_id: DocId,
fields: &DocumentFields,
) -> StorageBackendResult<()> {
for (field, snapshot) in fields {
batch.put(
&keys::metadata_key(self.table, field, doc_id)?,
&snapshot.metadata.to_bytes()?,
)?;
batch.put(
&keys::document_key(self.table, keys::LENGTH, doc_id, field)?,
&u64_value(snapshot.metadata.length),
)?;
batch.put(
&keys::document_key(self.table, keys::DOCUMENT, doc_id, field)?,
&crate::clustered_postings::encode_term_key_refs(snapshot.terms.keys())?,
)?;
}
Ok(())
}
pub(super) fn add_field_statistics(
totals: &mut BTreeMap<FieldName, FieldStats>,
fields: &DocumentFields,
) -> StorageBackendResult<()> {
for (field, snapshot) in fields {
let revision = snapshot.metadata.revision();
let stats = totals.entry(field.clone()).or_insert(FieldStats {
revision,
doc_count: 0,
total_length: 0,
});
if stats.revision != revision {
return Err(other_error(
"indexed field revisions disagree during replacement",
));
}
stats.doc_count = stats
.doc_count
.checked_add(1)
.ok_or_else(|| other_error("field document count overflow"))?;
stats.total_length = stats
.total_length
.checked_add(snapshot.metadata.length)
.ok_or_else(|| other_error("total field length overflow"))?;
}
Ok(())
}
pub(super) fn put_field_statistics(
&self,
batch: &mut dyn KeyValueBatch,
totals: BTreeMap<FieldName, FieldStats>,
) -> StorageBackendResult<()> {
for (field, stats) in totals {
let key = keys::field_prefix(self.table, keys::FIELD, &field)?;
if stats.doc_count == 0 {
if stats.total_length != 0 {
return Err(other_error("empty indexed field retains document length"));
}
batch.replace_occurrence_record(&key, None)?;
} else {
batch.replace_occurrence_record(&key, Some(&stats.to_bytes()?))?;
}
}
Ok(())
}
pub(super) fn add_documents(
&self,
batch: &mut dyn KeyValueBatch,
documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
mut visit: Option<&mut InvertedIndexChangeVisitor<'_>>,
) -> StorageBackendResult<()> {
self.require_graph_format()?;
let staged = self.stage_documents(documents, false)?;
let mut previous = BTreeMap::new();
let mut totals = BTreeMap::<FieldName, FieldStats>::new();
let mut changes = ClusterChanges::new();
for (doc_id, fields) in &staged {
batch.occurrence_document(self.table, *doc_id)?;
let old = self.old_document(*doc_id)?;
if let Some(visit) = visit.as_mut() {
visit_replacement(*visit, *doc_id, &old, fields)?;
}
for field in old.keys().chain(fields.keys()) {
if !totals.contains_key(field) {
if let Some(stats) = self.stored_field_stats(field)? {
totals.insert(field.clone(), stats);
}
}
}
for (field, snapshot) in &old {
let stats = totals
.get_mut(field)
.ok_or_else(|| other_error("indexed field revision is missing"))?;
stats.doc_count = stats
.doc_count
.checked_sub(1)
.ok_or_else(|| other_error("field document count underflow"))?;
stats.total_length = stats
.total_length
.checked_sub(snapshot.metadata.length)
.ok_or_else(|| other_error("total field length underflow"))?;
for term in snapshot.terms.keys() {
changes
.entry((field.clone(), term.clone(), cluster_id(*doc_id)))
.or_default()
.insert(*doc_id, None);
}
}
previous.insert(*doc_id, old);
}
for fields in staged.values() {
Self::add_field_statistics(&mut totals, fields)?;
}
if totals.is_empty() {
return Ok(());
}
self.invalidate_accelerators(batch)?;
for (doc_id, fields) in previous {
batch.delete_prefix(&keys::document_prefix(self.table, keys::DOCUMENT, doc_id)?)?;
batch.delete_prefix(&keys::document_prefix(self.table, keys::LENGTH, doc_id)?)?;
for field in fields.keys() {
batch.delete(&keys::metadata_key(self.table, field, doc_id)?)?;
}
}
for (doc_id, fields) in staged {
self.put_document(batch, doc_id, &fields)?;
for (field, snapshot) in fields {
for (term, occurrences) in snapshot.terms {
changes
.entry((field.clone(), term, cluster_id(doc_id)))
.or_default()
.insert(
doc_id,
Some(OccurrencePosting {
doc_id,
doc_length: snapshot.metadata.length,
occurrences,
}),
);
}
}
}
for ((field, term, cluster), updates) in changes {
let entries = self.load_cluster(&field, &term, cluster)?;
if cluster_replacement_changes(&entries, &updates, self.store.control())? {
let merged = merge_cluster_changes(entries, updates);
self.put_cluster(batch, &field, &term, cluster, &merged)?;
}
}
self.put_field_statistics(batch, totals)?;
batch.replace_occurrence_record(
&keys::kind_prefix(self.table, keys::FORMAT)?,
Some(keys::FORMAT_NAME),
)?;
Ok(())
}
}