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
15pub(super) const INF_DOC: u64 = u64::MAX;
16
17/// Min-heap entry by score for top-k selection.
18#[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        // `BinaryHeap` is a max-heap; flip the score comparison so the
41        // root holds the *minimum* score (the eviction candidate). On a
42        // score tie, the entry with the *larger* doc id sits at the
43        // root, matching the conventional "lower doc id wins" tie break
44        // applied at output time.
45        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/// Stats collected during a top-k pass; tests use these to assert the
78/// exit-criterion skip rates from the master plan.
79///
80/// Skip rate semantics are `1 - scored / total_candidates`. The materialized
81/// path reports the exact union of posting-list document ids. The score-cursor
82/// path deliberately avoids a complete pre-scan and reports the sum of term
83/// document frequencies, a safe upper bound on that union. `scored` counts
84/// documents for which the complete query score was evaluated;
85/// `cursor_advances` counts pivot-driven skips and is informational only.
86#[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}