uqa_scoring/wand/
cursor.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Score-only posting cursor state and exact WAND/BMW pivot loops.
8
9use std::collections::BinaryHeap;
10use std::sync::Arc;
11
12use uqa_core::{DocId, FieldName, Payload, PostingEntry, PostingList};
13use uqa_storage::{BlockMaxIndex, PostingCursor, StorageBackendResult};
14
15use crate::scorer::Scorer;
16
17use super::common::{
18    invalid_wand_input, require_nonnegative_finite, update_top_k, WANDResult, WANDStats, INF_DOC,
19};
20
21/// WAND query backed directly by score-only posting cursors.
22///
23/// Unlike [`super::materialized::WANDQuery`], this form never materializes positional payloads and
24/// carries document length beside term frequency in each cursor entry.
25pub struct CursorWANDQuery {
26    pub cursors: Vec<Box<dyn PostingCursor>>,
27    pub scorers: Vec<Arc<dyn Scorer>>,
28    pub fields: Vec<FieldName>,
29    pub terms: Vec<String>,
30    pub k: usize,
31}
32
33impl CursorWANDQuery {
34    pub fn new(
35        cursors: Vec<Box<dyn PostingCursor>>,
36        scorers: Vec<Arc<dyn Scorer>>,
37        fields: Vec<FieldName>,
38        terms: Vec<String>,
39        k: usize,
40    ) -> StorageBackendResult<Self> {
41        let expected = cursors.len();
42        if scorers.len() != expected || fields.len() != expected || terms.len() != expected {
43            return Err(invalid_wand_input(format!(
44                "cursor WAND term arrays must have equal lengths: cursors={expected}, scorers={}, fields={}, terms={}",
45                scorers.len(),
46                fields.len(),
47                terms.len()
48            )));
49        }
50        Ok(Self {
51            cursors,
52            scorers,
53            fields,
54            terms,
55            k,
56        })
57    }
58}
59
60struct ScoreTermCursor {
61    cursor: Box<dyn PostingCursor>,
62    upper_bound: f64,
63}
64
65impl ScoreTermCursor {
66    fn current_doc(&self) -> DocId {
67        self.cursor.current().map_or(INF_DOC, |entry| entry.doc_id)
68    }
69
70    fn block_ordinal(&self) -> StorageBackendResult<usize> {
71        usize::try_from(self.cursor.ordinal())
72            .map_err(|_| invalid_wand_input("posting cursor ordinal does not fit in usize"))
73    }
74}
75
76/// Standard WAND over score-only posting cursors.
77pub struct CursorWANDScorer<'a> {
78    query: &'a CursorWANDQuery,
79}
80
81impl<'a> CursorWANDScorer<'a> {
82    pub fn new(query: &'a CursorWANDQuery) -> Self {
83        Self { query }
84    }
85
86    pub fn score_top_k(&self) -> StorageBackendResult<WANDResult> {
87        validate_cursor_query(self.query)?;
88        let mut cursors = build_score_cursors(self.query)?;
89        run_cursor_pivot_loop(self.query, &mut cursors, |_, _, _| Ok(false))
90    }
91}
92
93/// Block-Max WAND over score-only posting cursors.
94pub struct CursorBlockMaxWANDScorer<'a> {
95    query: &'a CursorWANDQuery,
96    block_max_index: &'a BlockMaxIndex,
97    table: String,
98}
99
100impl<'a> CursorBlockMaxWANDScorer<'a> {
101    pub fn new(
102        query: &'a CursorWANDQuery,
103        block_max_index: &'a BlockMaxIndex,
104        table: impl Into<String>,
105    ) -> Self {
106        Self {
107            query,
108            block_max_index,
109            table: table.into(),
110        }
111    }
112
113    pub fn score_top_k(&self) -> StorageBackendResult<WANDResult> {
114        validate_cursor_query(self.query)?;
115        let mut cursors = build_score_cursors(self.query)?;
116        let query = self.query;
117        let block_max = self.block_max_index;
118        let suffix_bounds = query
119            .fields
120            .iter()
121            .zip(&query.terms)
122            .map(|(field, term)| {
123                let Some(blocks) = block_max.block_maxes(&self.table, field, term) else {
124                    return Vec::new();
125                };
126                let mut suffix = vec![0.0_f64; blocks.len()];
127                let mut maximum = 0.0_f64;
128                for (index, score) in blocks.iter().enumerate().rev() {
129                    maximum = maximum.max(*score);
130                    suffix[index] = maximum;
131                }
132                suffix
133            })
134            .collect::<Vec<_>>();
135        run_cursor_pivot_loop(query, &mut cursors, |sorted_terms, cursors, bounds| {
136            for &(doc_id, term_index) in sorted_terms {
137                if doc_id == INF_DOC {
138                    bounds.push(0.0);
139                    continue;
140                }
141                let block_index =
142                    block_max.block_index_for(cursors[term_index].block_ordinal()?)?;
143                let block_bound = suffix_bounds[term_index]
144                    .get(block_index)
145                    .copied()
146                    .unwrap_or(0.0);
147                bounds.push(if block_bound > 0.0 {
148                    block_bound
149                } else {
150                    cursors[term_index].upper_bound
151                });
152            }
153            Ok(true)
154        })
155    }
156}
157
158fn validate_cursor_query(query: &CursorWANDQuery) -> StorageBackendResult<()> {
159    let expected = query.cursors.len();
160    if query.scorers.len() == expected
161        && query.fields.len() == expected
162        && query.terms.len() == expected
163    {
164        Ok(())
165    } else {
166        Err(invalid_wand_input(format!(
167            "cursor WAND term arrays must have equal lengths: cursors={expected}, scorers={}, fields={}, terms={}",
168            query.scorers.len(),
169            query.fields.len(),
170            query.terms.len()
171        )))
172    }
173}
174
175fn build_score_cursors(query: &CursorWANDQuery) -> StorageBackendResult<Vec<ScoreTermCursor>> {
176    query
177        .cursors
178        .iter()
179        .cloned()
180        .zip(&query.scorers)
181        .map(|(cursor, scorer)| {
182            let upper_bound = scorer.term_upper_bound(cursor.doc_freq());
183            require_nonnegative_finite(upper_bound, "cursor WAND term upper bound")?;
184            Ok(ScoreTermCursor {
185                cursor,
186                upper_bound,
187            })
188        })
189        .collect()
190}
191
192fn cursor_candidate_upper_bound(query: &CursorWANDQuery) -> StorageBackendResult<u64> {
193    query.cursors.iter().try_fold(0_u64, |total, cursor| {
194        total
195            .checked_add(cursor.doc_freq())
196            .ok_or_else(|| invalid_wand_input("cursor candidate count overflowed"))
197    })
198}
199
200fn run_cursor_pivot_loop<F>(
201    query: &CursorWANDQuery,
202    cursors: &mut [ScoreTermCursor],
203    mut bound_provider: F,
204) -> StorageBackendResult<WANDResult>
205where
206    F: FnMut(&[(DocId, usize)], &[ScoreTermCursor], &mut Vec<f64>) -> StorageBackendResult<bool>,
207{
208    let total_candidates = cursor_candidate_upper_bound(query)?;
209    if cursors.is_empty() || query.k == 0 {
210        return Ok(WANDResult {
211            top_k: PostingList::new(),
212            stats: WANDStats {
213                total_candidates,
214                ..WANDStats::default()
215            },
216        });
217    }
218    let candidate_capacity = usize::try_from(total_candidates).unwrap_or(usize::MAX);
219    let mut top_k = BinaryHeap::with_capacity(query.k.min(candidate_capacity));
220    let mut threshold = 0.0_f64;
221    let mut stats = WANDStats {
222        total_candidates,
223        ..WANDStats::default()
224    };
225    let mut sorted_terms = cursors
226        .iter()
227        .enumerate()
228        .map(|(index, cursor)| (cursor.current_doc(), index))
229        .collect::<Vec<_>>();
230    sorted_terms.sort_unstable();
231    let mut bounds = Vec::with_capacity(cursors.len());
232    let mut term_scores = Vec::with_capacity(cursors.len());
233
234    while sorted_terms
235        .first()
236        .is_some_and(|(doc_id, _)| *doc_id != INF_DOC)
237    {
238        bounds.clear();
239        if !bound_provider(&sorted_terms, cursors, &mut bounds)? {
240            bounds.extend(sorted_terms.iter().map(|&(doc_id, term_index)| {
241                if doc_id == INF_DOC {
242                    0.0
243                } else {
244                    cursors[term_index].upper_bound
245                }
246            }));
247        }
248        let Some(pivot_index) = select_cursor_pivot(query, &sorted_terms, &bounds, threshold)?
249        else {
250            break;
251        };
252        let pivot_doc = sorted_terms[pivot_index].0;
253        if sorted_terms[0].0 == pivot_doc {
254            let score = score_cursor_document(query, cursors, pivot_doc, &mut term_scores)?;
255            stats.scored = stats
256                .scored
257                .checked_add(1)
258                .ok_or_else(|| invalid_wand_input("scored-document counter overflowed"))?;
259            update_top_k(&mut top_k, query.k, score, pivot_doc, &mut threshold);
260            for sorted in &mut sorted_terms {
261                let term_index = sorted.1;
262                if cursors[term_index].current_doc() == pivot_doc {
263                    cursors[term_index].cursor.advance()?;
264                    sorted.0 = cursors[term_index].current_doc();
265                }
266            }
267            sorted_terms.sort_unstable();
268        } else {
269            let term_index = sorted_terms[0].1;
270            cursors[term_index].cursor.advance_to(pivot_doc)?;
271            stats.cursor_advances = stats
272                .cursor_advances
273                .checked_add(1)
274                .ok_or_else(|| invalid_wand_input("cursor-advance counter overflowed"))?;
275            sorted_terms[0].0 = cursors[term_index].current_doc();
276            sorted_terms.sort_unstable();
277        }
278    }
279
280    let mut entries = top_k
281        .into_sorted_vec()
282        .into_iter()
283        .rev()
284        .map(|entry| PostingEntry::new(entry.doc_id, Payload::with_score(entry.score)))
285        .collect::<Vec<_>>();
286    entries.sort_by_key(|entry| entry.doc_id);
287    Ok(WANDResult {
288        top_k: PostingList::from_sorted_unchecked(entries),
289        stats,
290    })
291}
292
293fn select_cursor_pivot(
294    query: &CursorWANDQuery,
295    sorted_terms: &[(DocId, usize)],
296    bounds: &[f64],
297    threshold: f64,
298) -> StorageBackendResult<Option<usize>> {
299    if bounds.len() != sorted_terms.len() {
300        return Err(invalid_wand_input(format!(
301            "cursor bound provider returned {} bounds for {} terms",
302            bounds.len(),
303            sorted_terms.len()
304        )));
305    }
306    for bound in bounds {
307        require_nonnegative_finite(*bound, "cursor WAND pruning bound")?;
308    }
309    for (index, &(doc_id, _)) in sorted_terms.iter().enumerate() {
310        if doc_id == INF_DOC {
311            break;
312        }
313        let cumulative = query.scorers[0].finalize_upper_bound(&bounds[..=index]);
314        require_nonnegative_finite(cumulative, "cursor WAND cumulative upper bound")?;
315        if cumulative >= threshold {
316            return Ok(Some(index));
317        }
318    }
319    Ok(None)
320}
321
322fn score_cursor_document(
323    query: &CursorWANDQuery,
324    cursors: &[ScoreTermCursor],
325    target: DocId,
326    term_scores: &mut Vec<f64>,
327) -> StorageBackendResult<f64> {
328    term_scores.clear();
329    for (index, cursor) in cursors.iter().enumerate() {
330        let Some(entry) = cursor.cursor.current() else {
331            continue;
332        };
333        if entry.doc_id != target {
334            continue;
335        }
336        let term_score = query.scorers[index].term_score(
337            entry.term_freq,
338            entry.doc_length.max(entry.term_freq),
339            cursor.cursor.doc_freq(),
340        );
341        require_nonnegative_finite(term_score, "cursor WAND term score")?;
342        term_scores.push(term_score);
343    }
344    let score = query.scorers[0].finalize_score(term_scores);
345    require_nonnegative_finite(score, "cursor WAND finalized score")?;
346    Ok(score)
347}
uqa_scoring/wand/cursor.rs

uqa_scoring/wand/
cursor.rs
