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

uqa_scoring/text/
candidates.rs
