Skip to main content

uqa_core/
ranked_view.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Score-ordered views over document-id-ordered posting storage.
8
9use std::cmp::Ordering;
10use std::sync::OnceLock;
11
12use crate::{PostingEntry, PostingList};
13
14/// A score-ordered, non-owning view of a [`PostingList`].
15///
16/// Posting storage remains sorted by document id. The separate order by
17/// descending payload score is built lazily, with ascending document id as the
18/// deterministic tie-breaker. This lets support-preserving selections such as
19/// `select_top_k(len)` avoid ranking work entirely.
20#[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    /// Borrow every entry in rank order.
49    pub fn entries(&self) -> &[&'a PostingEntry] {
50        self.entries.get_or_init(|| Self::rank_entries(self.source))
51    }
52
53    /// Borrow at most the first `k` entries in rank order.
54    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    /// Materialize the top `k` selection as document-id-ordered posting
63    /// storage.
64    ///
65    /// Rank order intentionally does not leak into [`PostingList`]'s physical
66    /// ordering invariant. If the view has not already been ranked, this uses
67    /// linear-time selection rather than sorting entries that will immediately
68    /// be reordered by document id. Use [`Self::top_k`] when rank order itself
69    /// is needed.
70    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    /// Number of ranked entries.
92    pub fn len(&self) -> usize {
93        self.source.len()
94    }
95
96    /// Whether the view contains no entries.
97    pub fn is_empty(&self) -> bool {
98        self.source.is_empty()
99    }
100
101    /// Iterate in descending score order.
102    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}