uqa_scoring/text/
candidates.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Query-local scoring after a positional or other support predicate accepts a candidate.
8
9use super::TextSearchError;
10use crate::{BM25Params, BayesianBM25Params, ScoringMode};
11use uqa_core::{
12    memory::{BudgetedVec, MemoryBudget},
13    IndexStats,
14};
15
16enum CandidateMode {
17    BM25(BM25Params),
18    Bayesian(BayesianBM25Params),
19}
20impl CandidateMode {
21    fn params(&self) -> &BM25Params {
22        match self {
23            Self::BM25(params) => params,
24            Self::Bayesian(params) => &params.bm25,
25        }
26    }
27    fn finalize(&self, sum: f64) -> f64 {
28        match self {
29            Self::BM25(_) => sum,
30            Self::Bayesian(scorer) => scorer.calibrate_raw_value(sum),
31        }
32    }
33}
34
35/// Scores aligned emitted query terms, including duplicates, after support acceptance.
36pub struct TextCandidateScorer {
37    scorer: CandidateMode,
38    avg_doc_length: f64,
39    idfs: BudgetedVec<f64>,
40}
41
42impl TextCandidateScorer {
43    pub fn new(
44        mode: &ScoringMode,
45        stats: IndexStats,
46        document_frequencies: &[u64],
47    ) -> Result<Self, TextSearchError> {
48        Self::new_budgeted(
49            mode,
50            stats,
51            document_frequencies,
52            &MemoryBudget::new(usize::MAX),
53            || Ok(()),
54        )
55    }
56
57    /// Reserve query IDFs from the caller's shared runtime allowance.
58    ///
59    /// Frequencies arrive in emitted query order. Corpus scalars and scoring parameters stay inline, and unused index statistics are dropped. The callback covers construction and IDF preparation; partial owners are released on failure.
60    pub fn new_budgeted(
61        mode: &ScoringMode,
62        stats: IndexStats,
63        document_frequencies: &[u64],
64        budget: &MemoryBudget,
65        mut poll: impl FnMut() -> Result<(), TextSearchError>,
66    ) -> Result<Self, TextSearchError> {
67        poll()?;
68        let total_docs = stats.total_docs;
69        let avg_doc_length = stats.avg_doc_length;
70        drop(stats);
71        let scorer = match mode {
72            ScoringMode::BM25(params) => {
73                params.validate()?;
74                CandidateMode::BM25(*params)
75            }
76            ScoringMode::BayesianBM25(params) => {
77                let params = params.scaled_for_query_terms(document_frequencies.len());
78                crate::bayesian_bm25::validate_params(params, avg_doc_length)?;
79                CandidateMode::Bayesian(params)
80            }
81        };
82        let mut output = Self {
83            scorer,
84            avg_doc_length,
85            idfs: BudgetedVec::new(budget),
86        };
87        output.idfs.reserve(document_frequencies.len())?;
88        for frequency in document_frequencies {
89            poll()?;
90            output.idfs.push(crate::bm25::idf(total_docs, *frequency))?;
91        }
92        poll()?;
93        Ok(output)
94    }
95
96    pub fn score_document(
97        &mut self,
98        document_length: u64,
99        term_frequencies: &[u64],
100    ) -> Result<f64, TextSearchError> {
101        self.score_document_with_control(document_length, term_frequencies, || Ok(()))
102    }
103
104    /// Sum native term contributions in emitted order and apply calibration once, without a temporary score vector.
105    pub fn score_document_with_control(
106        &self,
107        document_length: u64,
108        term_frequencies: &[u64],
109        mut poll: impl FnMut() -> Result<(), TextSearchError>,
110    ) -> Result<f64, TextSearchError> {
111        poll()?;
112        if term_frequencies.len() != self.idfs.len() {
113            return Err(TextSearchError::InvalidIndex(
114                "candidate frequencies do not match the emitted query term count".into(),
115            ));
116        }
117        // The scalar sum uses the same initial value and left-to-right order as Iterator::sum.
118        let mut sum = -0.0;
119        for (idf, frequency) in self.idfs.iter().zip(term_frequencies) {
120            poll()?;
121            sum += self.scorer.params().score_with_idf(
122                self.avg_doc_length,
123                *frequency,
124                document_length,
125                *idf,
126            );
127        }
128        poll()?;
129        Ok(self.scorer.finalize(sum))
130    }
131}
132
133#[cfg(test)]
134mod tests;
uqa_scoring/text/candidates.rs

uqa_scoring/text/
candidates.rs
