risingwave_hummock_sdk/
key_range.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cmp;
16
17use bytes::Bytes;
18
19use super::key_cmp::KeyComparator;
20use crate::key::{FullKey, UserKey};
21
22#[derive(PartialEq, Eq, Clone, Debug, Default)]
23pub struct KeyRange {
24    pub left: Bytes,
25    pub right: Bytes,
26    pub right_exclusive: bool,
27}
28
29impl KeyRange {
30    pub fn new(left: Bytes, right: Bytes) -> Self {
31        Self {
32            left,
33            right,
34            right_exclusive: false,
35        }
36    }
37
38    pub fn inf() -> Self {
39        Self {
40            left: Bytes::new(),
41            right: Bytes::new(),
42            right_exclusive: false,
43        }
44    }
45
46    #[inline]
47    fn start_bound_inf(&self) -> bool {
48        self.left.is_empty()
49    }
50
51    #[inline]
52    fn end_bound_inf(&self) -> bool {
53        self.right.is_empty()
54    }
55
56    #[inline]
57    pub fn inf_key_range(&self) -> bool {
58        self.start_bound_inf() && self.end_bound_inf()
59    }
60}
61
62pub trait KeyRangeCommon {
63    fn full_key_overlap(&self, other: &Self) -> bool;
64    fn full_key_extend(&mut self, other: &Self);
65    fn sstable_overlap(&self, other: &Self) -> bool;
66    fn compare_right_with(&self, full_key: &[u8]) -> std::cmp::Ordering {
67        self.compare_right_with_user_key(FullKey::decode(full_key).user_key)
68    }
69    fn compare_right_with_user_key(&self, ukey: UserKey<&[u8]>) -> std::cmp::Ordering;
70}
71
72#[macro_export]
73macro_rules! impl_key_range_common {
74    ($T:ty) => {
75        impl KeyRangeCommon for $T {
76            fn full_key_overlap(&self, other: &Self) -> bool {
77                (self.end_bound_inf()
78                    || other.start_bound_inf()
79                    || KeyComparator::compare_encoded_full_key(&self.right, &other.left)
80                        != cmp::Ordering::Less)
81                    && (other.end_bound_inf()
82                        || self.start_bound_inf()
83                        || KeyComparator::compare_encoded_full_key(&other.right, &self.left)
84                            != cmp::Ordering::Less)
85            }
86
87            fn full_key_extend(&mut self, other: &Self) {
88                if !self.start_bound_inf()
89                    && (other.start_bound_inf()
90                        || KeyComparator::compare_encoded_full_key(&other.left, &self.left)
91                            == cmp::Ordering::Less)
92                {
93                    self.left = other.left.clone();
94                }
95                if !self.end_bound_inf()
96                    && (other.end_bound_inf()
97                        || KeyComparator::compare_encoded_full_key(&other.right, &self.right)
98                            == cmp::Ordering::Greater)
99                {
100                    self.right = other.right.clone();
101                    self.right_exclusive = other.right_exclusive;
102                }
103            }
104
105            fn sstable_overlap(&self, other: &Self) -> bool {
106                (self.end_bound_inf()
107                    || other.start_bound_inf()
108                    || self.compare_right_with(&other.left) != std::cmp::Ordering::Less)
109                    && (other.end_bound_inf()
110                        || self.start_bound_inf()
111                        || other.compare_right_with(&self.left) != std::cmp::Ordering::Less)
112            }
113
114            fn compare_right_with_user_key(
115                &self,
116                ukey: $crate::key::UserKey<&[u8]>,
117            ) -> std::cmp::Ordering {
118                use $crate::key::FullKey;
119                let ret = FullKey::decode(&self.right).user_key.cmp(&ukey);
120                if ret == cmp::Ordering::Equal && self.right_exclusive {
121                    cmp::Ordering::Less
122                } else {
123                    ret
124                }
125            }
126        }
127    };
128}
129
130#[macro_export]
131macro_rules! key_range_cmp {
132    ($left:expr, $right:expr) => {{
133        let ret = if $left.start_bound_inf() && $right.start_bound_inf() {
134            cmp::Ordering::Equal
135        } else if !$left.start_bound_inf() && !$right.start_bound_inf() {
136            KeyComparator::compare_encoded_full_key(&$left.left, &$right.left)
137        } else if $left.left.is_empty() {
138            cmp::Ordering::Less
139        } else {
140            cmp::Ordering::Greater
141        };
142        if ret != cmp::Ordering::Equal {
143            return ret;
144        }
145        if $left.end_bound_inf() && $right.end_bound_inf() {
146            cmp::Ordering::Equal
147        } else if !$left.end_bound_inf() && !$right.end_bound_inf() {
148            KeyComparator::compare_encoded_full_key(&$left.right, &$right.right)
149        } else if $left.end_bound_inf() {
150            cmp::Ordering::Greater
151        } else {
152            cmp::Ordering::Less
153        }
154    }};
155}
156
157impl_key_range_common!(KeyRange);
158
159impl Ord for KeyRange {
160    fn cmp(&self, other: &Self) -> cmp::Ordering {
161        key_range_cmp!(self, other)
162    }
163}
164
165impl PartialOrd for KeyRange {
166    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
167        Some(self.cmp(other))
168    }
169}
170
171impl From<KeyRange> for risingwave_pb::hummock::KeyRange {
172    fn from(kr: KeyRange) -> Self {
173        risingwave_pb::hummock::KeyRange {
174            left: kr.left.to_vec(),
175            right: kr.right.to_vec(),
176            right_exclusive: kr.right_exclusive,
177        }
178    }
179}
180
181impl From<&risingwave_pb::hummock::KeyRange> for KeyRange {
182    fn from(kr: &risingwave_pb::hummock::KeyRange) -> Self {
183        KeyRange {
184            left: Bytes::copy_from_slice(&kr.left),
185            right: Bytes::copy_from_slice(&kr.right),
186            right_exclusive: kr.right_exclusive,
187        }
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use crate::key::key_with_epoch;
195
196    #[test]
197    fn test_key_range_compare() {
198        let a1_slice = &key_with_epoch(Vec::from("a"), 1);
199        let a2_slice = &key_with_epoch(Vec::from("a"), 2);
200        let b1_slice = &key_with_epoch(Vec::from("b"), 1);
201        let a1 = Bytes::copy_from_slice(a1_slice);
202        let a2 = Bytes::copy_from_slice(a2_slice);
203        let b1 = Bytes::copy_from_slice(b1_slice);
204        assert_eq!(
205            KeyRange::new(a1.clone(), a2.clone()).cmp(&KeyRange::new(a2.clone(), a2.clone())),
206            cmp::Ordering::Greater
207        );
208        assert_eq!(
209            KeyRange::new(a1.clone(), a2).partial_cmp(&KeyRange::new(a1, b1)),
210            Some(cmp::Ordering::Less)
211        );
212    }
213}
risingwave_hummock_sdk/key_range.rs

risingwave_hummock_sdk/
key_range.rs
