Skip to main content

uqa_scoring/wand/
common.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Shared top-k heap, result statistics, and bound validation.
8
9use std::cmp::Ordering;
10use std::collections::BinaryHeap;
11
12use uqa_core::{DocId, PostingList};
13use uqa_storage::{StorageBackendError, StorageBackendResult};
14
15/// Min-heap entry by score for top-k selection.
16#[derive(Debug, Clone, Copy)]
17pub(super) struct HeapEntry {
18    pub(super) score: f64,
19    pub(super) doc_id: DocId,
20}
21
22impl PartialEq for HeapEntry {
23    fn eq(&self, other: &Self) -> bool {
24        self.score == other.score && self.doc_id == other.doc_id
25    }
26}
27
28impl Eq for HeapEntry {}
29
30impl PartialOrd for HeapEntry {
31    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
32        Some(self.cmp(other))
33    }
34}
35
36impl Ord for HeapEntry {
37    fn cmp(&self, other: &Self) -> Ordering {
38        // `BinaryHeap` is a max-heap; flip the score comparison so the
39        // root holds the *minimum* score (the eviction candidate). On a
40        // score tie, the entry with the *larger* doc id sits at the
41        // root, matching the conventional "lower doc id wins" tie break
42        // applied at output time.
43        match other.score.total_cmp(&self.score) {
44            Ordering::Equal => self.doc_id.cmp(&other.doc_id),
45            ord => ord,
46        }
47    }
48}
49
50pub(super) fn update_top_k(
51    top_k: &mut BinaryHeap<HeapEntry>,
52    k: usize,
53    score: f64,
54    doc_id: DocId,
55    threshold: &mut f64,
56) {
57    let candidate = HeapEntry { score, doc_id };
58    if top_k.len() < k {
59        top_k.push(candidate);
60        if top_k.len() == k {
61            *threshold = top_k.peek().map_or(0.0, |entry| entry.score);
62        }
63        return;
64    }
65    let Some(eviction) = top_k.peek() else {
66        return;
67    };
68    if score > eviction.score || (score == eviction.score && doc_id < eviction.doc_id) {
69        top_k.pop();
70        top_k.push(candidate);
71        *threshold = top_k.peek().map_or(*threshold, |entry| entry.score);
72    }
73}
74
75/// Stats collected during a top-k pass; tests use these to assert the
76/// exit-criterion skip rates from the master plan.
77///
78/// Skip rate semantics are `1 - scored / total_candidates`. The materialized
79/// path reports the exact union of posting-list document ids. The score-cursor
80/// path deliberately avoids a complete pre-scan and reports the sum of term
81/// document frequencies, a safe upper bound on that union. `scored` counts
82/// documents for which the complete query score was evaluated;
83/// `cursor_advances` counts pivot-driven skips and is informational only.
84#[derive(Debug, Default, Clone, Copy, PartialEq)]
85pub struct WANDStats {
86    pub scored: u64,
87    pub total_candidates: u64,
88    pub cursor_advances: u64,
89}
90
91impl WANDStats {
92    pub fn skip_rate(&self) -> f64 {
93        if self.total_candidates == 0 {
94            0.0
95        } else {
96            1.0 - (self.scored as f64 / self.total_candidates as f64)
97        }
98    }
99}
100
101#[derive(Debug, Clone)]
102pub struct WANDResult {
103    pub top_k: PostingList,
104    pub stats: WANDStats,
105}
106
107pub(super) fn invalid_wand_input(message: impl Into<String>) -> StorageBackendError {
108    StorageBackendError::Other(format!("invalid WAND input: {}", message.into()))
109}
110
111pub(super) fn require_nonnegative_finite(value: f64, name: &str) -> StorageBackendResult<()> {
112    if value.is_finite() && value >= 0.0 {
113        Ok(())
114    } else {
115        Err(invalid_wand_input(format!(
116            "{name} must be finite and non-negative, got {value}"
117        )))
118    }
119}