uqa_scoring/wand/
common.rs1use std::cmp::Ordering;
10use std::collections::BinaryHeap;
11
12use uqa_core::{DocId, PostingList};
13use uqa_storage::{StorageBackendError, StorageBackendResult};
14
15#[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 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#[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}