sparse_vector/wand/
sink.rs

1//! Where scored records go: a top-k tracker or a plain collector.
2
3use std::cmp::{Ordering, Reverse};
4use std::collections::BinaryHeap;
5
6use ordered_float::OrderedFloat;
7
8use super::RecordId;
9
10/// A scored record. Ordered so that a *greater* hit is a better one: higher
11/// score first, then lower id.
12#[derive(Clone, Copy, Debug, PartialEq)]
13pub struct Hit {
14    pub id: RecordId,
15    pub score: f32,
16}
17
18impl Eq for Hit {}
19
20impl PartialOrd for Hit {
21    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
22        Some(self.cmp(other))
23    }
24}
25
26impl Ord for Hit {
27    fn cmp(&self, other: &Self) -> Ordering {
28        OrderedFloat(self.score)
29            .cmp(&OrderedFloat(other.score))
30            .then_with(|| other.id.cmp(&self.id))
31    }
32}
33
34/// Consumer of `(id, score)` pairs produced by the search loop.
35///
36/// The loop offers records in increasing id order and asks for the
37/// [`threshold`](Self::threshold) before scoring a window: a record whose
38/// score cannot *strictly* exceed it is not worth scoring. Because later
39/// records have larger ids and ties go to the lower id, a score equal to the
40/// threshold can never displace what the sink already holds.
41pub trait ScoreSink {
42    /// Offer one scored record.
43    fn offer(&mut self, id: RecordId, score: f32);
44
45    /// Score a new record must strictly beat to be retained, or `None` if
46    /// every record is still welcome.
47    fn threshold(&self) -> Option<f32>;
48
49    /// Retained records, sorted by score descending then id ascending.
50    fn into_results(self) -> Vec<(RecordId, f32)>;
51}
52
53/// Keeps the `k` best hits (highest score, ties to the lower id).
54#[derive(Debug)]
55pub struct TopKSink {
56    k: usize,
57    /// Min-heap on `Hit`: the top is the worst retained hit.
58    heap: BinaryHeap<Reverse<Hit>>,
59}
60
61impl TopKSink {
62    pub fn new(k: usize) -> Self {
63        Self {
64            k,
65            heap: BinaryHeap::with_capacity(k.saturating_add(1)),
66        }
67    }
68
69    pub fn len(&self) -> usize {
70        self.heap.len()
71    }
72
73    pub fn is_empty(&self) -> bool {
74        self.heap.is_empty()
75    }
76
77    pub fn capacity(&self) -> usize {
78        self.k
79    }
80
81    fn worst(&self) -> Option<Hit> {
82        self.heap.peek().map(|r| r.0)
83    }
84}
85
86impl ScoreSink for TopKSink {
87    fn offer(&mut self, id: RecordId, score: f32) {
88        if self.k == 0 {
89            return;
90        }
91        let hit = Hit { id, score };
92        if self.heap.len() < self.k {
93            self.heap.push(Reverse(hit));
94        } else if let Some(worst) = self.worst() {
95            if hit > worst {
96                self.heap.pop();
97                self.heap.push(Reverse(hit));
98            }
99        }
100    }
101
102    fn threshold(&self) -> Option<f32> {
103        if self.k == 0 {
104            return Some(f32::INFINITY);
105        }
106        if self.heap.len() < self.k {
107            None
108        } else {
109            self.worst().map(|h| h.score)
110        }
111    }
112
113    fn into_results(self) -> Vec<(RecordId, f32)> {
114        let mut hits: Vec<Hit> = self.heap.into_iter().map(|r| r.0).collect();
115        hits.sort_by(|a, b| b.cmp(a));
116        hits.into_iter().map(|h| (h.id, h.score)).collect()
117    }
118}
119
120/// Keeps everything it is offered.
121#[derive(Debug, Default)]
122pub struct CollectAll {
123    hits: Vec<Hit>,
124}
125
126impl CollectAll {
127    pub fn new() -> Self {
128        Self::default()
129    }
130
131    pub fn len(&self) -> usize {
132        self.hits.len()
133    }
134
135    pub fn is_empty(&self) -> bool {
136        self.hits.is_empty()
137    }
138}
139
140impl ScoreSink for CollectAll {
141    fn offer(&mut self, id: RecordId, score: f32) {
142        self.hits.push(Hit { id, score });
143    }
144
145    fn threshold(&self) -> Option<f32> {
146        None
147    }
148
149    fn into_results(mut self) -> Vec<(RecordId, f32)> {
150        self.hits.sort_by(|a, b| b.cmp(a));
151        self.hits.into_iter().map(|h| (h.id, h.score)).collect()
152    }
153}
sparse_vector/wand/sink.rs

sparse_vector/wand/
sink.rs
