1use std::cmp::Ordering;
10use std::sync::OnceLock;
11
12use crate::{PostingEntry, PostingList};
13
14#[derive(Debug)]
21pub struct RankedView<'a> {
22 source: &'a PostingList,
23 entries: OnceLock<Vec<&'a PostingEntry>>,
24}
25
26impl<'a> RankedView<'a> {
27 pub(crate) fn new(posting_list: &'a PostingList) -> Self {
28 Self {
29 source: posting_list,
30 entries: OnceLock::new(),
31 }
32 }
33
34 fn compare_rank(left: &PostingEntry, right: &PostingEntry) -> Ordering {
35 right
36 .payload
37 .score
38 .total_cmp(&left.payload.score)
39 .then_with(|| left.doc_id.cmp(&right.doc_id))
40 }
41
42 fn rank_entries(posting_list: &'a PostingList) -> Vec<&'a PostingEntry> {
43 let mut entries: Vec<&PostingEntry> = posting_list.entries().iter().collect();
44 entries.sort_by(|left, right| Self::compare_rank(left, right));
45 entries
46 }
47
48 pub fn entries(&self) -> &[&'a PostingEntry] {
50 self.entries.get_or_init(|| Self::rank_entries(self.source))
51 }
52
53 pub fn top_k(&self, k: usize) -> &[&'a PostingEntry] {
55 if k == 0 {
56 return &[];
57 }
58 let entries = self.entries();
59 &entries[..k.min(entries.len())]
60 }
61
62 pub fn select_top_k(self, k: usize) -> PostingList {
71 if k == 0 {
72 return PostingList::new();
73 }
74 if k >= self.source.len() {
75 return self.source.clone();
76 }
77
78 let source = self.source;
79 let selected = if let Some(ranked) = self.entries.into_inner() {
80 ranked
81 } else {
82 let mut candidates: Vec<&PostingEntry> = source.entries().iter().collect();
83 candidates.select_nth_unstable_by(k, |left, right| Self::compare_rank(left, right));
84 candidates
85 };
86 let mut entries: Vec<PostingEntry> = selected.into_iter().take(k).cloned().collect();
87 entries.sort_by_key(|entry| entry.doc_id);
88 PostingList::from_sorted_unchecked(entries)
89 }
90
91 pub fn len(&self) -> usize {
93 self.source.len()
94 }
95
96 pub fn is_empty(&self) -> bool {
98 self.source.is_empty()
99 }
100
101 pub fn iter(
103 &self,
104 ) -> impl ExactSizeIterator<Item = &'a PostingEntry> + DoubleEndedIterator + '_ {
105 self.entries().iter().copied()
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use crate::{Payload, PostingEntry, PostingList};
112
113 #[test]
114 fn ranking_is_separate_from_posting_storage_order() {
115 let posting = PostingList::from_unsorted(vec![
116 PostingEntry::new(1, Payload::with_score(0.2)),
117 PostingEntry::new(2, Payload::with_score(0.9)),
118 PostingEntry::new(3, Payload::with_score(0.9)),
119 ]);
120
121 let ranked = posting.ranked();
122 let ranked_ids: Vec<_> = ranked.iter().map(|entry| entry.doc_id).collect();
123 assert_eq!(ranked_ids, vec![2, 3, 1]);
124
125 let selected = ranked.select_top_k(2);
126 assert_eq!(selected.doc_ids().collect::<Vec<_>>(), vec![2, 3]);
127 }
128
129 #[test]
130 fn selecting_the_entire_source_preserves_every_payload() {
131 let posting = PostingList::from_unsorted(vec![
132 PostingEntry::new(1, Payload::with_score(0.2)),
133 PostingEntry::new(2, Payload::with_score(0.9)),
134 PostingEntry::new(3, Payload::with_score(0.5)),
135 ]);
136
137 assert_eq!(posting.ranked().select_top_k(posting.len()), posting);
138 }
139
140 #[test]
141 fn selecting_zero_entries_is_empty() {
142 let posting =
143 PostingList::from_unsorted(vec![PostingEntry::new(1, Payload::with_score(0.2))]);
144
145 assert!(posting.ranked().top_k(0).is_empty());
146 assert!(posting.ranked().select_top_k(0).is_empty());
147 }
148
149 #[test]
150 fn materialized_selection_matches_the_full_rank_order() {
151 let posting = PostingList::from_unsorted(vec![
152 PostingEntry::new(1, Payload::with_score(f64::NAN)),
153 PostingEntry::new(2, Payload::with_score(0.9)),
154 PostingEntry::new(3, Payload::with_score(0.9)),
155 PostingEntry::new(4, Payload::with_score(f64::INFINITY)),
156 PostingEntry::new(5, Payload::with_score(f64::NEG_INFINITY)),
157 ]);
158 let ranked_entries = posting.ranked().entries().to_vec();
159
160 for k in 0..=posting.len() + 1 {
161 let mut expected: Vec<_> = ranked_entries
162 .iter()
163 .take(k)
164 .map(|entry| (entry.doc_id, entry.payload.score.to_bits()))
165 .collect();
166 expected.sort_unstable_by_key(|(doc_id, _)| *doc_id);
167
168 let actual: Vec<_> = posting
169 .ranked()
170 .select_top_k(k)
171 .entries()
172 .iter()
173 .map(|entry| (entry.doc_id, entry.payload.score.to_bits()))
174 .collect();
175 assert_eq!(actual, expected);
176 }
177 }
178}