uqa_storage/inverted_index/

memory.rs

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Memory provider mutations, projections, and exact analyzer bindings.

use super::{
    checked_sum_u64, counter_error, usize_to_u64, Analyzer, AnalyzerBindings, AnalyzerPhase, Arc,
    BTreeMap, BTreeSet, DocId, FieldName, IndexStats, IndexedFieldMetadata, InvertedIndex,
    MaterializedPostingCursor, MemoryInvertedIndex, PostingCursor, PostingEntry, PostingList,
    PostingScore, StorageBackendError, StorageBackendResult, TokenOccurrence, TokenTermKey,
};

impl InvertedIndex for MemoryInvertedIndex {
    fn field_stats_scalar_budgeted(
        &self,
        field: &str,
        control: &crate::read_control::StorageReadControl,
    ) -> StorageBackendResult<IndexStats> {
        control.check()?;
        let stats = self.field_stats_scalar(field)?;
        control.check()?;
        Ok(stats)
    }

    fn posting_read_cursor_key_budgeted<'a>(
        &'a self,
        field: &'a str,
        term: &'a TokenTermKey,
        control: &crate::read_control::StorageReadControl,
    ) -> StorageBackendResult<crate::clustered_postings::BudgetedPostingReadCursor<'a>> {
        let cursor =
            super::read_cursor::MemoryPostingReadCursor::with_control(self, field, term, control)?;
        crate::clustered_postings::BudgetedPostingReadCursor::new(cursor, control)
    }

    fn get_occurrences_budgeted(
        &self,
        doc_id: DocId,
        field: &str,
        term: &TokenTermKey,
        control: &crate::read_control::StorageReadControl,
    ) -> StorageBackendResult<uqa_core::memory::Budgeted<Vec<TokenOccurrence>>> {
        let postings = super::read_cursor::controlled_postings(self, field, term, control)?;
        let mut output = uqa_core::memory::BudgetedVec::new(control.memory());
        if let Some(posting) = postings.and_then(|postings| postings.get(&doc_id)) {
            output.reserve(posting.occurrences.len())?;
            for occurrence in &posting.occurrences {
                control.check()?;
                output.push(*occurrence)?;
            }
        }
        control.check()?;
        let (values, memory) = output.into_parts();
        Ok(uqa_core::memory::Budgeted::new(values, memory))
    }

    fn posting_read_cursor_key<'a>(
        &'a self,
        field: &'a str,
        term: &TokenTermKey,
    ) -> StorageBackendResult<Box<dyn crate::clustered_postings::PostingReadCursor + 'a>> {
        Ok(Box::new(super::read_cursor::MemoryPostingReadCursor::new(
            self, field, term,
        )?))
    }

    fn analyzer(&self) -> &Analyzer {
        self.bindings.default_configuration()
    }

    fn add_document(
        &mut self,
        doc_id: DocId,
        fields: BTreeMap<FieldName, String>,
    ) -> StorageBackendResult<()> {
        // Resolve and analyze every field before touching postings. A deferred default can fail even when another field already has a valid revision.
        let staged = self.stage_document(doc_id, fields)?;
        let plan = self.state.plan_replacement(doc_id, &staged.fields)?;
        Arc::make_mut(&mut self.state).apply_replacement(doc_id, staged, plan)
    }

    fn remove_document(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
        let Some(keys) = self.state.doc_terms.get(&doc_id).cloned() else {
            if self.state.doc_fields.contains_key(&doc_id) {
                return Err(StorageBackendError::Other(format!(
                    "inverted-index document {doc_id} has lengths but no reverse postings"
                )));
            }
            return Ok(());
        };
        let lengths = self.state.doc_fields.get(&doc_id).cloned().ok_or_else(|| {
            StorageBackendError::Other(format!(
                "inverted-index document {doc_id} has reverse postings but no lengths"
            ))
        })?;
        let next_doc_count = self
            .state
            .doc_count
            .checked_sub(1)
            .ok_or_else(|| counter_error("document count"))?;
        for key in &keys {
            if !self
                .state
                .index
                .get(key)
                .is_some_and(|postings| postings.contains_key(&doc_id))
            {
                return Err(StorageBackendError::Other(format!(
                    "inverted-index document {doc_id} references a missing posting"
                )));
            }
        }
        let mut next_field_counters = BTreeMap::new();
        for (field, metadata) in &lengths {
            let total = self
                .state
                .total_length
                .get(field)
                .copied()
                .unwrap_or(0)
                .checked_sub(metadata.length)
                .ok_or_else(|| counter_error("total field length"))?;
            let field_docs = self
                .state
                .field_doc_counts
                .get(field)
                .copied()
                .unwrap_or(0)
                .checked_sub(1)
                .ok_or_else(|| counter_error("field document count"))?;
            next_field_counters.insert(field.clone(), (total, field_docs));
        }

        let state = Arc::make_mut(&mut self.state);
        for key in keys {
            let inner = state.index.get_mut(&key).ok_or_else(|| {
                StorageBackendError::Other(format!(
                    "inverted-index document {doc_id} lost a validated posting before removal"
                ))
            })?;
            inner.remove(&doc_id);
            if inner.is_empty() {
                state.index.remove(&key);
            }
        }
        state.doc_terms.remove(&doc_id);
        state.doc_fields.remove(&doc_id);
        for (field, (total, field_docs)) in next_field_counters {
            if field_docs == 0 {
                state.total_length.remove(&field);
                state.field_doc_counts.remove(&field);
            } else {
                state.total_length.insert(field.clone(), total);
                state.field_doc_counts.insert(field, field_docs);
            }
        }
        state.doc_count = next_doc_count;
        Ok(())
    }

    fn try_rebuild_documents(
        &mut self,
        documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
    ) -> StorageBackendResult<()> {
        self.rebuild_documents_inner(documents, None)
    }

    fn try_rebuild_documents_cancellable(
        &mut self,
        documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
        cancellation: &uqa_core::CancellationToken,
    ) -> StorageBackendResult<()> {
        self.rebuild_documents_inner(documents, Some(cancellation))
    }

    fn try_add_documents(
        &mut self,
        documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
    ) -> StorageBackendResult<()> {
        self.add_document_batch(documents)
    }

    fn clear(&mut self) -> StorageBackendResult<()> {
        // Clearing a shared snapshot must not copy the state it is discarding.
        if let Some(state) = Arc::get_mut(&mut self.state) {
            *state = super::MemoryIndexState::default();
        } else {
            self.state = Arc::default();
        }
        Ok(())
    }

    fn get_posting_list(&self, field: &str, term: &str) -> StorageBackendResult<PostingList> {
        self.get_posting_list_key(field, &TokenTermKey::from_text(term))
    }

    fn get_posting_list_key(
        &self,
        field: &str,
        term: &TokenTermKey,
    ) -> StorageBackendResult<PostingList> {
        let entries = self
            .state
            .index
            .get(&(field.to_owned(), term.clone()))
            .into_iter()
            .flat_map(|postings| postings.values())
            .map(|posting| posting.projection.clone())
            .collect();
        Ok(PostingList::from_sorted_unchecked(entries))
    }

    fn posting_cursor(
        &self,
        field: &str,
        term: &str,
    ) -> StorageBackendResult<Box<dyn PostingCursor>> {
        self.posting_cursor_key(field, &TokenTermKey::from_text(term))
    }

    fn posting_cursor_key(
        &self,
        field: &str,
        term: &TokenTermKey,
    ) -> StorageBackendResult<Box<dyn PostingCursor>> {
        let entries = self
            .state
            .index
            .get(&(field.to_owned(), term.clone()))
            .into_iter()
            .flat_map(|postings| postings.values())
            .map(|posting| {
                Ok(PostingScore {
                    doc_id: posting.projection.doc_id,
                    term_freq: usize_to_u64(posting.occurrences.len(), "term frequency")?,
                    doc_length: self.get_doc_length(posting.projection.doc_id, field)?,
                })
            })
            .collect::<StorageBackendResult<Vec<_>>>()?;
        Ok(Box::new(MaterializedPostingCursor::new(entries)?))
    }

    fn get_occurrence_postings(
        &self,
        field: &str,
        term: &TokenTermKey,
    ) -> StorageBackendResult<Vec<crate::clustered_postings::OccurrencePosting>> {
        self.state
            .index
            .get(&(field.to_owned(), term.clone()))
            .into_iter()
            .flat_map(|postings| postings.values())
            .map(|posting| {
                Ok(crate::clustered_postings::OccurrencePosting {
                    doc_id: posting.projection.doc_id,
                    doc_length: self.get_doc_length(posting.projection.doc_id, field)?,
                    occurrences: posting.occurrences.clone(),
                })
            })
            .collect()
    }

    fn get_occurrences(
        &self,
        doc_id: DocId,
        field: &str,
        term: &TokenTermKey,
    ) -> StorageBackendResult<Vec<TokenOccurrence>> {
        Ok(self
            .state
            .index
            .get(&(field.to_owned(), term.clone()))
            .and_then(|postings| postings.get(&doc_id))
            .map_or_else(Vec::new, |posting| posting.occurrences.clone()))
    }

    fn indexed_field_metadata(
        &self,
        doc_id: DocId,
        field: &str,
    ) -> StorageBackendResult<Option<IndexedFieldMetadata>> {
        Ok(self
            .state
            .doc_fields
            .get(&doc_id)
            .and_then(|fields| fields.get(field))
            .copied())
    }

    fn for_each_posting(
        &self,
        field: &str,
        term: &str,
        visit: &mut dyn FnMut(&PostingEntry),
    ) -> StorageBackendResult<()> {
        if let Some(postings) = self
            .state
            .index
            .get(&(field.to_owned(), TokenTermKey::from_text(term)))
        {
            for posting in postings.values() {
                visit(&posting.projection);
            }
        }
        Ok(())
    }

    fn for_each_term_freq(
        &self,
        field: &str,
        term: &str,
        visit: &mut dyn FnMut(DocId, u64),
    ) -> StorageBackendResult<()> {
        if let Some(postings) = self
            .state
            .index
            .get(&(field.to_owned(), TokenTermKey::from_text(term)))
        {
            for posting in postings.values() {
                visit(
                    posting.projection.doc_id,
                    usize_to_u64(posting.occurrences.len(), "term frequency")?,
                );
            }
        }
        Ok(())
    }

    fn doc_freq(&self, field: &str, term: &str) -> StorageBackendResult<u64> {
        self.doc_freq_key(field, &TokenTermKey::from_text(term))
    }

    fn doc_freq_key(&self, field: &str, term: &TokenTermKey) -> StorageBackendResult<u64> {
        self.state
            .index
            .get(&(field.to_owned(), term.clone()))
            .map_or(Ok(0), |postings| {
                usize_to_u64(postings.len(), "document frequency")
            })
    }

    fn get_doc_length(&self, doc_id: DocId, field: &str) -> StorageBackendResult<u64> {
        Ok(self
            .state
            .doc_fields
            .get(&doc_id)
            .and_then(|fields| fields.get(field))
            .map_or(0, |metadata| metadata.length))
    }

    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))
    }

    fn get_term_freq_key(
        &self,
        doc_id: DocId,
        field: &str,
        term: &TokenTermKey,
    ) -> StorageBackendResult<u64> {
        self.state
            .index
            .get(&(field.to_owned(), term.clone()))
            .and_then(|postings| postings.get(&doc_id))
            .map_or(Ok(0), |posting| {
                usize_to_u64(posting.occurrences.len(), "term frequency")
            })
    }

    fn doc_count(&self) -> StorageBackendResult<u64> {
        Ok(self.state.doc_count)
    }

    fn total_field_length(&self, field: &str) -> StorageBackendResult<u64> {
        Ok(self.state.total_length.get(field).copied().unwrap_or(0))
    }

    fn vocabulary_terms(&self, field: &str) -> StorageBackendResult<Vec<String>> {
        self.vocabulary_keys(field)?
            .into_iter()
            .map(|key| Ok(key.to_term().into_string()?))
            .collect()
    }

    fn vocabulary_keys(&self, field: &str) -> StorageBackendResult<Vec<TokenTermKey>> {
        Ok(self
            .state
            .index
            .keys()
            .filter(|(indexed_field, _)| indexed_field == field)
            .map(|(_, term)| term.clone())
            .collect())
    }

    fn stats(&self) -> StorageBackendResult<IndexStats> {
        let mut s = IndexStats::default();
        s.total_docs = self.state.doc_count;
        if self.state.doc_count > 0 {
            let total = checked_sum_u64(
                self.state.total_length.values().copied(),
                "total document length",
            )?;
            s.avg_doc_length = total as f64 / self.state.doc_count as f64;
        }
        for ((field, term), inner) in &self.state.index {
            let term = term.to_term();
            let frequency = usize_to_u64(inner.len(), "document frequency")?;
            if let Some(text) = term.as_str() {
                s.set_doc_freq(field.clone(), text, frequency);
            } else {
                s.set_doc_freq_utf16(field.clone(), term.into_utf16(), frequency);
            }
        }
        Ok(s)
    }

    fn posting_count(&self, field: Option<&str>) -> StorageBackendResult<u64> {
        checked_sum_u64(
            self.state
                .index
                .iter()
                .filter(|((f, _), _)| field.is_none_or(|target| f == target))
                .map(|(_, postings)| usize_to_u64(postings.len(), "posting count"))
                .collect::<StorageBackendResult<Vec<_>>>()?,
            "posting count",
        )
    }

    fn doc_length_count(&self, field: Option<&str>) -> StorageBackendResult<u64> {
        Ok(match field {
            Some(target) => self
                .state
                .field_doc_counts
                .get(target)
                .copied()
                .unwrap_or(0),
            None => checked_sum_u64(
                self.state.field_doc_counts.values().copied(),
                "document-length row count",
            )?,
        })
    }

    fn term_count(&self, field: Option<&str>) -> StorageBackendResult<u64> {
        usize_to_u64(
            self.state
                .index
                .keys()
                .filter(|(f, _)| field.is_none_or(|target| f == target))
                .map(|(_, term)| term)
                .collect::<BTreeSet<_>>()
                .len(),
            "term count",
        )
    }

    fn snapshot(&self) -> StorageBackendResult<Arc<dyn InvertedIndex>> {
        Ok(Arc::new(self.shared_snapshot()))
    }

    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn InvertedIndex>> {
        Ok(Box::new(self.shared_snapshot()))
    }

    fn field_names(&self) -> StorageBackendResult<Vec<FieldName>> {
        Ok(self.state.total_length.keys().cloned().collect())
    }

    fn set_field_analyzer(
        &mut self,
        field: &str,
        analyzer: Analyzer,
        phase: AnalyzerPhase,
    ) -> Result<(), String> {
        let mut candidate = self.bindings.clone();
        candidate
            .bind(field, &analyzer, phase)
            .map_err(|error| error.to_string())?;
        self.validate_index_revision_change(field, &candidate)?;
        self.bindings = candidate;
        Ok(())
    }

    fn remove_field_analyzers(&mut self, field: &str) -> Result<(), String> {
        let mut candidate = self.bindings.clone();
        candidate.remove(field);
        self.validate_index_revision_change(field, &candidate)?;
        self.bindings = candidate;
        Ok(())
    }

    fn get_field_analyzer(&self, field: &str) -> Analyzer {
        self.bindings.index_configuration(field).clone()
    }

    fn get_search_analyzer(&self, field: &str) -> Analyzer {
        self.bindings.search_configuration(field).clone()
    }

    fn index_analyzer_revision(
        &self,
        field: &str,
    ) -> StorageBackendResult<Arc<uqa_analysis::CompiledAnalyzer>> {
        Ok(self.bindings.index_revision(field)?)
    }

    fn search_analyzer_revision(
        &self,
        field: &str,
    ) -> StorageBackendResult<Arc<uqa_analysis::CompiledAnalyzer>> {
        Ok(self.bindings.search_revision(field)?)
    }

    fn set_field_analyzer_revision(
        &mut self,
        field: &str,
        revision: Arc<uqa_analysis::CompiledAnalyzer>,
        phase: AnalyzerPhase,
    ) -> Result<(), String> {
        let mut candidate = self.bindings.clone();
        candidate
            .bind_revision(field, revision, phase)
            .map_err(|error| error.to_string())?;
        self.validate_index_revision_change(field, &candidate)?;
        self.bindings = candidate;
        Ok(())
    }

    fn set_field_analyzer_revisions(
        &mut self,
        field: &str,
        index: Arc<uqa_analysis::CompiledAnalyzer>,
        search: Arc<uqa_analysis::CompiledAnalyzer>,
    ) -> Result<(), String> {
        let mut candidate = self.bindings.clone();
        candidate
            .bind_revisions(field, index, search)
            .map_err(|error| error.to_string())?;
        self.validate_index_revision_change(field, &candidate)?;
        self.bindings = candidate;
        Ok(())
    }

    fn rebuild_with_analyzer_revision(
        &mut self,
        field: &str,
        revision: Arc<uqa_analysis::CompiledAnalyzer>,
        phase: AnalyzerPhase,
        documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
    ) -> StorageBackendResult<()> {
        let mut replacement = Self::with_bindings(self.bindings.clone());
        replacement
            .set_field_analyzer_revision(field, revision, phase)
            .map_err(crate::StorageBackendError::Other)?;
        replacement.try_rebuild_documents(documents)?;
        *self = replacement;
        Ok(())
    }

    fn rebuild_with_analyzer_revision_cancellable(
        &mut self,
        field: &str,
        revision: Arc<uqa_analysis::CompiledAnalyzer>,
        phase: AnalyzerPhase,
        documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
        cancellation: &uqa_core::CancellationToken,
    ) -> StorageBackendResult<()> {
        cancellation.check()?;
        let mut replacement = Self::with_bindings(self.bindings.clone());
        replacement
            .set_field_analyzer_revision(field, revision, phase)
            .map_err(crate::StorageBackendError::Other)?;
        replacement.rebuild_documents_inner(documents, Some(cancellation))?;
        *self = replacement;
        Ok(())
    }
}

impl MemoryInvertedIndex {
    fn rebuild_documents_inner(
        &mut self,
        documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
        cancellation: Option<&uqa_core::CancellationToken>,
    ) -> StorageBackendResult<()> {
        let mut replacement = Self::with_bindings(self.bindings.clone());
        for (doc_id, fields) in documents {
            if let Some(cancellation) = cancellation {
                cancellation.check()?;
            }
            if !fields.is_empty() {
                let staged = replacement.stage_document_inner(doc_id, fields, cancellation)?;
                let plan = replacement.state.plan_replacement(doc_id, &staged.fields)?;
                Arc::make_mut(&mut replacement.state).apply_replacement(doc_id, staged, plan)?;
            }
        }
        if let Some(cancellation) = cancellation {
            cancellation.check()?;
        }
        *self = replacement;
        Ok(())
    }

    fn validate_index_revision_change(
        &self,
        field: &str,
        candidate: &AnalyzerBindings,
    ) -> Result<(), String> {
        if self.state.field_doc_counts.get(field).copied().unwrap_or(0) > 0 {
            let current = self
                .bindings
                .index_revision(field)
                .map_err(|error| error.to_string())?;
            let proposed = candidate
                .index_revision(field)
                .map_err(|error| error.to_string())?;
            if current.descriptor().fingerprint() != proposed.descriptor().fingerprint() {
                return Err(format!("field `{field}` has indexed documents; changing its index analyzer requires an atomic source rebuild"));
            }
        }
        Ok(())
    }
}