uqa_scoring/wand/
common.rs1use std::cmp::Ordering;
10use std::collections::BinaryHeap;
11
12use uqa_core::{DocId, PostingList};
13use uqa_storage::{StorageBackendError, StorageBackendResult};
14
15pub(super) const INF_DOC: u64 = u64::MAX;
16
17#[derive(Debug, Clone, Copy)]
19pub(super) struct HeapEntry {
20 pub(super) score: f64,
21 pub(super) doc_id: DocId,
22}
23
24impl PartialEq for HeapEntry {
25 fn eq(&self, other: &Self) -> bool {
26 self.score == other.score && self.doc_id == other.doc_id
27 }
28}
29
30impl Eq for HeapEntry {}
31
32impl PartialOrd for HeapEntry {
33 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
34 Some(self.cmp(other))
35 }
36}
37
38impl Ord for HeapEntry {
39 fn cmp(&self, other: &Self) -> Ordering {
40 match other.score.total_cmp(&self.score) {
46 Ordering::Equal => self.doc_id.cmp(&other.doc_id),
47 ord => ord,
48 }
49 }
50}
51
52pub(super) fn update_top_k(
53 top_k: &mut BinaryHeap<HeapEntry>,
54 k: usize,
55 score: f64,
56 doc_id: DocId,
57 threshold: &mut f64,
58) {
59 let candidate = HeapEntry { score, doc_id };
60 if top_k.len() < k {
61 top_k.push(candidate);
62 if top_k.len() == k {
63 *threshold = top_k.peek().map_or(0.0, |entry| entry.score);
64 }
65 return;
66 }
67 let Some(eviction) = top_k.peek() else {
68 return;
69 };
70 if score > eviction.score || (score == eviction.score && doc_id < eviction.doc_id) {
71 top_k.pop();
72 top_k.push(candidate);
73 *threshold = top_k.peek().map_or(*threshold, |entry| entry.score);
74 }
75}
76
77#[derive(Debug, Default, Clone, Copy, PartialEq)]
87pub struct WANDStats {
88 pub scored: u64,
89 pub total_candidates: u64,
90 pub cursor_advances: u64,
91}
92
93impl WANDStats {
94 pub fn skip_rate(&self) -> f64 {
95 if self.total_candidates == 0 {
96 0.0
97 } else {
98 1.0 - (self.scored as f64 / self.total_candidates as f64)
99 }
100 }
101}
102
103#[derive(Debug, Clone)]
104pub struct WANDResult {
105 pub top_k: PostingList,
106 pub stats: WANDStats,
107}
108
109pub(super) fn invalid_wand_input(message: impl Into<String>) -> StorageBackendError {
110 StorageBackendError::Other(format!("invalid WAND input: {}", message.into()))
111}
112
113pub(super) fn require_nonnegative_finite(value: f64, name: &str) -> StorageBackendResult<()> {
114 if value.is_finite() && value >= 0.0 {
115 Ok(())
116 } else {
117 Err(invalid_wand_input(format!(
118 "{name} must be finite and non-negative, got {value}"
119 )))
120 }
121}