refactor: optimize compaction's pick

Signed-off-by: luofucong <luofc@foxmail.com>
This commit is contained in:
luofucong
2026-05-28 19:15:32 +08:00
parent 123524474d
commit 3534a04ad9
5 changed files with 194 additions and 45 deletions

1
Cargo.lock generated
View File

@@ -8320,6 +8320,7 @@ dependencies = [
"datafusion-common",
"datafusion-expr",
"datatypes",
"derive_more",
"dotenv",
"either",
"futures",

View File

@@ -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

View File

@@ -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,130 @@ where
runs
}
pub(crate) fn find_sorted_runs_by_time_range<T>(items: &mut [T]) -> Vec<SortedRun<T>>
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<T: Item> {
i: usize,
#[partial_eq(skip)]
run: SortedRun<T>,
}
impl<T: Item> Run<T> {
fn new(i: usize, item: &T) -> Run<T> {
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<T: Item> PartialOrd for Run<T> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
/// Sort by run's `end` desc then `start` asc.
impl<T: Item> Ord for Run<T> {
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)
})
}
}
/// Wrapper around the `Run` above, to support sorting them by their creation sequence `i`.
#[derive(PartialEq, Eq)]
struct Wrapper<T: Item>(Run<T>);
impl<T: Item> PartialOrd for Wrapper<T> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<T: Item> Ord for Wrapper<T> {
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::<Run<T>>::new();
let mut runs_sort_by_index = BinaryHeap::<Wrapper<T>>::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<T: Item>(mut runs: Vec<SortedRun<T>>) -> Vec<T> {
@@ -599,6 +726,8 @@ mod tests {
expected_runs: &[Vec<(i64, i64)>],
) -> Vec<SortedRun<MockFile>> {
let mut files = build_items(ranges);
let mut files_clone = files.clone();
let runs = find_sorted_runs(&mut files);
let result_file_ranges: Vec<Vec<_>> = runs
@@ -606,6 +735,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<Vec<_>> = runs_by_time_range
.iter()
.map(|r| r.items.iter().map(|f| f.range()).collect())
.collect();
assert_eq!(&expected_runs, &results);
runs
}

View File

@@ -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<i64, Window>,
active_window: Option<i64>,
) -> Vec<CompactionOutput> {
let mut output = vec![];
for (window, files) in time_windows {
if files.files.is_empty() {
continue;
}
let find_inputs = |files: &Window,
windows: &BTreeMap<i64, Window>|
-> (Vec<FileGroup>, 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::<Vec<_>>();
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::<Vec<_>>()
{
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<Option<NonZeroU64>, 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::<Vec<_>>();
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<i64, Window>) -> 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 {

View File

@@ -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))]