Skip to main content

risingwave_meta/hummock/compaction/picker/
min_overlap_compaction_picker.rs

1// Copyright 2023 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::sync::Arc;
16
17use risingwave_hummock_sdk::level::{InputLevel, Levels};
18use risingwave_hummock_sdk::sstable_info::SstableInfo;
19use risingwave_pb::hummock::LevelType;
20
21use super::{CompactionInput, CompactionPicker, LocalPickerStatistic};
22use crate::hummock::compaction::overlap_strategy::OverlapStrategy;
23use crate::hummock::level_handler::LevelHandler;
24
25pub struct MinOverlappingPicker {
26    level: usize,
27    target_level: usize,
28    max_select_bytes: u64,
29    vnode_partition_count: u32,
30    overlap_strategy: Arc<dyn OverlapStrategy>,
31}
32
33impl MinOverlappingPicker {
34    pub fn new(
35        level: usize,
36        target_level: usize,
37        max_select_bytes: u64,
38        vnode_partition_count: u32,
39        overlap_strategy: Arc<dyn OverlapStrategy>,
40    ) -> MinOverlappingPicker {
41        MinOverlappingPicker {
42            level,
43            target_level,
44            max_select_bytes,
45            vnode_partition_count,
46            overlap_strategy,
47        }
48    }
49
50    pub fn pick_tables(
51        &self,
52        select_tables: &[SstableInfo],
53        target_tables: &[SstableInfo],
54        level_handlers: &[LevelHandler],
55    ) -> (Vec<SstableInfo>, Vec<SstableInfo>) {
56        let mut select_file_ranges = vec![];
57        // Both levels are ordered and non-overlapping, so overlap boundaries only
58        // move forward, including across pending source SSTs.
59        let mut previous_overlap = None;
60        for (idx, sst) in select_tables.iter().enumerate() {
61            if level_handlers[self.level].is_pending_compact(&sst.sst_id) {
62                continue;
63            }
64            let overlap_files_range = self.overlap_strategy.check_overlap_range_with_hint(
65                &sst.key_range,
66                target_tables,
67                previous_overlap,
68            );
69            previous_overlap = Some(overlap_files_range.clone());
70
71            if overlap_files_range.is_empty() {
72                return (vec![sst.clone()], vec![]);
73            }
74            select_file_ranges.push((idx, overlap_files_range));
75        }
76        select_file_ranges.retain(|(_, range)| {
77            let mut pending_compact = false;
78            for other in &target_tables[range.clone()] {
79                if level_handlers[self.target_level].is_pending_compact(&other.sst_id) {
80                    pending_compact = true;
81                    break;
82                }
83            }
84            !pending_compact
85        });
86
87        let mut min_score = u64::MAX;
88        let mut min_score_select_range = 0..0;
89        let mut min_score_target_range = 0..0;
90        let mut min_score_select_file_size = 0;
91        for left in 0..select_file_ranges.len() {
92            let mut select_file_size = 0;
93            let mut target_level_overlap_range = select_file_ranges[left].1.clone();
94            let mut total_file_size = 0;
95            for other in &target_tables[target_level_overlap_range.clone()] {
96                total_file_size += other.sst_size;
97            }
98            let start_idx = select_file_ranges[left].0;
99            let mut end_idx = start_idx + 1;
100            for (idx, range) in select_file_ranges.iter().skip(left) {
101                if select_file_size > self.max_select_bytes
102                    || *idx > end_idx
103                    || range.start >= target_level_overlap_range.end
104                {
105                    break;
106                }
107                select_file_size += select_tables[*idx].sst_size;
108                if range.end > target_level_overlap_range.end {
109                    for other in &target_tables[target_level_overlap_range.end..range.end] {
110                        total_file_size += other.sst_size;
111                    }
112                    target_level_overlap_range.end = range.end;
113                }
114                let score = (total_file_size * 100)
115                    .checked_div(select_file_size)
116                    .unwrap_or(total_file_size);
117                end_idx = idx + 1;
118                if score < min_score
119                    || (score == min_score && select_file_size < min_score_select_file_size)
120                {
121                    min_score = score;
122                    min_score_select_range = start_idx..end_idx;
123                    min_score_target_range = target_level_overlap_range.clone();
124                    min_score_select_file_size = select_file_size;
125                }
126            }
127        }
128        if min_score == u64::MAX {
129            return (vec![], vec![]);
130        }
131        let select_input_ssts = select_tables[min_score_select_range].to_vec();
132        let target_input_ssts = target_tables[min_score_target_range].to_vec();
133        (select_input_ssts, target_input_ssts)
134    }
135}
136
137impl CompactionPicker for MinOverlappingPicker {
138    fn pick_compaction(
139        &mut self,
140        levels: &Levels,
141        level_handlers: &[LevelHandler],
142        stats: &mut LocalPickerStatistic,
143    ) -> Option<CompactionInput> {
144        assert!(self.level > 0);
145        let (select_input_ssts, target_input_ssts) = self.pick_tables(
146            &levels.get_level(self.level).table_infos,
147            &levels.get_level(self.target_level).table_infos,
148            level_handlers,
149        );
150        if select_input_ssts.is_empty() {
151            stats.skip_by_pending_files += 1;
152            return None;
153        }
154        Some(CompactionInput {
155            select_input_size: select_input_ssts.iter().map(|sst| sst.sst_size).sum(),
156            target_input_size: target_input_ssts.iter().map(|sst| sst.sst_size).sum(),
157            total_file_count: (select_input_ssts.len() + target_input_ssts.len()) as u64,
158            input_levels: vec![
159                InputLevel {
160                    level_idx: self.level as u32,
161                    level_type: LevelType::Nonoverlapping,
162                    table_infos: select_input_ssts,
163                },
164                InputLevel {
165                    level_idx: self.target_level as u32,
166                    level_type: LevelType::Nonoverlapping,
167                    table_infos: target_input_ssts,
168                },
169            ],
170            target_level: self.target_level,
171            vnode_partition_count: self.vnode_partition_count,
172            ..Default::default()
173        })
174    }
175}
176
177#[cfg(test)]
178pub mod tests {
179    use risingwave_hummock_sdk::key_range::KeyRangeCommon;
180    use risingwave_hummock_sdk::level::Level;
181
182    use super::*;
183    use crate::hummock::compaction::overlap_strategy::RangeOverlapStrategy;
184    use crate::hummock::compaction::selector::tests::{
185        generate_l0_nonoverlapping_sublevels, generate_table,
186    };
187
188    fn sized_table(id: u64, left: usize, right: usize, size: u64) -> SstableInfo {
189        let mut table =
190            crate::hummock::compaction::selector::tests::generate_table_impl(id, 1, left, right, 1);
191        table.sst_size = size;
192        // Fixed-width keys preserve byte ordering beyond the shared helper's five-digit range.
193        let key = |idx: usize| {
194            risingwave_hummock_sdk::key::FullKey::for_test(
195                1.into(),
196                (idx as u64).to_be_bytes(),
197                risingwave_common::util::epoch::test_epoch(1),
198            )
199            .encode()
200        };
201        table.key_range.left = key(left).into();
202        table.key_range.right = key(right).into();
203        table.into()
204    }
205
206    #[test]
207    fn test_overlap_hint_boundaries_and_sparse_queries() {
208        let strategy = RangeOverlapStrategy::default();
209        let targets: Vec<SstableInfo> = (0..512)
210            .map(|i| {
211                let mut table = (*sized_table(i as u64, i * 4, i * 4 + 2, 1)).clone();
212                table.key_range.right_exclusive = i % 2 == 0;
213                table.into()
214            })
215            .collect();
216        for targets in [targets.as_slice(), &[]] {
217            for exclusive in [false, true] {
218                let mut previous = None;
219                // Empty overlaps, touching endpoints, repeated bounds, large gaps,
220                // and exhaustion of the target slice all preserve exact indices.
221                for (left, right) in [
222                    (0, 1),
223                    (1, 2),
224                    (2, 3),
225                    (2, 3),
226                    (3, 4),
227                    (4, 6),
228                    (6, 8),
229                    (8, 16),
230                    (1000, 1002),
231                    (2046, 2048),
232                    (4096, 4097),
233                ] {
234                    let mut query = sized_table(1000, left, right, 1).key_range.clone();
235                    query.right_exclusive = exclusive;
236                    let mut info = strategy.create_overlap_info();
237                    info.update(&query);
238                    let expected = info.check_multiple_overlap(targets);
239                    let actual = strategy.check_overlap_range_with_hint(&query, targets, previous);
240                    assert_eq!(actual, expected);
241                    assert_eq!(
242                        strategy.check_overlap_range_with_hint(&query, targets, None),
243                        expected
244                    );
245                    previous = Some(actual);
246                }
247            }
248        }
249    }
250
251    // Run with cargo test --release -p risingwave_meta bench_min_overlap_leveled -- --ignored --nocapture --test-threads=1.
252    #[test]
253    #[ignore = "microbenchmark"]
254    fn bench_min_overlap_leveled() {
255        use std::hint::black_box;
256        use std::time::Duration;
257
258        let mut c = criterion::Criterion::default()
259            .sample_size(20)
260            .warm_up_time(Duration::from_millis(300))
261            .measurement_time(Duration::from_millis(600))
262            .without_plots();
263        let mib = 1024 * 1024_u64;
264        // Match the default min-overlap source-byte limit, including its pre-add check.
265        let limit = 256 * mib;
266        for (name, n, source_mib, target_mib, shift, varied, pending) in [
267            ("aligned_32m_10k", 10000_usize, 32, 32, 0, false, false),
268            ("staggered_32m_10k", 10000, 32, 32, 5, false, false),
269            ("staggered_32m_100k", 100000, 32, 32, 5, false, false),
270            ("staggered_32m_64m_10k", 10000, 32, 64, 5, false, false),
271            ("fragmented_1m_10k", 10000, 1, 1, 5, false, false),
272            ("varied_24_40m_10k", 10000, 32, 32, 5, true, false),
273            ("pending_32m_10k", 10000, 32, 32, 5, false, true),
274        ] {
275            // Width is proportional to bytes within each level; lower density is 10x.
276            let mut end = shift;
277            let source: Vec<_> = (0..n)
278                .map(|i| {
279                    let size = if varied {
280                        [24, 32, 40][i % 3]
281                    } else {
282                        source_mib
283                    };
284                    let start = end;
285                    end += (size * 100) as usize;
286                    sized_table(i as u64, start, end - 1, size * mib)
287                })
288                .collect();
289            let target_width = (target_mib * 10) as usize;
290            let target: Vec<_> = (0..end.div_ceil(target_width))
291                .map(|i| {
292                    sized_table(
293                        1000000 + i as u64,
294                        i * target_width,
295                        (i + 1) * target_width - 1,
296                        target_mib * mib,
297                    )
298                })
299                .collect();
300            for tables in [&source, &target] {
301                assert!(tables.windows(2).all(|pair| {
302                    pair[0]
303                        .key_range
304                        .compare_right_with(&pair[1].key_range.left)
305                        == std::cmp::Ordering::Less
306                }));
307            }
308            let strategy = Arc::new(RangeOverlapStrategy::default());
309            let ranges: Vec<_> = source
310                .iter()
311                .map(|sst| {
312                    let mut info = strategy.create_overlap_info();
313                    info.update(&sst.key_range);
314                    info.check_multiple_overlap(&target)
315                })
316                .collect();
317            assert!(ranges.iter().all(|r| !r.is_empty()));
318            // Adjacent source SSTs overlap different lower ranges in these leveled layouts.
319            assert!(ranges.windows(2).all(|pair| pair[0] != pair[1]));
320            let mut handlers: Vec<_> = (0..3).map(LevelHandler::new).collect();
321            if pending {
322                for i in (99..target.len()).step_by(100) {
323                    handlers[2].test_add_pending_sst(target[i].sst_id, 1);
324                }
325            }
326            let picker = MinOverlappingPicker::new(1, 2, limit, 0, strategy);
327            let (selected, overlapped) = picker.pick_tables(&source, &target, &handlers);
328            assert!(!selected.is_empty() && !overlapped.is_empty());
329            eprintln!(
330                "{name}: source={} target={} selected={} overlapped={}",
331                source.len(),
332                target.len(),
333                selected.len(),
334                overlapped.len()
335            );
336            c.bench_function(&format!("min_overlap_leveled/{name}"), |b| {
337                b.iter(|| {
338                    black_box(picker.pick_tables(
339                        black_box(&source),
340                        black_box(&target),
341                        black_box(&handlers),
342                    ))
343                })
344            });
345        }
346        c.final_summary();
347    }
348
349    #[test]
350    fn test_compact_l1() {
351        let mut picker =
352            MinOverlappingPicker::new(1, 2, 10000, 0, Arc::new(RangeOverlapStrategy::default()));
353        let levels = vec![
354            Level {
355                level_idx: 1,
356                level_type: LevelType::Nonoverlapping,
357                table_infos: vec![
358                    generate_table(0, 1, 0, 100, 1),
359                    generate_table(1, 1, 101, 200, 1),
360                    generate_table(2, 1, 222, 300, 1),
361                ],
362                ..Default::default()
363            },
364            Level {
365                level_idx: 2,
366                level_type: LevelType::Nonoverlapping,
367                table_infos: vec![
368                    generate_table(4, 1, 0, 100, 1),
369                    generate_table(5, 1, 101, 150, 1),
370                    generate_table(6, 1, 151, 201, 1),
371                    generate_table(7, 1, 501, 800, 1),
372                    generate_table(8, 2, 301, 400, 1),
373                ],
374                ..Default::default()
375            },
376        ];
377        let levels = Levels {
378            levels,
379            l0: generate_l0_nonoverlapping_sublevels(vec![]),
380            ..Default::default()
381        };
382        let mut level_handlers = vec![
383            LevelHandler::new(0),
384            LevelHandler::new(1),
385            LevelHandler::new(2),
386        ];
387
388        // pick a non-overlapping files. It means that this file could be trivial move to next
389        // level.
390        let mut local_stats = LocalPickerStatistic::default();
391        let ret = picker
392            .pick_compaction(&levels, &level_handlers, &mut local_stats)
393            .unwrap();
394        assert_eq!(ret.input_levels[0].level_idx, 1);
395        assert_eq!(ret.target_level, 2);
396        assert_eq!(ret.input_levels[0].table_infos.len(), 1);
397        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 2);
398        assert_eq!(ret.input_levels[1].table_infos.len(), 0);
399        ret.add_pending_task(0, &mut level_handlers);
400
401        let ret = picker
402            .pick_compaction(&levels, &level_handlers, &mut local_stats)
403            .unwrap();
404        assert_eq!(ret.input_levels[0].level_idx, 1);
405        assert_eq!(ret.target_level, 2);
406        assert_eq!(ret.input_levels[0].table_infos.len(), 1);
407        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 0);
408        assert_eq!(ret.input_levels[1].table_infos.len(), 1);
409        assert_eq!(ret.input_levels[1].table_infos[0].sst_id, 4);
410        ret.add_pending_task(1, &mut level_handlers);
411
412        let ret = picker
413            .pick_compaction(&levels, &level_handlers, &mut local_stats)
414            .unwrap();
415        assert_eq!(ret.input_levels[0].table_infos.len(), 1);
416        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 1);
417        assert_eq!(ret.input_levels[1].table_infos.len(), 2);
418        assert_eq!(ret.input_levels[1].table_infos[0].sst_id, 5);
419    }
420
421    #[test]
422    fn test_expand_l1_files() {
423        let mut picker =
424            MinOverlappingPicker::new(1, 2, 10000, 0, Arc::new(RangeOverlapStrategy::default()));
425        let levels = vec![
426            Level {
427                level_idx: 1,
428                level_type: LevelType::Nonoverlapping,
429                table_infos: vec![
430                    generate_table(0, 1, 50, 99, 2),
431                    generate_table(1, 1, 100, 149, 2),
432                    generate_table(2, 1, 150, 249, 2),
433                ],
434                ..Default::default()
435            },
436            Level {
437                level_idx: 2,
438                level_type: LevelType::Nonoverlapping,
439                table_infos: vec![
440                    generate_table(4, 1, 50, 199, 1),
441                    generate_table(5, 1, 200, 399, 1),
442                ],
443                ..Default::default()
444            },
445        ];
446        let levels = Levels {
447            levels,
448            l0: generate_l0_nonoverlapping_sublevels(vec![]),
449            ..Default::default()
450        };
451        let levels_handler = vec![
452            LevelHandler::new(0),
453            LevelHandler::new(1),
454            LevelHandler::new(2),
455        ];
456
457        // pick a non-overlapping files. It means that this file could be trivial move to next
458        // level.
459        let ret = picker
460            .pick_compaction(
461                &levels,
462                &levels_handler,
463                &mut LocalPickerStatistic::default(),
464            )
465            .unwrap();
466        assert_eq!(ret.input_levels[0].level_idx, 1);
467        assert_eq!(ret.input_levels[1].level_idx, 2);
468
469        assert_eq!(ret.input_levels[0].table_infos.len(), 2);
470        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 0);
471        assert_eq!(ret.input_levels[0].table_infos[1].sst_id, 1);
472
473        assert_eq!(ret.input_levels[1].table_infos.len(), 1);
474        assert_eq!(ret.input_levels[1].table_infos[0].sst_id, 4);
475    }
476
477    #[test]
478    fn test_trivial_move_bug() {
479        let levels = [
480            Level {
481                level_idx: 1,
482                level_type: LevelType::Nonoverlapping,
483                table_infos: vec![generate_table(0, 1, 400, 500, 2)],
484                total_file_size: 100,
485                ..Default::default()
486            },
487            Level {
488                level_idx: 2,
489                level_type: LevelType::Nonoverlapping,
490                table_infos: vec![
491                    generate_table(1, 1, 100, 200, 1),
492                    generate_table(2, 1, 600, 700, 1),
493                ],
494                total_file_size: 200,
495                ..Default::default()
496            },
497            Level {
498                level_idx: 3,
499                level_type: LevelType::Nonoverlapping,
500                table_infos: vec![
501                    generate_table(3, 1, 100, 300, 2),
502                    generate_table(4, 1, 600, 800, 1),
503                ],
504                total_file_size: 400,
505                ..Default::default()
506            },
507        ];
508
509        let levels_handlers = vec![
510            LevelHandler::new(0),
511            LevelHandler::new(1),
512            LevelHandler::new(2),
513            LevelHandler::new(3),
514        ];
515        // no limit
516        let picker =
517            MinOverlappingPicker::new(2, 3, 1000, 0, Arc::new(RangeOverlapStrategy::default()));
518        let (select_files, target_files) = picker.pick_tables(
519            &levels[1].table_infos,
520            &levels[2].table_infos,
521            &levels_handlers,
522        );
523        let overlap_strategy = Arc::new(RangeOverlapStrategy::default());
524        let mut overlap_info = overlap_strategy.create_overlap_info();
525        for sst in &select_files {
526            overlap_info.update(&sst.key_range);
527        }
528        let range = overlap_info.check_multiple_overlap(&levels[0].table_infos);
529        assert!(range.is_empty());
530        assert_eq!(select_files.len(), 1);
531        assert_eq!(target_files.len(), 1);
532    }
533}