subms_merge_iterator/features/
priority.rs

1//! Priority-aware k-way merge. Each source carries an explicit
2//! `priority` integer. On key tie, the highest-priority source wins.
3//! Ties between equal priorities fall through to the source's
4//! registration order (higher source index breaks the tie - matches
5//! the latest-source-wins shape used by `dedup` and `tombstones`).
6//!
7//! This generalises `dedup`: pass `(priority = source_index, ...)` for
8//! the same behaviour. The reason to reach for `priority` is when the
9//! producer order on the wire doesn't match the recency / authority
10//! ordering you want (e.g. an in-memory memtable should beat every
11//! on-disk SSTable level even though it was registered first).
12
13use std::cmp::Reverse;
14use std::collections::BinaryHeap;
15
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct PriorityEntry<K, V> {
18    pub key: K,
19    pub value: V,
20}
21
22impl<K, V> PriorityEntry<K, V> {
23    pub fn new(key: K, value: V) -> Self {
24        Self { key, value }
25    }
26}
27
28/// A source + the priority it carries. Higher `priority` wins on key
29/// ties. Equal-priority ties fall through to registration order (the
30/// index in the `new(...)` argument list).
31pub struct PrioritySource<I> {
32    pub priority: i32,
33    pub stream: I,
34}
35
36impl<I> PrioritySource<I> {
37    pub fn new(priority: i32, stream: I) -> Self {
38        Self { priority, stream }
39    }
40}
41
42pub struct PriorityMergeIterator<K, V, I>
43where
44    K: Ord,
45    I: Iterator<Item = PriorityEntry<K, V>>,
46{
47    streams: Vec<I>,
48    priorities: Vec<i32>,
49    heap: BinaryHeap<Reverse<HeapItem<K, V>>>,
50}
51
52struct HeapItem<K, V> {
53    key: K,
54    /// Higher = wins. Negated so the min-heap pops it first on tie.
55    priority: i32,
56    source: usize,
57    value: V,
58}
59
60impl<K: Ord, V> PartialEq for HeapItem<K, V> {
61    fn eq(&self, other: &Self) -> bool {
62        self.key == other.key && self.source == other.source
63    }
64}
65impl<K: Ord, V> Eq for HeapItem<K, V> {}
66
67impl<K: Ord, V> PartialOrd for HeapItem<K, V> {
68    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
69        Some(self.cmp(other))
70    }
71}
72impl<K: Ord, V> Ord for HeapItem<K, V> {
73    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
74        // Key asc, then priority DESC (higher pops first off min-heap
75        // via Reverse), then source DESC (latest registration wins on
76        // priority tie).
77        self.key
78            .cmp(&other.key)
79            .then(other.priority.cmp(&self.priority))
80            .then(other.source.cmp(&self.source))
81    }
82}
83
84impl<K, V, I> PriorityMergeIterator<K, V, I>
85where
86    K: Ord,
87    I: Iterator<Item = PriorityEntry<K, V>>,
88{
89    pub fn new<S: IntoIterator<Item = PrioritySource<I>>>(sources: S) -> Self {
90        let sources: Vec<PrioritySource<I>> = sources.into_iter().collect();
91        let mut streams: Vec<I> = Vec::with_capacity(sources.len());
92        let mut priorities: Vec<i32> = Vec::with_capacity(sources.len());
93        let mut heap: BinaryHeap<Reverse<HeapItem<K, V>>> =
94            BinaryHeap::with_capacity(sources.len());
95        for (i, src) in sources.into_iter().enumerate() {
96            streams.push(src.stream);
97            priorities.push(src.priority);
98            if let Some(e) = streams[i].next() {
99                heap.push(Reverse(HeapItem {
100                    key: e.key,
101                    priority: src.priority,
102                    source: i,
103                    value: e.value,
104                }));
105            }
106        }
107        Self {
108            streams,
109            priorities,
110            heap,
111        }
112    }
113
114    fn advance(&mut self, source: usize) {
115        if let Some(e) = self.streams[source].next() {
116            let p = self.priorities[source];
117            self.heap.push(Reverse(HeapItem {
118                key: e.key,
119                priority: p,
120                source,
121                value: e.value,
122            }));
123        }
124    }
125}
126
127impl<K, V, I> Iterator for PriorityMergeIterator<K, V, I>
128where
129    K: Ord,
130    I: Iterator<Item = PriorityEntry<K, V>>,
131{
132    type Item = PriorityEntry<K, V>;
133
134    fn next(&mut self) -> Option<PriorityEntry<K, V>> {
135        let Reverse(HeapItem {
136            key: winning_key,
137            source,
138            value: winning_value,
139            ..
140        }) = self.heap.pop()?;
141        self.advance(source);
142        while let Some(Reverse(item)) = self.heap.peek() {
143            if item.key == winning_key {
144                let Reverse(item) = self.heap.pop().unwrap();
145                self.advance(item.source);
146            } else {
147                break;
148            }
149        }
150        Some(PriorityEntry {
151            key: winning_key,
152            value: winning_value,
153        })
154    }
155}
156
157#[cfg(test)]
158#[path = "priority_tests.rs"]
159mod tests;
subms_merge_iterator/features/priority.rs

subms_merge_iterator/features/
priority.rs
