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

uqa_scoring/wand/
cursor.rs
