risingwave_storage/hummock/iterator/
backward_merge.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
15#[cfg(test)]
16mod test {
17    use crate::hummock::BackwardSstableIterator;
18    use crate::hummock::iterator::test_utils::{
19        TEST_KEYS_COUNT, default_builder_opt_for_test, gen_iterator_test_sstable_base,
20        iterator_test_key_of, iterator_test_value_of, mock_sstable_store,
21    };
22    use crate::hummock::iterator::{HummockIterator, MergeIterator};
23
24    #[tokio::test]
25    async fn test_backward_merge_basic() {
26        let sstable_store = mock_sstable_store().await;
27        let (table0, sstable_info_0) = gen_iterator_test_sstable_base(
28            0,
29            default_builder_opt_for_test(),
30            |x| x * 3 + 1,
31            sstable_store.clone(),
32            TEST_KEYS_COUNT,
33        )
34        .await;
35        let (table1, sstable_info_1) = gen_iterator_test_sstable_base(
36            1,
37            default_builder_opt_for_test(),
38            |x| x * 3 + 2,
39            sstable_store.clone(),
40            TEST_KEYS_COUNT,
41        )
42        .await;
43        let (table2, sstable_info_2) = gen_iterator_test_sstable_base(
44            2,
45            default_builder_opt_for_test(),
46            |x| x * 3 + 3,
47            sstable_store.clone(),
48            TEST_KEYS_COUNT,
49        )
50        .await;
51        let iters = vec![
52            BackwardSstableIterator::new(table0, sstable_store.clone(), &sstable_info_0),
53            BackwardSstableIterator::new(table1, sstable_store.clone(), &sstable_info_1),
54            BackwardSstableIterator::new(table2, sstable_store, &sstable_info_2),
55        ];
56
57        let mut mi = MergeIterator::new(iters);
58        let mut i = 3 * TEST_KEYS_COUNT;
59        mi.rewind().await.unwrap();
60        while mi.is_valid() {
61            let key = mi.key();
62            let val = mi.value();
63            assert_eq!(key, iterator_test_key_of(i).to_ref());
64            assert_eq!(
65                val.into_user_value().unwrap(),
66                iterator_test_value_of(i).as_slice()
67            );
68            i -= 1;
69            mi.next().await.unwrap();
70            if i == 0 {
71                assert!(!mi.is_valid());
72                break;
73            }
74        }
75    }
76
77    #[tokio::test]
78    async fn test_backward_merge_seek() {
79        let sstable_store = mock_sstable_store().await;
80        let (table0, sstable_info_0) = gen_iterator_test_sstable_base(
81            0,
82            default_builder_opt_for_test(),
83            |x| x * 3 + 1,
84            sstable_store.clone(),
85            TEST_KEYS_COUNT,
86        )
87        .await;
88        let (table1, sstable_info_1) = gen_iterator_test_sstable_base(
89            1,
90            default_builder_opt_for_test(),
91            |x| x * 3 + 2,
92            sstable_store.clone(),
93            TEST_KEYS_COUNT,
94        )
95        .await;
96        let (table2, sstable_info_2) = gen_iterator_test_sstable_base(
97            2,
98            default_builder_opt_for_test(),
99            |x| x * 3 + 3,
100            sstable_store.clone(),
101            TEST_KEYS_COUNT,
102        )
103        .await;
104        let iters = vec![
105            BackwardSstableIterator::new(table0, sstable_store.clone(), &sstable_info_0),
106            BackwardSstableIterator::new(table1, sstable_store.clone(), &sstable_info_1),
107            BackwardSstableIterator::new(table2, sstable_store, &sstable_info_2),
108        ];
109
110        let mut mi = MergeIterator::new(iters);
111
112        // right edge case
113        mi.seek(iterator_test_key_of(0).to_ref()).await.unwrap();
114        assert!(!mi.is_valid());
115
116        // normal case
117        mi.seek(iterator_test_key_of(TEST_KEYS_COUNT + 4).to_ref())
118            .await
119            .unwrap();
120        let k = mi.key();
121        let v = mi.value();
122        assert_eq!(
123            v.into_user_value().unwrap(),
124            iterator_test_value_of(TEST_KEYS_COUNT + 4).as_slice()
125        );
126        assert_eq!(k, iterator_test_key_of(TEST_KEYS_COUNT + 4).to_ref());
127
128        mi.seek(iterator_test_key_of(2 * TEST_KEYS_COUNT + 7).to_ref())
129            .await
130            .unwrap();
131        let k = mi.key();
132        let v = mi.value();
133        assert_eq!(
134            v.into_user_value().unwrap(),
135            iterator_test_value_of(2 * TEST_KEYS_COUNT + 7).as_slice()
136        );
137        assert_eq!(k, iterator_test_key_of(2 * TEST_KEYS_COUNT + 7).to_ref());
138
139        // left edge case
140        mi.seek(iterator_test_key_of(3 * TEST_KEYS_COUNT).to_ref())
141            .await
142            .unwrap();
143        let k = mi.key();
144        let v = mi.value();
145        assert_eq!(
146            v.into_user_value().unwrap(),
147            iterator_test_value_of(3 * TEST_KEYS_COUNT).as_slice()
148        );
149        assert_eq!(k, iterator_test_key_of(3 * TEST_KEYS_COUNT).to_ref());
150    }
151
152    #[tokio::test]
153    async fn test_backward_merge_invalidate_reset() {
154        let sstable_store = mock_sstable_store().await;
155        let (table0, sstable_info_0) = gen_iterator_test_sstable_base(
156            0,
157            default_builder_opt_for_test(),
158            |x| x + 1,
159            sstable_store.clone(),
160            TEST_KEYS_COUNT,
161        )
162        .await;
163        let (table1, sstable_info_1) = gen_iterator_test_sstable_base(
164            1,
165            default_builder_opt_for_test(),
166            |x| TEST_KEYS_COUNT + x + 1,
167            sstable_store.clone(),
168            TEST_KEYS_COUNT,
169        )
170        .await;
171        let iters = vec![
172            BackwardSstableIterator::new(table1, sstable_store.clone(), &sstable_info_1),
173            BackwardSstableIterator::new(table0, sstable_store, &sstable_info_0),
174        ];
175
176        let mut mi = MergeIterator::new(iters);
177
178        mi.rewind().await.unwrap();
179        let mut count = 0;
180        while mi.is_valid() {
181            count += 1;
182            mi.next().await.unwrap();
183        }
184        assert_eq!(count, TEST_KEYS_COUNT * 2);
185
186        mi.rewind().await.unwrap();
187        let mut count = 0;
188        while mi.is_valid() {
189            count += 1;
190            mi.next().await.unwrap();
191        }
192        assert_eq!(count, TEST_KEYS_COUNT * 2);
193    }
194}