uqa_scoring/wand/
materialized.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Materialized posting-list cursor state and exact WAND/BMW pivot loops.
8
9use std::cmp::Reverse;
10use std::collections::BinaryHeap;
11use std::sync::Arc;
12
13use uqa_core::{DocId, FieldName, Payload, PostingEntry, PostingList};
14use uqa_storage::{BlockMaxIndex, InvertedIndex, StorageBackendResult};
15
16use crate::scorer::Scorer;
17
18use super::common::{
19    invalid_wand_input, require_nonnegative_finite, update_top_k, HeapEntry, WANDResult, WANDStats,
20    INF_DOC,
21};
22
23/// Common per-term cursor state: the entry slice and current position.
24/// Field, term, and scorer live on [`WANDQuery`] so a single cursor
25/// stays small and pivot reordering touches just one cache line.
26struct TermCursor<'a> {
27    entries: &'a [PostingEntry],
28    position: usize,
29    doc_freq: u64,
30    upper_bound: f64,
31}
32
33impl<'a> TermCursor<'a> {
34    fn current_doc(&self) -> u64 {
35        self.entries
36            .get(self.position)
37            .map_or(INF_DOC, |e| e.doc_id)
38    }
39
40    fn current(&self) -> Option<&'a PostingEntry> {
41        self.entries.get(self.position)
42    }
43
44    /// Binary search advance to the first entry with `doc_id >= target`.
45    fn advance_to(&mut self, target: u64) {
46        let mut lo = self.position;
47        let mut hi = self.entries.len();
48        while lo < hi {
49            let mid = lo + (hi - lo) / 2;
50            if self.entries[mid].doc_id < target {
51                lo = mid + 1;
52            } else {
53                hi = mid;
54            }
55        }
56        self.position = lo;
57    }
58}
59
60/// WAND configuration shared by both algorithms.
61pub struct WANDQuery {
62    pub posting_lists: Vec<PostingList>,
63    pub scorers: Vec<Arc<dyn Scorer>>,
64    pub fields: Vec<FieldName>,
65    pub terms: Vec<String>,
66    pub k: usize,
67}
68
69impl WANDQuery {
70    pub fn new(
71        posting_lists: Vec<PostingList>,
72        scorers: Vec<Arc<dyn Scorer>>,
73        fields: Vec<FieldName>,
74        terms: Vec<String>,
75        k: usize,
76    ) -> StorageBackendResult<Self> {
77        let expected = posting_lists.len();
78        if scorers.len() != expected || fields.len() != expected || terms.len() != expected {
79            return Err(invalid_wand_input(format!(
80                "WAND term arrays must have equal lengths: posting_lists={expected}, scorers={}, fields={}, terms={}",
81                scorers.len(),
82                fields.len(),
83                terms.len()
84            )));
85        }
86        Ok(Self {
87            posting_lists,
88            scorers,
89            fields,
90            terms,
91            k,
92        })
93    }
94}
95
96/// Standard WAND with per-term `term_upper_bound(df)` pruning.
97pub struct WANDScorer<'a> {
98    query: &'a WANDQuery,
99    inverted_index: Option<&'a dyn InvertedIndex>,
100}
101
102impl<'a> WANDScorer<'a> {
103    pub fn new(query: &'a WANDQuery, inverted_index: Option<&'a dyn InvertedIndex>) -> Self {
104        Self {
105            query,
106            inverted_index,
107        }
108    }
109
110    pub fn score_top_k(&self) -> StorageBackendResult<WANDResult> {
111        validate_query(self.query)?;
112        let mut cursors = build_cursors(self.query)?;
113        run_pivot_loop(self.query, &mut cursors, self.inverted_index, |_, _, _| {
114            Ok(false)
115        })
116    }
117}
118
119/// Block-Max WAND: pivot pruning uses per-block max scores from
120/// [`BlockMaxIndex`] for tighter bounds.
121pub struct BlockMaxWANDScorer<'a> {
122    query: &'a WANDQuery,
123    inverted_index: Option<&'a dyn InvertedIndex>,
124    block_max_index: &'a BlockMaxIndex,
125    table: String,
126}
127
128impl<'a> BlockMaxWANDScorer<'a> {
129    pub fn new(
130        query: &'a WANDQuery,
131        inverted_index: Option<&'a dyn InvertedIndex>,
132        block_max_index: &'a BlockMaxIndex,
133        table: impl Into<String>,
134    ) -> Self {
135        Self {
136            query,
137            inverted_index,
138            block_max_index,
139            table: table.into(),
140        }
141    }
142
143    pub fn score_top_k(&self) -> StorageBackendResult<WANDResult> {
144        validate_query(self.query)?;
145        let mut cursors = build_cursors(self.query)?;
146        let q = self.query;
147        let bmi = self.block_max_index;
148        let table = &self.table;
149        let suffix_bounds = q
150            .fields
151            .iter()
152            .zip(&q.terms)
153            .map(|(field, term)| {
154                let Some(blocks) = bmi.block_maxes(table, field, term) else {
155                    return Vec::new();
156                };
157                let mut suffix = vec![0.0_f64; blocks.len()];
158                let mut maximum = 0.0_f64;
159                for (index, score) in blocks.iter().enumerate().rev() {
160                    maximum = maximum.max(*score);
161                    suffix[index] = maximum;
162                }
163                suffix
164            })
165            .collect::<Vec<_>>();
166        run_pivot_loop(
167            q,
168            &mut cursors,
169            self.inverted_index,
170            |sorted_terms, cursors, bounds| {
171                for &(doc_val, ti) in sorted_terms {
172                    if doc_val == INF_DOC {
173                        bounds.push(0.0);
174                        continue;
175                    }
176                    let cur_block = bmi.block_index_for(cursors[ti].position)?;
177                    // Take the max block-max across the remaining blocks
178                    // for this term so the bound stays valid for any
179                    // pivot doc the cursor could still reach. Anchoring
180                    // on just `cur_block` under-counts a later block
181                    // whose max score exceeds the current block, which
182                    // would prune candidates BMW must still consider.
183                    let bm = suffix_bounds[ti].get(cur_block).copied().unwrap_or(0.0);
184                    // Fall back to the per-term `term_upper_bound(df)` if no
185                    // block was recorded; an unindexed term must not get
186                    // pruned more aggressively than plain WAND.
187                    let bound = if bm > 0.0 {
188                        bm
189                    } else {
190                        cursors[ti].upper_bound
191                    };
192                    bounds.push(bound);
193                }
194                Ok(true)
195            },
196        )
197    }
198}
199
200fn build_cursors(query: &WANDQuery) -> StorageBackendResult<Vec<TermCursor<'_>>> {
201    let mut cursors = Vec::with_capacity(query.posting_lists.len());
202    for i in 0..query.posting_lists.len() {
203        let entries = query.posting_lists[i].entries();
204        let df = u64::try_from(entries.len())
205            .map_err(|_| invalid_wand_input("posting-list length does not fit in u64"))?;
206        let upper_bound = query.scorers[i].term_upper_bound(df);
207        require_nonnegative_finite(upper_bound, "WAND term upper bound")?;
208        cursors.push(TermCursor {
209            entries,
210            position: 0,
211            doc_freq: df,
212            upper_bound,
213        });
214    }
215    Ok(cursors)
216}
217
218fn validate_query(query: &WANDQuery) -> StorageBackendResult<()> {
219    let expected = query.posting_lists.len();
220    if query.scorers.len() == expected
221        && query.fields.len() == expected
222        && query.terms.len() == expected
223    {
224        Ok(())
225    } else {
226        Err(invalid_wand_input(format!(
227            "WAND term arrays must have equal lengths: posting_lists={expected}, scorers={}, fields={}, terms={}",
228            query.scorers.len(),
229            query.fields.len(),
230            query.terms.len()
231        )))
232    }
233}
234
235fn build_field_slots(fields: &[FieldName]) -> (Vec<usize>, usize) {
236    let mut unique_fields = Vec::<&str>::with_capacity(fields.len());
237    let slots = fields
238        .iter()
239        .map(|field| {
240            if let Some(slot) = unique_fields.iter().position(|known| *known == field) {
241                slot
242            } else {
243                let slot = unique_fields.len();
244                unique_fields.push(field);
245                slot
246            }
247        })
248        .collect();
249    (slots, unique_fields.len())
250}
251
252/// Core pivot loop shared by WAND and BMW. `bound_provider` fills a reusable
253/// per-term bound vector aligned with the current `sorted_terms` order and
254/// returns `true`; `false` selects each cursor's precomputed upper bound.
255fn run_pivot_loop<F>(
256    query: &WANDQuery,
257    cursors: &mut [TermCursor<'_>],
258    inverted_index: Option<&dyn InvertedIndex>,
259    mut bound_provider: F,
260) -> StorageBackendResult<WANDResult>
261where
262    F: FnMut(&[(u64, usize)], &[TermCursor<'_>], &mut Vec<f64>) -> StorageBackendResult<bool>,
263{
264    let num_terms = query.posting_lists.len();
265    let total_candidates = candidate_union(&query.posting_lists)?;
266    if num_terms == 0 || query.k == 0 {
267        return Ok(WANDResult {
268            top_k: PostingList::new(),
269            stats: WANDStats {
270                total_candidates,
271                ..WANDStats::default()
272            },
273        });
274    }
275    let candidate_capacity = usize::try_from(total_candidates).unwrap_or(usize::MAX);
276    let mut top_k: BinaryHeap<HeapEntry> =
277        BinaryHeap::with_capacity(query.k.min(candidate_capacity));
278    let mut threshold = 0.0_f64;
279    let mut stats = WANDStats {
280        total_candidates,
281        ..WANDStats::default()
282    };
283
284    let mut sorted_terms: Vec<(u64, usize)> = (0..num_terms)
285        .map(|i| (cursors[i].current_doc(), i))
286        .collect();
287    sorted_terms.sort_unstable();
288    let mut bounds = Vec::with_capacity(num_terms);
289    let mut term_scores = Vec::with_capacity(num_terms);
290    let (field_slots, field_count) = build_field_slots(&query.fields);
291    let mut doc_lengths = vec![None; field_count];
292
293    while !sorted_terms.is_empty() {
294        if sorted_terms[0].0 == INF_DOC {
295            break;
296        }
297
298        bounds.clear();
299        if !bound_provider(&sorted_terms, cursors, &mut bounds)? {
300            bounds.extend(sorted_terms.iter().map(|&(doc_val, ti)| {
301                if doc_val == INF_DOC {
302                    0.0
303                } else {
304                    cursors[ti].upper_bound
305                }
306            }));
307        }
308        let Some(pivot_idx) = select_pivot(query, &sorted_terms, &bounds, threshold)? else {
309            break;
310        };
311
312        let pivot_doc = sorted_terms[pivot_idx].0;
313        let first_doc = sorted_terms[0].0;
314
315        if first_doc == pivot_doc {
316            let actual_score = score_document(
317                query,
318                cursors,
319                inverted_index,
320                pivot_doc as DocId,
321                &field_slots,
322                &mut doc_lengths,
323                &mut term_scores,
324            )?;
325            stats.scored = stats
326                .scored
327                .checked_add(1)
328                .ok_or_else(|| invalid_wand_input("scored-document counter overflowed"))?;
329
330            update_top_k(
331                &mut top_k,
332                query.k,
333                actual_score,
334                pivot_doc as DocId,
335                &mut threshold,
336            );
337
338            // Advance every cursor at pivot_doc.
339            for st in &mut sorted_terms {
340                let ti = st.1;
341                if cursors[ti].current_doc() == pivot_doc {
342                    cursors[ti].position += 1;
343                    st.0 = cursors[ti].current_doc();
344                }
345            }
346            sorted_terms.sort_unstable();
347        } else {
348            // Skip first cursor forward to pivot_doc.
349            let first_term = sorted_terms[0].1;
350            cursors[first_term].advance_to(pivot_doc);
351            stats.cursor_advances = stats
352                .cursor_advances
353                .checked_add(1)
354                .ok_or_else(|| invalid_wand_input("cursor-advance counter overflowed"))?;
355            sorted_terms[0].0 = cursors[first_term].current_doc();
356            sorted_terms.sort_unstable();
357        }
358    }
359
360    let mut entries: Vec<PostingEntry> = top_k
361        .into_sorted_vec()
362        .into_iter()
363        .rev()
364        .map(|h| PostingEntry::new(h.doc_id, Payload::with_score(h.score)))
365        .collect();
366    entries.sort_by_key(|e| e.doc_id);
367    Ok(WANDResult {
368        top_k: PostingList::from_sorted_unchecked(entries),
369        stats,
370    })
371}
372
373fn select_pivot(
374    query: &WANDQuery,
375    sorted_terms: &[(u64, usize)],
376    bounds: &[f64],
377    threshold: f64,
378) -> StorageBackendResult<Option<usize>> {
379    if bounds.len() != sorted_terms.len() {
380        return Err(invalid_wand_input(format!(
381            "bound provider returned {} bounds for {} terms",
382            bounds.len(),
383            sorted_terms.len()
384        )));
385    }
386    for bound in bounds {
387        require_nonnegative_finite(*bound, "WAND pruning bound")?;
388    }
389    for (index, &(doc_id, _)) in sorted_terms.iter().enumerate() {
390        if doc_id == INF_DOC {
391            break;
392        }
393        let cumulative = query.scorers[0].finalize_upper_bound(&bounds[..=index]);
394        require_nonnegative_finite(cumulative, "WAND cumulative upper bound")?;
395        if cumulative >= threshold {
396            return Ok(Some(index));
397        }
398    }
399    Ok(None)
400}
401
402pub(super) fn candidate_union(posting_lists: &[PostingList]) -> StorageBackendResult<u64> {
403    let mut positions = vec![0_usize; posting_lists.len()];
404    let mut next = BinaryHeap::<Reverse<(DocId, usize)>>::with_capacity(posting_lists.len());
405    for (list_index, posting) in posting_lists.iter().enumerate() {
406        if let Some(entry) = posting.entries().first() {
407            next.push(Reverse((entry.doc_id, list_index)));
408        }
409    }
410    let mut count = 0_u64;
411    let mut previous = None;
412    while let Some(Reverse((doc_id, list_index))) = next.pop() {
413        if previous != Some(doc_id) {
414            count = count
415                .checked_add(1)
416                .ok_or_else(|| invalid_wand_input("candidate union length does not fit in u64"))?;
417            previous = Some(doc_id);
418        }
419        let entries = posting_lists[list_index].entries();
420        let position = &mut positions[list_index];
421        while entries
422            .get(*position)
423            .is_some_and(|entry| entry.doc_id == doc_id)
424        {
425            *position += 1;
426        }
427        if let Some(entry) = entries.get(*position) {
428            next.push(Reverse((entry.doc_id, list_index)));
429        }
430    }
431    Ok(count)
432}
433
434/// Score a single document against every term cursor. Cursors that
435/// point at a different `doc_id` contribute nothing; cursors that point
436/// at `target` contribute a raw term score. The query score is finalized
437/// once after all term contributions have been collected.
438fn score_document(
439    query: &WANDQuery,
440    cursors: &[TermCursor<'_>],
441    inverted_index: Option<&dyn InvertedIndex>,
442    target: DocId,
443    field_slots: &[usize],
444    doc_lengths: &mut [Option<u64>],
445    term_scores: &mut Vec<f64>,
446) -> StorageBackendResult<f64> {
447    doc_lengths.fill(None);
448    term_scores.clear();
449    for (i, cursor) in cursors.iter().enumerate() {
450        let Some(entry) = cursor.current() else {
451            continue;
452        };
453        if entry.doc_id != target {
454            continue;
455        }
456        let tf = if entry.payload.positions.is_empty() {
457            1
458        } else {
459            u64::try_from(entry.payload.positions.len())
460                .map_err(|_| invalid_wand_input("term frequency does not fit in u64"))?
461        };
462        let df = cursor.doc_freq;
463        let doc_length = match inverted_index {
464            Some(idx) => {
465                let slot = field_slots[i];
466                let length = if let Some(length) = doc_lengths[slot] {
467                    length
468                } else {
469                    let length = idx.get_doc_length(target, &query.fields[i])?;
470                    doc_lengths[slot] = Some(length);
471                    length
472                };
473                length.max(tf)
474            }
475            None => tf,
476        };
477        let term_score = query.scorers[i].term_score(tf, doc_length, df);
478        require_nonnegative_finite(term_score, "WAND term score")?;
479        term_scores.push(term_score);
480    }
481    let score = query.scorers[0].finalize_score(term_scores);
482    require_nonnegative_finite(score, "WAND finalized score")?;
483    Ok(score)
484}
uqa_scoring/wand/materialized.rs

uqa_scoring/wand/
materialized.rs
