use super::{
decode_occurrence_cluster, decode_term_keys, decode_u64_value, keys, other_error,
AnalyzerBindings, BTreeMap, DocId, DocumentFields, FieldName, FieldSnapshot, FieldStats,
IndexedFieldMetadata, IndexedFieldRevision, OccurrencePosting, OccurrenceRead, StagedDocuments,
StorageBackendResult, TokenTermKey,
};
#[cfg(test)]
mod tests;
pub(super) type Revisions = BTreeMap<FieldName, std::sync::Arc<uqa_analysis::CompiledAnalyzer>>;
pub(super) fn analyze_document(
bindings: &AnalyzerBindings,
control: &crate::read_control::StorageReadControl,
fields: BTreeMap<FieldName, String>,
revisions: &mut Revisions,
cancellation: Option<&uqa_core::CancellationToken>,
) -> StorageBackendResult<DocumentFields> {
control.check()?;
if let Some(cancellation) = cancellation {
cancellation.check()?;
}
let mut snapshot = DocumentFields::new();
for (field, text) in fields {
if !revisions.contains_key(&field) {
revisions.insert(field.clone(), bindings.index_revision(&field)?);
}
let revision = &revisions[&field];
let analyzed = crate::inverted_index::analyze_index_field_with_scratch(
revision,
&text,
control.memory(),
|| {
control
.cancellation()
.check()
.and_then(|()| cancellation.map_or(Ok(()), uqa_core::CancellationToken::check))
.map_err(|_| uqa_analysis::AnalysisError::Cancelled)
},
)
.map_err(|error| match error {
crate::StorageBackendError::Analysis(uqa_analysis::AnalysisError::Memory(memory)) => {
crate::StorageBackendError::Memory(memory)
}
crate::StorageBackendError::Analysis(uqa_analysis::AnalysisError::Cancelled) => {
crate::StorageBackendError::Cancelled(uqa_core::QueryCancelled::USER_REQUEST)
}
other => other,
})?;
let metadata = IndexedFieldMetadata::new(revision, &analyzed);
snapshot.insert(
field,
FieldSnapshot {
metadata,
terms: analyzed.terms,
},
);
}
Ok(snapshot)
}
impl FieldStats {
pub(in crate::key_value) fn to_bytes(self) -> StorageBackendResult<[u8; 56]> {
if self.doc_count == 0 {
return Err(other_error("empty field statistics must be removed"));
}
let mut bytes = [0; 56];
bytes[..40].copy_from_slice(&self.revision.to_bytes()?);
bytes[40..48].copy_from_slice(&self.doc_count.to_le_bytes());
bytes[48..].copy_from_slice(&self.total_length.to_le_bytes());
Ok(bytes)
}
pub(in crate::key_value) fn from_bytes(bytes: &[u8]) -> StorageBackendResult<Self> {
if bytes.len() != 56 {
return Err(other_error("invalid occurrence field statistics"));
}
let stats = Self {
revision: IndexedFieldRevision::from_bytes(&bytes[..40])?,
doc_count: u64::from_le_bytes(
bytes[40..48]
.try_into()
.map_err(|_| other_error("invalid field document count"))?,
),
total_length: u64::from_le_bytes(
bytes[48..]
.try_into()
.map_err(|_| other_error("invalid total field length"))?,
),
};
stats.to_bytes()?;
Ok(stats)
}
}
impl OccurrenceRead<'_> {
pub(super) fn stored_field_stats(
&self,
field: &str,
) -> StorageBackendResult<Option<FieldStats>> {
self.store
.get(&keys::field_prefix(self.table, keys::FIELD, field)?)?
.map(|bytes| FieldStats::from_bytes(&bytes))
.transpose()
}
pub(super) fn validate_index_revision_change(
&self,
field: &str,
candidate: &AnalyzerBindings,
) -> StorageBackendResult<()> {
if self.needs_source_rebuild()? {
return Ok(());
}
if let Some(stored) = self.stored_field_stats(field)? {
let revision = candidate.index_revision(field)?;
if stored.revision != IndexedFieldRevision::new(&revision) {
return Err(other_error(
"changing a populated field's index revision requires an atomic source rebuild",
));
}
}
Ok(())
}
pub(super) fn stage_documents(
&self,
documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
rebuilding: bool,
) -> StorageBackendResult<StagedDocuments> {
self.stage_documents_inner(documents, rebuilding, None)
}
pub(super) fn stage_documents_inner(
&self,
documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
rebuilding: bool,
cancellation: Option<&uqa_core::CancellationToken>,
) -> StorageBackendResult<StagedDocuments> {
let mut staged = BTreeMap::new();
let mut revisions = BTreeMap::new();
for (doc_id, fields) in documents {
let snapshot = self.stage_document(fields, &mut revisions, rebuilding, cancellation)?;
staged.insert(doc_id, snapshot);
}
Ok(staged)
}
pub(super) fn stage_document(
&self,
fields: BTreeMap<FieldName, String>,
revisions: &mut Revisions,
rebuilding: bool,
cancellation: Option<&uqa_core::CancellationToken>,
) -> StorageBackendResult<DocumentFields> {
if !rebuilding {
for field in fields.keys() {
if !revisions.contains_key(field) {
self.validate_index_revision_change(field, self.bindings)?;
}
}
}
analyze_document(
self.bindings,
self.store.control(),
fields,
revisions,
cancellation,
)
}
pub(super) fn old_document(&self, doc_id: DocId) -> StorageBackendResult<DocumentFields> {
let mut fields = DocumentFields::new();
for (key, value) in self
.scan_prefix(&keys::document_prefix(self.table, keys::LENGTH, doc_id)?)?
.iter()
{
let (_, field) = keys::read_document(key, keys::LENGTH)?;
let metadata = self
.read_field_metadata(doc_id, &field)?
.ok_or_else(|| other_error("indexed field end metadata is missing"))?;
if metadata.length != decode_u64_value(value)? {
return Err(other_error(
"indexed field length disagrees with its source metadata",
));
}
let reverse = self
.store
.get(&keys::document_key(
self.table,
keys::DOCUMENT,
doc_id,
&field,
)?)?
.ok_or_else(|| other_error("indexed field reverse terms are missing"))?;
fields.insert(
field,
FieldSnapshot {
metadata,
terms: decode_term_keys(&reverse)?
.into_iter()
.map(|term| (term, Vec::new()))
.collect(),
},
);
}
Ok(fields)
}
pub(super) fn read_field_metadata(
&self,
doc_id: DocId,
field: &str,
) -> StorageBackendResult<Option<IndexedFieldMetadata>> {
let Some(bytes) = self
.store
.get(&keys::metadata_key(self.table, field, doc_id)?)?
else {
return Ok(None);
};
let metadata = IndexedFieldMetadata::from_bytes(&bytes)?;
let stats = self
.stored_field_stats(field)?
.ok_or_else(|| other_error("indexed field revision is missing"))?;
if stats.revision != metadata.revision() {
return Err(other_error("indexed field revisions disagree"));
}
Ok(Some(metadata))
}
pub(super) fn load_cluster(
&self,
field: &str,
term: &TokenTermKey,
cluster: u64,
) -> StorageBackendResult<Vec<OccurrencePosting>> {
let score = self.store.get(&keys::cluster_key(
self.table,
keys::SCORE,
field,
term,
cluster,
)?)?;
let positions = self.store.get(&keys::cluster_key(
self.table,
keys::POSITIONS,
field,
term,
cluster,
)?)?;
match (score, positions) {
(None, None) => Ok(Vec::new()),
(Some(score), Some(positions)) => {
decode_occurrence_cluster(cluster, &score, &positions)
}
_ => Err(other_error("occurrence score and graph payloads disagree")),
}
}
}