use super::super::codec::usize_to_u64;
use super::queries::require_score_version;
use super::{
cluster_id, decode_all_scores, keys, other_error, score_count, BTreeMap, DocId, FieldName,
IndexStats, IndexedFieldMetadata, OccurrencePosting, OccurrenceRead, Payload, PostingCursor,
PostingEntry, PostingList, StorageBackendResult, TokenTermKey,
};
#[cfg(test)]
mod tests;
impl OccurrenceRead<'_> {
pub(super) fn visit_score_clusters(
&self,
field: &str,
term: &TokenTermKey,
after: Option<u64>,
limit: usize,
control: &crate::read_control::StorageReadControl,
visit: &mut crate::clustered_postings::ScoreClusterVisitor<'_>,
) -> StorageBackendResult<()> {
self.visit_clusters_budgeted(field, term, after, limit, control, visit)
}
pub(super) fn get_occurrences_budgeted(
&self,
doc_id: DocId,
field: &str,
term: &TokenTermKey,
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<uqa_core::memory::Budgeted<Vec<uqa_core::TokenOccurrence>>> {
self.occurrences_budgeted(doc_id, field, term, control)
}
pub(super) fn field_stats_scalar_budgeted(
&self,
field: &str,
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<IndexStats> {
self.scalar_stats_budgeted(field, control)
}
pub(super) fn source_rebuild_required(&self) -> StorageBackendResult<bool> {
self.needs_source_rebuild()
}
pub(super) fn get_posting_list(
&self,
field: &str,
term: &str,
) -> StorageBackendResult<PostingList> {
self.get_posting_list_key(field, &TokenTermKey::from_text(term))
}
pub(super) fn get_posting_list_key(
&self,
field: &str,
term: &TokenTermKey,
) -> StorageBackendResult<PostingList> {
let entries = self
.occurrence_postings(field, term)?
.into_iter()
.map(|entry| {
PostingEntry::new(
entry.doc_id,
Payload {
positions: entry.positions(),
score: 0.0,
fields: BTreeMap::new(),
},
)
})
.collect();
Ok(PostingList::from_sorted_unchecked(entries))
}
pub(super) fn posting_cursor(
&self,
field: &str,
term: &str,
) -> StorageBackendResult<Box<dyn PostingCursor>> {
self.cursor_for_term(field, &TokenTermKey::from_text(term))
}
pub(super) fn posting_cursor_key(
&self,
field: &str,
term: &TokenTermKey,
) -> StorageBackendResult<Box<dyn PostingCursor>> {
self.cursor_for_term(field, term)
}
pub(super) fn get_occurrence_postings(
&self,
field: &str,
term: &TokenTermKey,
) -> StorageBackendResult<Vec<OccurrencePosting>> {
self.occurrence_postings(field, term)
}
pub(super) fn get_occurrences(
&self,
doc_id: DocId,
field: &str,
term: &TokenTermKey,
) -> StorageBackendResult<Vec<uqa_core::TokenOccurrence>> {
self.require_graph_format()?;
let entries = self.load_cluster(field, term, cluster_id(doc_id))?;
let Some(posting) = entries.into_iter().find(|entry| entry.doc_id == doc_id) else {
return Ok(Vec::new());
};
self.validate_posting_metadata(field, &posting)?;
Ok(posting.occurrences)
}
pub(super) fn indexed_field_metadata(
&self,
doc_id: DocId,
field: &str,
) -> StorageBackendResult<Option<IndexedFieldMetadata>> {
self.require_graph_format()?;
self.read_field_metadata(doc_id, field)
}
pub(super) fn doc_freq(&self, field: &str, term: &str) -> StorageBackendResult<u64> {
self.doc_freq_key(field, &TokenTermKey::from_text(term))
}
pub(super) fn doc_freq_key(
&self,
field: &str,
term: &TokenTermKey,
) -> StorageBackendResult<u64> {
self.require_graph_format()?;
self.scan_prefix(&keys::term_prefix(self.table, keys::SCORE, field, term)?)?
.iter()
.try_fold(0_u64, |total, (key, score)| {
keys::read_cluster(key, keys::SCORE)?;
require_score_version(score)?;
total
.checked_add(score_count(score)?)
.ok_or_else(|| other_error("document frequency overflow"))
})
}
pub(super) fn get_doc_length(&self, doc_id: DocId, field: &str) -> StorageBackendResult<u64> {
self.document_length(doc_id, field)
}
pub(super) fn get_scoring_inputs_bulk(
&self,
doc_ids: &[DocId],
field: &str,
terms: &[String],
) -> StorageBackendResult<Vec<(u64, Vec<u64>)>> {
let keys = terms
.iter()
.map(|term| TokenTermKey::from_text(term))
.collect::<Vec<_>>();
self.get_scoring_inputs_keys_bulk(doc_ids, field, &keys)
}
pub(super) fn get_scoring_inputs_keys_bulk(
&self,
doc_ids: &[DocId],
field: &str,
terms: &[TokenTermKey],
) -> StorageBackendResult<Vec<(u64, Vec<u64>)>> {
let mut output = doc_ids
.iter()
.map(|id| Ok((self.document_length(*id, field)?, vec![0; terms.len()])))
.collect::<StorageBackendResult<Vec<_>>>()?;
let mut positions = BTreeMap::<DocId, Vec<usize>>::new();
for (position, doc_id) in doc_ids.iter().copied().enumerate() {
positions.entry(doc_id).or_default().push(position);
}
for (term_index, term) in terms.iter().enumerate() {
let mut cursor = self.posting_cursor_key(field, term)?;
while let Some(entry) = cursor.current() {
if let Some(output_positions) = positions.get(&entry.doc_id) {
for position in output_positions {
output[*position].0 = entry.doc_length;
output[*position].1[term_index] = entry.term_freq;
}
}
cursor.advance()?;
}
}
Ok(output)
}
pub(super) fn get_term_freq(
&self,
doc_id: DocId,
field: &str,
term: &str,
) -> StorageBackendResult<u64> {
self.get_term_freq_key(doc_id, field, &TokenTermKey::from_text(term))
}
pub(super) fn get_term_freq_key(
&self,
doc_id: DocId,
field: &str,
term: &TokenTermKey,
) -> StorageBackendResult<u64> {
self.require_graph_format()?;
let cluster = cluster_id(doc_id);
self.store
.get(&keys::cluster_key(
self.table,
keys::SCORE,
field,
term,
cluster,
)?)?
.map_or(Ok(0), |score| {
require_score_version(&score)?;
let entries = decode_all_scores(cluster, &score)?;
Ok(entries
.binary_search_by_key(&doc_id, |entry| entry.doc_id)
.ok()
.map_or(0, |position| entries[position].term_freq))
})
}
pub(super) fn doc_count(&self) -> StorageBackendResult<u64> {
self.require_graph_format()?;
let prefix = keys::kind_prefix(self.table, keys::LENGTH)?;
let mut previous = None;
let mut count = 0_u64;
self.store.visit_keys_after(
&prefix,
None,
usize::MAX,
self.store.control(),
&mut |key| {
self.store.control().check()?;
let doc_id = keys::read_document(key, keys::LENGTH)?.0;
if previous != Some(doc_id) {
count = count
.checked_add(1)
.ok_or_else(|| other_error("document count overflow"))?;
previous = Some(doc_id);
}
Ok(())
},
)?;
Ok(count)
}
pub(super) fn total_field_length(&self, field: &str) -> StorageBackendResult<u64> {
self.require_graph_format()?;
Ok(self
.stored_field_stats(field)?
.map_or(0, |stats| stats.total_length))
}
pub(super) fn field_doc_count(&self, field: &str) -> StorageBackendResult<u64> {
self.require_graph_format()?;
Ok(self
.stored_field_stats(field)?
.map_or(0, |stats| stats.doc_count))
}
pub(super) fn vocabulary_terms(&self, field: &str) -> StorageBackendResult<Vec<String>> {
self.vocabulary_keys(field)?
.into_iter()
.map(|key| Ok(key.to_term().into_string()?))
.collect()
}
pub(super) fn vocabulary_keys(&self, field: &str) -> StorageBackendResult<Vec<TokenTermKey>> {
self.indexed_terms(Some(field))
}
pub(super) fn stats(&self) -> StorageBackendResult<IndexStats> {
self.index_statistics()
}
pub(super) fn posting_count(&self, field: Option<&str>) -> StorageBackendResult<u64> {
self.scan_prefix(&self.score_prefix(field)?)?
.iter()
.try_fold(0_u64, |total, (key, value)| {
keys::read_cluster(key, keys::SCORE)?;
require_score_version(value)?;
total
.checked_add(score_count(value)?)
.ok_or_else(|| other_error("posting count overflow"))
})
}
pub(super) fn doc_length_count(&self, field: Option<&str>) -> StorageBackendResult<u64> {
self.require_graph_format()?;
if let Some(field) = field {
return self.field_doc_count(field);
}
let mut count = 0_u64;
for (key, _) in self
.scan_prefix(&keys::kind_prefix(self.table, keys::LENGTH)?)?
.iter()
{
keys::read_document(key, keys::LENGTH)?;
count = count
.checked_add(1)
.ok_or_else(|| other_error("document-length row count overflow"))?;
}
Ok(count)
}
pub(super) fn term_count(&self, field: Option<&str>) -> StorageBackendResult<u64> {
usize_to_u64(self.indexed_terms(field)?.len(), "term count")
}
pub(super) fn field_names(&self) -> StorageBackendResult<Vec<FieldName>> {
self.require_graph_format()?;
let mut fields = Vec::new();
for (key, value) in self
.scan_prefix(&keys::kind_prefix(self.table, keys::FIELD)?)?
.iter()
{
super::FieldStats::from_bytes(value)?;
fields.push(keys::read_field(key)?);
}
fields.sort();
Ok(fields)
}
}