use super::{
keys, other_error, BTreeMap, DocId, KeyValueBatch, KeyValueInvertedIndex, OccurrenceRead,
StorageBackendResult, TokenTermKey,
};
use crate::block_max_index::{BlockMaxIndex, BlockMaxScorer, DEFAULT_BLOCK_SIZE};
use crate::key_value::{
codec::{decode_u64_value, u64_value},
occurrence_format::BlockMaxValue,
};
impl KeyValueInvertedIndex {
pub fn flush_skip_pointers(&self) -> StorageBackendResult<()> {
self.mutate(|view, batch| view.rebuild_skips(batch))
}
pub fn skip_to(
&self,
field: &str,
term: &str,
target: DocId,
) -> StorageBackendResult<(DocId, usize)> {
self.read(|view| view.skip_to(field, &TokenTermKey::from_text(term), target))
}
pub fn build_block_max_scores_key<S: BlockMaxScorer + ?Sized>(
&self,
field: &str,
term: &TokenTermKey,
scorer: &S,
fingerprint: &str,
) -> StorageBackendResult<()> {
self.mutate(|view, batch| {
view.build_bounds(batch, field, Some(term), scorer, fingerprint, false)
})
}
pub fn build_all_block_max_scores<S: BlockMaxScorer + ?Sized>(
&self,
field: &str,
scorer: &S,
) -> StorageBackendResult<()> {
self.mutate(|view, batch| view.build_bounds(batch, field, None, scorer, "", false))
}
pub(super) fn rebuild_block_max(
&self,
field: &str,
scorer: &dyn BlockMaxScorer,
fingerprint: &str,
) -> StorageBackendResult<bool> {
if fingerprint.is_empty() {
return Err(other_error(
"persisted block-max scorer fingerprint must not be empty",
));
}
self.mutate(|view, batch| {
view.build_bounds(batch, field, None, scorer, fingerprint, true)
})?;
Ok(true)
}
pub fn get_block_max_score(
&self,
field: &str,
term: &str,
block: usize,
) -> StorageBackendResult<f64> {
self.read(|view| {
view.require_graph_format()?;
let block = u64::try_from(block).map_err(|_| other_error("block index exceeds u64"))?;
let key = keys::cluster_key(
view.table,
keys::BLOCK_MAX,
field,
&TokenTermKey::from_text(term),
block,
)?;
view.store
.get(&key)?
.map_or(Ok(0.0), |bytes| Ok(BlockMaxValue::decode(&bytes)?.score))
})
}
pub fn get_all_block_max_scores_key(
&self,
field: &str,
term: &TokenTermKey,
) -> StorageBackendResult<Vec<f64>> {
self.read(|view| Ok(view.bounds(field, term, None)?.unwrap_or_default()))
}
pub fn get_versioned_block_max_scores_keys_bulk(
&self,
field: &str,
terms: &[TokenTermKey],
fingerprint: &str,
) -> StorageBackendResult<Vec<Option<Vec<f64>>>> {
self.read(|view| {
terms
.iter()
.map(|term| view.bounds(field, term, Some(fingerprint)))
.collect()
})
}
pub fn load_block_max_into(&self, target: &mut BlockMaxIndex) -> StorageBackendResult<()> {
self.read(|view| {
view.require_graph_format()?;
let prefix = keys::kind_prefix(view.table, keys::BLOCK_MAX)?;
let mut by_term = BTreeMap::<(String, TokenTermKey), Vec<f64>>::new();
view.store.visit_prefix(&prefix, &mut |key, bytes| {
view.store.control().check()?;
let (field, term, ordinal) = keys::read_cluster(key, keys::BLOCK_MAX)?;
let scores = by_term.entry((field, term)).or_default();
if ordinal != scores.len() as u64 {
return Err(other_error("non-contiguous occurrence block-max bounds"));
}
scores.push(BlockMaxValue::decode(bytes)?.score);
Ok(())
})?;
for ((field, term), scores) in by_term {
target.set_block_maxes_key(view.table, &field, &term, scores)?;
}
Ok(())
})
}
}
impl OccurrenceRead<'_> {
pub(super) fn invalidate_accelerators(
&self,
batch: &mut dyn KeyValueBatch,
) -> StorageBackendResult<()> {
for kind in [keys::SKIP, keys::BLOCK_MAX] {
batch.invalidate_occurrence_prefix(&keys::kind_prefix(self.table, kind)?)?;
}
Ok(())
}
fn fence_accelerators(&self, batch: &mut dyn KeyValueBatch) -> StorageBackendResult<()> {
batch.put(
&keys::kind_prefix(self.table, keys::FORMAT)?,
keys::FORMAT_NAME,
)
}
fn rebuild_skips(&self, batch: &mut dyn KeyValueBatch) -> StorageBackendResult<()> {
self.require_graph_format()?;
batch.delete_prefix(&keys::kind_prefix(self.table, keys::SKIP)?)?;
let fields = self.field_names()?;
for field in &fields {
for term in self.vocabulary_keys(field)? {
let mut cursor = self.posting_cursor_key(field, &term)?;
while let Some(entry) = cursor.current() {
self.store.control().check()?;
let offset = cursor.ordinal();
if offset % DEFAULT_BLOCK_SIZE as u64 == 0 {
batch.put(
&keys::cluster_key(self.table, keys::SKIP, field, &term, entry.doc_id)?,
&u64_value(offset),
)?;
}
cursor.advance()?;
}
}
}
if !fields.is_empty() {
self.fence_accelerators(batch)?;
}
Ok(())
}
fn skip_to(
&self,
field: &str,
term: &TokenTermKey,
target: DocId,
) -> StorageBackendResult<(DocId, usize)> {
self.require_graph_format()?;
let prefix = keys::term_prefix(self.table, keys::SKIP, field, term)?;
let mut found = (0, 0);
self.store.visit_prefix(&prefix, &mut |key, bytes| {
self.store.control().check()?;
let (_, _, doc_id) = keys::read_cluster(key, keys::SKIP)?;
let offset = usize::try_from(decode_u64_value(bytes)?)
.map_err(|_| other_error("skip offset exceeds usize"))?;
if doc_id <= target {
found = (doc_id, offset);
}
Ok(())
})?;
Ok(found)
}
fn build_bounds<S: BlockMaxScorer + ?Sized>(
&self,
batch: &mut dyn KeyValueBatch,
field: &str,
term: Option<&TokenTermKey>,
scorer: &S,
fingerprint: &str,
replace: bool,
) -> StorageBackendResult<()> {
self.require_graph_format()?;
if replace {
batch.delete_prefix(&keys::field_prefix(self.table, keys::BLOCK_MAX, field)?)?;
}
if self.field_doc_count(field)? == 0 {
return Ok(());
}
if let Some(term) = term {
self.build_term_bounds(batch, field, term, scorer, fingerprint)?;
} else {
for term in self.vocabulary_keys(field)? {
self.build_term_bounds(batch, field, &term, scorer, fingerprint)?;
}
}
self.fence_accelerators(batch)
}
fn build_term_bounds<S: BlockMaxScorer + ?Sized>(
&self,
batch: &mut dyn KeyValueBatch,
field: &str,
term: &TokenTermKey,
scorer: &S,
fingerprint: &str,
) -> StorageBackendResult<()> {
batch.delete_prefix(&keys::term_prefix(
self.table,
keys::BLOCK_MAX,
field,
term,
)?)?;
let mut cursor = self.posting_cursor_key(field, term)?;
let mut count = 0usize;
let mut block = 0u64;
let mut maximum = 0.0_f64;
while let Some(entry) = cursor.current() {
self.store.control().check()?;
let score = scorer.score(entry.term_freq, entry.doc_length, cursor.doc_freq());
crate::block_max_index::validate_score(score)?;
maximum = maximum.max(score);
count += 1;
if count == DEFAULT_BLOCK_SIZE {
self.put_bound(batch, field, term, block, maximum, fingerprint)?;
block = block
.checked_add(1)
.ok_or_else(|| other_error("block index overflow"))?;
count = 0;
maximum = 0.0;
}
cursor.advance()?;
}
if count != 0 {
self.put_bound(batch, field, term, block, maximum, fingerprint)?;
}
Ok(())
}
fn put_bound(
&self,
batch: &mut dyn KeyValueBatch,
field: &str,
term: &TokenTermKey,
block: u64,
score: f64,
fingerprint: &str,
) -> StorageBackendResult<()> {
let value = BlockMaxValue { score, fingerprint }.encode(self.store.control())?;
batch.put(
&keys::cluster_key(self.table, keys::BLOCK_MAX, field, term, block)?,
&value,
)
}
fn bounds(
&self,
field: &str,
term: &TokenTermKey,
fingerprint: Option<&str>,
) -> StorageBackendResult<Option<Vec<f64>>> {
self.require_graph_format()?;
let prefix = keys::term_prefix(self.table, keys::BLOCK_MAX, field, term)?;
let mut scores = Vec::new();
let mut matches = true;
self.store.visit_prefix(&prefix, &mut |key, bytes| {
self.store.control().check()?;
let (_, _, block) = keys::read_cluster(key, keys::BLOCK_MAX)?;
if block != scores.len() as u64 {
return Err(other_error("non-contiguous occurrence block-max bounds"));
}
let bound = BlockMaxValue::decode(bytes)?;
matches &= fingerprint.is_none_or(|fingerprint| bound.fingerprint == fingerprint);
scores.push(bound.score);
Ok(())
})?;
Ok((matches && !scores.is_empty()).then_some(scores))
}
}