From 252ce87447090f612b1ec997f4ef95622506567b Mon Sep 17 00:00:00 2001 From: luofucong Date: Thu, 28 May 2026 19:15:32 +0800 Subject: [PATCH] refactor: optimize compaction's pick Signed-off-by: luofucong --- Cargo.lock | 1 + src/mito2/Cargo.toml | 1 + src/mito2/src/compaction/run.rs | 139 +++++++++++++++++++++++++++++++ src/mito2/src/compaction/twcs.rs | 100 ++++++++++++---------- src/mito2/src/lib.rs | 1 + 5 files changed, 197 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 63ba289947..41b1a2ed2d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8320,6 +8320,7 @@ dependencies = [ "datafusion-common", "datafusion-expr", "datatypes", + "derive_more", "dotenv", "either", "futures", diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 3e3a18a24d..99e3439879 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -50,6 +50,7 @@ datafusion-common.workspace = true datafusion-expr.workspace = true datatypes.workspace = true dashmap.workspace = true +derive_more.workspace = true dotenv.workspace = true either.workspace = true futures.workspace = true diff --git a/src/mito2/src/compaction/run.rs b/src/mito2/src/compaction/run.rs index ecd19666e2..1ba4569021 100644 --- a/src/mito2/src/compaction/run.rs +++ b/src/mito2/src/compaction/run.rs @@ -15,6 +15,9 @@ //! This file contains code to find sorted runs in a set if ranged items and //! along with the best way to merge these items to satisfy the desired run count. +use std::cmp::Ordering; +use std::collections::BinaryHeap; + use bytes::{Buf, Bytes}; use common_base::BitVec; use common_base::readable_size::ReadableSize; @@ -423,6 +426,133 @@ where runs } +pub(crate) fn find_sorted_runs_by_time_range(items: &mut [T]) -> Vec> +where + T: Item, +{ + if items.is_empty() { + return vec![]; + } + sort_ranged_items(items); + + use derive_more::{Eq, PartialEq}; + + /// `SortedRun` with a creation sequence `i`. + #[derive(PartialEq, Eq)] + struct Run { + i: usize, + #[partial_eq(skip)] + run: SortedRun, + } + + impl Run { + fn new(i: usize, item: &T) -> Run { + let mut run = SortedRun::default(); + run.push_item(item.clone()); + Run { i, run } + } + + fn push_item(&mut self, item: &T) { + self.run.push_item(item.clone()); + } + } + + impl PartialOrd for Run { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + /// Sort by run's `end` desc then `start` asc. + impl Ord for Run { + fn cmp(&self, other: &Self) -> Ordering { + let l_run = &self.run; + let r_run = &other.run; + + // Safety: `start` and `end` must both exist because it's guaranteed that whenever a + // `Run` is created, an item is pushed into it immediately (see its `new` method above). + // And there are no other ways to create a `Run` beyond its `new` method in this + // function's scope. + let l_end = l_run.end.unwrap(); + let r_end = r_run.end.unwrap(); + r_end + .cmp(&l_end) + .then_with(|| { + let l_start = l_run.start.unwrap(); + let r_start = r_run.start.unwrap(); + l_start.cmp(&r_start) + }) + .then_with(|| self.i.cmp(&other.i)) + } + } + + /// Wrapper around the `Run` above, to support sorting them by their creation sequence `i`. + #[derive(PartialEq, Eq)] + struct Wrapper(Run); + + impl PartialOrd for Wrapper { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + impl Ord for Wrapper { + fn cmp(&self, other: &Self) -> Ordering { + other.0.i.cmp(&self.0.i) + } + } + + // Two heaps for finding a run that is both: + // 1. not overlapping with item's range, + // 2. and is created earliest, + // when iterating the items. + // + // Heap 1 (`runs_sorted_by_end`) is for storing the runs of which top has the minimal "end" + // just about to overlap with the current selected item. + // + // Heap 2 (`runs_sort_by_index`) is for storing the runs that all have "end"s non-overlap with + // the current selected item, and of which top is the earliest created run. + // + // The finding of a suitable run basically works like this: + // 1. moves the runs in heap 1 to heap 2, until the top is overlapping with the current item; + // 2. now heap 2 has all the runs that can accept the current item, pop its top; + // 3. the top is the earliest created run, push the current item; + // 4. because the run has changed, push it back to heap 1; + // 5. check the next item. Important: we don't need to push the runs in heap 2 to 1, because + // the items are sorted by "start". When checking the next item, heap 2's runs must all have + // "end"s smaller than next item's "start". + // + // Actually the heap 2 is only for aligning with the runs selection outcomes in the original + // `find_sorted_runs` implementation. If we just need the invariant that each run has the + // non-overlapping items, we can get rid of heap 2 and make the codes simpler. + + let mut runs_sort_by_end = BinaryHeap::>::new(); + let mut runs_sort_by_index = BinaryHeap::>::new(); + let mut i = 0; + + for item in items { + let (start, _) = item.range(); + + while let Some(run) = runs_sort_by_end.pop_if(|x| x.run.end.unwrap() <= start) { + runs_sort_by_index.push(Wrapper(run)); + } + + let Some(mut run) = runs_sort_by_index.pop() else { + i += 1; + runs_sort_by_end.push(Run::new(i, item)); + continue; + }; + + run.0.push_item(item); + runs_sort_by_end.push(run.0); + } + + let mut runs = runs_sort_by_end.into_vec(); + runs.extend(runs_sort_by_index.into_vec().into_iter().map(|x| x.0)); + runs.sort_unstable_by_key(|run| run.i); + runs.into_iter().map(|x| x.run).collect() +} + /// Finds a set of files with minimum penalty to merge that can reduce the total num of runs. /// The penalty of merging is defined as the size of all overlapping files between two runs. pub fn reduce_runs(mut runs: Vec>) -> Vec { @@ -599,6 +729,8 @@ mod tests { expected_runs: &[Vec<(i64, i64)>], ) -> Vec> { let mut files = build_items(ranges); + let mut files_clone = files.clone(); + let runs = find_sorted_runs(&mut files); let result_file_ranges: Vec> = runs @@ -606,6 +738,13 @@ mod tests { .map(|r| r.items.iter().map(|f| f.range()).collect()) .collect(); assert_eq!(&expected_runs, &result_file_ranges); + + let runs_by_time_range = find_sorted_runs_by_time_range(&mut files_clone); + let results: Vec> = runs_by_time_range + .iter() + .map(|r| r.items.iter().map(|f| f.range()).collect()) + .collect(); + assert_eq!(&expected_runs, &results); runs } diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index 87dcac7279..cfa44d5045 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -22,14 +22,15 @@ use common_telemetry::{debug, info}; use common_time::Timestamp; use common_time::timestamp::TimeUnit; use common_time::timestamp_millis::BucketAligned; +use rayon::prelude::*; use store_api::storage::RegionId; use crate::compaction::buckets::infer_time_bucket; use crate::compaction::compactor::CompactionRegion; use crate::compaction::picker::{Picker, PickerOutput}; use crate::compaction::run::{ - FileGroup, Item, Ranged, find_sorted_runs, merge_primary_key_ranges, merge_seq_files, - primary_key_ranges_overlap, reduce_runs, + FileGroup, Item, Ranged, find_sorted_runs, find_sorted_runs_by_time_range, + merge_primary_key_ranges, merge_seq_files, primary_key_ranges_overlap, reduce_runs, }; use crate::compaction::{CompactionOutput, get_expired_ssts}; use crate::sst::file::{FileHandle, Level, overlaps}; @@ -64,11 +65,10 @@ impl TwcsPicker { time_windows: &mut BTreeMap, active_window: Option, ) -> Vec { - let mut output = vec![]; - for (window, files) in time_windows { - if files.files.is_empty() { - continue; - } + let find_inputs = |files: &Window, + windows: &BTreeMap| + -> (Vec, bool) { + let window = &files.time_window; let mut files_to_merge: Vec<_> = files.files().cloned().collect(); // Filter out large files in append mode - they won't benefit from compaction @@ -88,13 +88,18 @@ impl TwcsPicker { ); } - let sorted_runs = find_sorted_runs(&mut files_to_merge); + let sorted_runs = if files_to_merge.len() < 1024 { + find_sorted_runs(&mut files_to_merge) + } else { + find_sorted_runs_by_time_range(&mut files_to_merge) + }; let found_runs = sorted_runs.len(); // We only remove deletion markers if we found less than 2 runs and not in append mode. // because after compaction there will be no overlapping files. - let filter_deleted = !files.overlapping && found_runs <= 2 && !self.append_mode; + let filter_deleted = + found_runs <= 2 && !self.append_mode && !window_has_overlap(files, windows); if found_runs == 0 { - continue; + return (vec![], filter_deleted); } let mut inputs = if found_runs > 1 { @@ -102,7 +107,7 @@ impl TwcsPicker { } else { let run = sorted_runs.last().unwrap(); if run.items().len() < self.trigger_file_num { - continue; + return (vec![], filter_deleted); } // no overlapping files, try merge small files merge_seq_files(run.items(), self.max_output_file_size) @@ -144,6 +149,26 @@ impl TwcsPicker { filter_deleted, &inputs, ); + } + (inputs, filter_deleted) + }; + + let mut output = vec![]; + let windows = time_windows + .values() + .filter(|w| !w.files.is_empty()) + .collect::>(); + let chunk_size = self.max_background_tasks.unwrap_or(windows.len()).max(1); + 'chunks: for chunk in windows.chunks(chunk_size) { + for (inputs, filter_deleted) in chunk + .par_iter() // parallelly calculate the inputs + .map(|window| find_inputs(window, time_windows)) + .collect::>() + { + if inputs.is_empty() { + continue; + } + output.push(CompactionOutput { output_level: LEVEL_COMPACTED, // always compact to l1 inputs: inputs.into_iter().flat_map(|fg| fg.into_files()).collect(), @@ -158,7 +183,7 @@ impl TwcsPicker { "Region ({:?}) compaction task size larger than max background tasks({}), remaining tasks discarded", region_id, max_background_tasks ); - break; + break 'chunks; } } } @@ -268,7 +293,6 @@ struct Window { // created from the same compaction task. files: HashMap, FileGroup>, time_window: i64, - overlapping: bool, primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } @@ -283,7 +307,6 @@ impl Window { end, files, time_window: 0, - overlapping: false, primary_key_range, } } @@ -346,37 +369,21 @@ fn assign_to_windows<'a>( } } } - if windows.is_empty() { - return BTreeMap::new(); - } + windows.into_iter().collect() +} - let mut windows = windows.into_values().collect::>(); - windows.sort_unstable_by(|l, r| l.start.cmp(&r.start).then(l.end.cmp(&r.end).reverse())); - - for idx in 0..windows.len() { - let lhs_range = windows[idx].range(); - for next_idx in idx + 1..windows.len() { - let rhs_range = windows[next_idx].range(); - if rhs_range.0 > lhs_range.1 { - break; - } - - let windows_overlap = overlaps(&lhs_range, &rhs_range) - && match ( - &windows[idx].primary_key_range, - &windows[next_idx].primary_key_range, - ) { - (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), +fn window_has_overlap(this: &Window, windows: &BTreeMap) -> bool { + windows + .values() + .filter(|that| this.time_window != that.time_window) + .any(|that| { + overlaps(&this.range(), &that.range()) && { + match (&this.primary_key_range, &that.primary_key_range) { + (Some(l), Some(r)) => primary_key_ranges_overlap(l, r), _ => true, - }; - if windows_overlap { - windows[idx].overlapping = true; - windows[next_idx].overlapping = true; + } } - } - } - - windows.into_iter().map(|w| (w.time_window, w)).collect() + }) } /// Finds the latest active writing window among all files. @@ -606,7 +613,8 @@ mod tests { for (expected_window, overlapping, window_files) in expected_files { let actual_window = windows.get(expected_window).unwrap(); - assert_eq!(*overlapping, actual_window.overlapping); + let actual_overlapping = window_has_overlap(actual_window, &windows); + assert_eq!(*overlapping, actual_overlapping); let mut file_ranges = actual_window .files .values() @@ -744,7 +752,8 @@ mod tests { let windows = assign_to_windows(files.iter(), 2); - assert!(!windows.get(&2).unwrap().overlapping); + let overlapping = window_has_overlap(windows.get(&2).unwrap(), &windows); + assert!(!overlapping); } #[test] @@ -773,7 +782,8 @@ mod tests { let windows = assign_to_windows(files.iter(), 2); - assert!(!windows.get(&4).unwrap().overlapping); + let overlapping = window_has_overlap(windows.get(&4).unwrap(), &windows); + assert!(!overlapping); } struct CompactionPickerTestCase { diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index 7d43685ded..a452e106c7 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -18,6 +18,7 @@ #![feature(debug_closure_helpers)] #![feature(duration_constructors)] +#![feature(binary_heap_pop_if)] #[cfg(any(test, feature = "test"))] #[cfg_attr(feature = "test", allow(unused))]