1use std::collections::hash_map::Entry;
16use std::collections::{BTreeMap, HashMap, HashSet};
17use std::fmt::Debug;
18use std::sync::{Arc, LazyLock};
19
20use common_base::readable_size::ReadableSize;
21use common_telemetry::{debug, info};
22use common_time::Timestamp;
23use common_time::range::TimestampRange;
24use common_time::timestamp::TimeUnit;
25use common_time::timestamp_millis::BucketAligned;
26use snafu::ResultExt;
27use store_api::storage::RegionId;
28
29use crate::compaction::CompactionOutput;
30use crate::compaction::buckets::infer_time_bucket;
31use crate::compaction::compactor::CompactionRegion;
32use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts};
33use crate::compaction::run::{
34 Ranged, SortedRun, find_sorted_runs, find_sorted_runs_by_time_range, merge_primary_key_ranges,
35 primary_key_ranges_overlap,
36};
37use crate::error::{JoinSnafu, Result};
38use crate::sst::file::{FileHandle, Level, overlaps};
39use crate::sst::version::LevelMeta;
40
41const LEVEL_COMPACTED: Level = 1;
42
43const MAX_L1_L0_ROW_RATIO: usize = 2;
45
46const DEFAULT_MAX_INPUT_FILES: usize = 32;
48
49const MAX_INPUT_FILES_ENV: &str = "GREPTIME_TWCS_MAX_INPUT_FILES";
50
51static MAX_INPUT_FILES: LazyLock<usize> = LazyLock::new(|| {
54 let env_value = std::env::var(MAX_INPUT_FILES_ENV).ok();
55 parse_max_input_files(env_value.as_deref())
56});
57
58fn parse_max_input_files(env_value: Option<&str>) -> usize {
59 env_value
60 .and_then(|env_value| env_value.parse().ok())
61 .filter(|max_input_files| *max_input_files >= 2)
62 .unwrap_or(DEFAULT_MAX_INPUT_FILES)
63}
64
65#[derive(Clone, Debug)]
68pub struct TwcsPicker {
69 pub trigger_file_num: usize,
71 pub time_window_seconds: Option<i64>,
73 pub max_output_file_size: Option<u64>,
76 pub append_mode: bool,
78 pub max_background_tasks: Option<usize>,
80 pub(crate) time_range: Option<TimestampRange>,
82}
83
84impl TwcsPicker {
85 async fn build_output_with_time_range(
86 &self,
87 region_id: RegionId,
88 time_windows: BTreeMap<i64, Window>,
89 active_window: Option<i64>,
90 time_window_size: Option<i64>,
91 ) -> Result<Vec<CompactionOutput>> {
92 let mut output = vec![];
93 let windows = time_windows
94 .values()
95 .rev()
96 .filter(|window| {
97 !window.files.is_empty()
98 && self.time_range.as_ref().is_none_or(|time_range| {
99 time_window_size.is_none_or(|time_window_size| {
100 time_window_intersects_range(
101 window.time_window,
102 time_window_size,
103 time_range,
104 )
105 })
106 })
107 })
108 .map(|window| window.time_window)
109 .collect::<Vec<_>>();
110 let time_windows = Arc::new(time_windows);
111 let chunk_size = self.max_background_tasks.unwrap_or(windows.len()).max(1);
112 'chunks: for chunk in windows.chunks(chunk_size) {
113 let mut handles = Vec::with_capacity(chunk.len());
114 for window in chunk {
115 let picker = self.clone();
116 let time_windows = time_windows.clone();
117 let window = *window;
118 handles.push(common_runtime::spawn_blocking_compact(move || {
119 time_windows.get(&window).map(|window| {
120 picker.find_inputs(region_id, active_window, window, &time_windows)
121 })
122 }));
123 tokio::task::yield_now().await;
124 }
125 for result in futures::future::join_all(handles).await {
126 let Some((inputs, filter_deleted)) = result.context(JoinSnafu)? else {
127 continue;
128 };
129 if inputs.is_empty() {
130 continue;
131 }
132
133 output.push(CompactionOutput {
134 output_level: LEVEL_COMPACTED, inputs,
136 filter_deleted,
137 output_time_range: None, });
139
140 if let Some(max_background_tasks) = self.max_background_tasks
141 && output.len() >= max_background_tasks
142 {
143 debug!(
144 "Region ({:?}) compaction task size larger than max background tasks({}), remaining tasks discarded",
145 region_id, max_background_tasks
146 );
147 break 'chunks;
148 }
149 }
150 }
151 Ok(output)
152 }
153
154 fn find_inputs(
155 &self,
156 region_id: RegionId,
157 active_window: Option<i64>,
158 files: &Window,
159 windows: &BTreeMap<i64, Window>,
160 ) -> (Vec<FileHandle>, bool) {
161 if files.files.len() < self.trigger_file_num {
162 return (vec![], false);
163 }
164
165 let window = &files.time_window;
166 let mut files_to_merge: Vec<_> = files.files().cloned().collect();
167
168 if self.append_mode
170 && let Some(max_size) = self.max_output_file_size
171 {
172 let (kept_files, ignored_files) = files_to_merge
173 .into_iter()
174 .partition(|file| file.size() <= max_size);
175 files_to_merge = kept_files;
176 if !ignored_files.is_empty() {
177 info!(
178 "Skipped {} large files in append mode for region {}, window {}, max_size: {}",
179 ignored_files.len(),
180 region_id,
181 window,
182 max_size
183 );
184 }
185 }
186
187 let (mut l0_files, l1_files): (Vec<_>, Vec<_>) = files_to_merge
188 .into_iter()
189 .partition(|file| file.level() == 0);
190 let num_l0_files = l0_files.len();
191 let num_l1_files = l1_files.len();
192 let (inputs, found_runs) = if num_l0_files >= self.trigger_file_num {
196 let l0_pick =
197 pick_candidate_files(l0_files, self.max_output_file_size, pick_count_first);
198 if l0_pick.0.is_empty() && num_l1_files >= self.trigger_file_num {
199 pick_candidate_files(l1_files, self.max_output_file_size, pick_count_first)
200 } else {
201 l0_pick
202 }
203 } else if num_l1_files >= self.trigger_file_num {
204 pick_candidate_files(l1_files, self.max_output_file_size, pick_count_first)
205 } else {
206 l0_files.extend(l1_files);
207 let picker = if num_l0_files > 0 && num_l1_files > 0 {
208 pick_mixed_count_first
209 } else {
210 pick_count_first
211 };
212 pick_candidate_files(l0_files, self.max_output_file_size, picker)
213 };
214 let filter_deleted = !self.append_mode
215 && !window_has_overlap(files, windows)
216 && !selected_overlaps_unselected(&inputs, files);
217
218 if inputs.len() > 1 {
219 log_pick_result(
221 region_id,
222 *window,
223 active_window,
224 found_runs,
225 files.files.len(),
226 self.max_output_file_size,
227 filter_deleted,
228 &inputs,
229 );
230 }
231 (inputs, filter_deleted)
232 }
233}
234
235fn pick_candidate_files(
236 mut files: Vec<FileHandle>,
237 max_output_file_size: Option<u64>,
238 picker: fn(Vec<SortedRun<FileHandle>>, Option<u64>) -> Vec<FileHandle>,
239) -> (Vec<FileHandle>, usize) {
240 let sorted_runs = if files.len() < 1024 {
241 find_sorted_runs(&mut files)
242 } else {
243 find_sorted_runs_by_time_range(&mut files)
244 };
245 let found_runs = sorted_runs.len();
246 (picker(sorted_runs, max_output_file_size), found_runs)
247}
248
249#[derive(Debug)]
250struct OrderedFile<'a> {
251 file: &'a FileHandle,
252 run_id: usize,
253 position_in_run: usize,
254}
255
256#[derive(Debug, Default)]
259struct Candidate {
260 num_files: usize,
262 total_size: usize,
264 largest_file_size: usize,
266 overlap_participants: usize,
269 l0_rows: usize,
271 l1_rows: usize,
273 has_l0: bool,
274 has_l1: bool,
275 has_unknown_rows: bool,
277}
278
279impl Candidate {
280 fn absorb(
284 &mut self,
285 file: &OrderedFile,
286 preceding: &[OrderedFile],
287 participations: &mut Vec<bool>,
288 ) {
289 self.num_files += 1;
290 let file_size = file.file.size() as usize;
291 self.total_size += file_size;
292 self.largest_file_size = self.largest_file_size.max(file_size);
293 self.absorb_level_rows(file.file);
294
295 let mut participates = false;
296 for (offset, other) in preceding.iter().enumerate() {
297 if file.run_id != other.run_id && file.file.overlap_inclusive(other.file) {
298 if !participations[offset] {
299 participations[offset] = true;
300 self.overlap_participants += 1;
301 }
302 participates = true;
303 }
304 }
305 participations.push(participates);
306 if participates {
307 self.overlap_participants += 1;
308 }
309 }
310
311 fn absorb_level_rows(&mut self, file: &FileHandle) {
312 let num_rows = file.num_rows();
313 self.has_unknown_rows |= num_rows == 0;
314 if file.level() == 0 {
315 self.has_l0 = true;
316 self.l0_rows = self.l0_rows.saturating_add(num_rows);
317 } else {
318 self.has_l1 = true;
319 self.l1_rows = self.l1_rows.saturating_add(num_rows);
320 }
321 }
322
323 fn predicted_output_files(&self, max_output_file_size: Option<u64>) -> usize {
328 match max_output_file_size {
329 Some(max) if max > 0 => self.total_size.div_ceil(max as usize).max(1),
330 _ => 1,
331 }
332 }
333
334 fn file_reduction(&self, max_output_file_size: Option<u64>) -> usize {
338 self.num_files
339 .saturating_sub(self.predicted_output_files(max_output_file_size))
340 }
341
342 fn is_balanced(&self) -> bool {
345 self.largest_file_size <= self.total_size - self.largest_file_size
346 }
347
348 fn has_balanced_level_rows(&self) -> bool {
351 !self.has_l0
352 || !self.has_l1
353 || self.has_unknown_rows
354 || self.l1_rows <= self.l0_rows.saturating_mul(MAX_L1_L0_ROW_RATIO)
355 }
356
357 fn has_mixed_levels(&self) -> bool {
358 self.has_l0 && self.has_l1
359 }
360
361 fn makes_progress(&self, max_output_file_size: Option<u64>) -> bool {
366 self.file_reduction(max_output_file_size) > 0 || self.overlap_participants > 0
367 }
368
369 fn score(&self, max_output_file_size: Option<u64>) -> CandidateScore {
370 CandidateScore {
371 file_reduction: self.file_reduction(max_output_file_size),
372 overlap_participants: self.overlap_participants,
373 total_size: self.total_size,
374 }
375 }
376}
377
378#[derive(Debug)]
382struct CandidateScore {
383 file_reduction: usize,
384 overlap_participants: usize,
385 total_size: usize,
386}
387
388impl CandidateScore {
389 fn is_better_than(&self, other: &Self) -> bool {
390 self.file_reduction
391 .cmp(&other.file_reduction)
392 .then_with(|| self.overlap_participants.cmp(&other.overlap_participants))
393 .then_with(|| other.total_size.cmp(&self.total_size))
394 .is_gt()
395 }
396}
397
398fn pick_count_first(
413 sorted_runs: Vec<SortedRun<FileHandle>>,
414 max_output_file_size: Option<u64>,
415) -> Vec<FileHandle> {
416 pick_count_first_where(sorted_runs, max_output_file_size, |_| true)
417}
418
419fn pick_mixed_count_first(
420 sorted_runs: Vec<SortedRun<FileHandle>>,
421 max_output_file_size: Option<u64>,
422) -> Vec<FileHandle> {
423 pick_count_first_where(
424 sorted_runs,
425 max_output_file_size,
426 Candidate::has_mixed_levels,
427 )
428}
429
430fn pick_count_first_where(
431 sorted_runs: Vec<SortedRun<FileHandle>>,
432 max_output_file_size: Option<u64>,
433 is_eligible: impl Fn(&Candidate) -> bool,
434) -> Vec<FileHandle> {
435 let files = ordered_files(&sorted_runs);
436
437 let mut best = None;
438 for left in 0..files.len() {
439 let mut candidate = Candidate::default();
440 let right_bound = left.saturating_add(*MAX_INPUT_FILES).min(files.len());
441 let mut participations: Vec<bool> = Vec::with_capacity(right_bound - left);
442 for right in left..right_bound {
443 candidate.absorb(&files[right], &files[left..right], &mut participations);
444 if candidate.num_files < 2
445 || !candidate.is_balanced()
446 || !is_eligible(&candidate)
447 || !candidate.has_balanced_level_rows()
448 || !candidate.makes_progress(max_output_file_size)
449 {
450 continue;
451 }
452
453 let score = candidate.score(max_output_file_size);
454 if best
455 .as_ref()
456 .is_none_or(|(best_score, _)| score.is_better_than(best_score))
457 {
458 best = Some((score, &files[left..=right]));
459 }
460 }
461 }
462
463 let Some((_, best)) = best else {
464 return vec![];
465 };
466 best.iter().map(|file| file.file.clone()).collect()
467}
468
469fn ordered_files(sorted_runs: &[SortedRun<FileHandle>]) -> Vec<OrderedFile<'_>> {
472 let mut files = sorted_runs
473 .iter()
474 .enumerate()
475 .flat_map(|(run_id, run)| {
476 run.items()
477 .iter()
478 .enumerate()
479 .map(move |(position_in_run, file)| OrderedFile {
480 file,
481 run_id,
482 position_in_run,
483 })
484 })
485 .collect::<Vec<_>>();
486 files.sort_unstable_by(|lhs, rhs| {
487 let (lhs_start, lhs_end) = lhs.file.range();
488 let (rhs_start, rhs_end) = rhs.file.range();
489 lhs_start
490 .cmp(&rhs_start)
491 .then_with(|| rhs_end.cmp(&lhs_end))
492 .then_with(|| lhs.run_id.cmp(&rhs.run_id))
493 .then_with(|| lhs.position_in_run.cmp(&rhs.position_in_run))
494 });
495 files
496}
497
498fn selected_overlaps_unselected(selected: &[FileHandle], window: &Window) -> bool {
499 let Some((span_start, span_end)) = selected
502 .iter()
503 .map(Ranged::range)
504 .reduce(|(start_a, end_a), (start_b, end_b)| (start_a.min(start_b), end_a.max(end_b)))
505 else {
506 return false;
507 };
508 let selected_file_ids = selected
509 .iter()
510 .map(FileHandle::file_id)
511 .collect::<HashSet<_>>();
512 window
513 .files()
514 .filter(|file| {
515 let (start, end) = file.range();
516 start <= span_end && span_start <= end
517 })
518 .filter(|file| !selected_file_ids.contains(&file.file_id()))
519 .any(|unselected| {
520 selected
521 .iter()
522 .any(|selected| selected.overlap_inclusive(unselected))
523 })
524}
525
526#[allow(clippy::too_many_arguments)]
527fn log_pick_result(
528 region_id: RegionId,
529 window: i64,
530 active_window: Option<i64>,
531 found_runs: usize,
532 file_num: usize,
533 max_output_file_size: Option<u64>,
534 filter_deleted: bool,
535 inputs: &[FileHandle],
536) {
537 let input_file_str: Vec<String> = inputs
538 .iter()
539 .map(|f| {
540 let range = f.range();
541 let start = range.0.to_iso8601_string();
542 let end = range.1.to_iso8601_string();
543 let num_rows = f.num_rows();
544 format!(
545 "File{{id: {:?}, range: ({}, {}), size: {}, num rows: {} }}",
546 f.file_id(),
547 start,
548 end,
549 ReadableSize(f.size()),
550 num_rows
551 )
552 })
553 .collect();
554 let window_str = Timestamp::new_second(window).to_iso8601_string();
555 let active_window_str = active_window.map(|s| Timestamp::new_second(s).to_iso8601_string());
556 let max_output_file_size = max_output_file_size.map(|size| ReadableSize(size).to_string());
557 info!(
558 "Region ({:?}) compaction pick result: current window: {}, active window: {:?}, \
559 found runs: {}, file num: {}, max output file size: {:?}, filter deleted: {}, \
560 input files: {:?}",
561 region_id,
562 window_str,
563 active_window_str,
564 found_runs,
565 file_num,
566 max_output_file_size,
567 filter_deleted,
568 input_file_str
569 );
570}
571
572#[async_trait::async_trait]
573impl Picker for TwcsPicker {
574 async fn pick(&self, compaction_region: &CompactionRegion) -> Result<Option<PickerOutput>> {
575 let region_id = compaction_region.region_id;
576 let picker = self.clone();
577 let compaction_region = compaction_region.clone();
578 let (expired_ssts, time_window_size, active_window, windows) =
579 common_runtime::spawn_blocking_compact(move || {
580 let levels = compaction_region.current_version.ssts.levels();
581 let expired_ssts = get_expired_ssts(
582 levels,
583 compaction_region.ttl,
584 Timestamp::current_millis(),
585 );
586 if !expired_ssts.is_empty() {
587 info!("Expired SSTs in region {}: {:?}", region_id, expired_ssts);
588 }
589 let expired_file_ids = expired_ssts
590 .iter()
591 .map(|file| file.file_id())
592 .collect::<HashSet<_>>();
593
594 let compaction_time_window = compaction_region
595 .current_version
596 .compaction_time_window
597 .map(|window| window.as_secs() as i64);
598 let time_window_size = compaction_time_window
599 .or(picker.time_window_seconds)
600 .unwrap_or_else(|| {
601 let inferred = infer_time_bucket(levels[0].files());
602 info!(
603 "Compaction window for region {} is not present, inferring from files: {:?}",
604 region_id, inferred
605 );
606 inferred
607 });
608
609 let active_window =
610 find_latest_window_in_seconds(levels[0].files(), time_window_size);
611 let windows = assign_to_windows(
612 levels
613 .iter()
614 .flat_map(LevelMeta::files)
615 .filter(|file| !expired_file_ids.contains(&file.file_id())),
616 time_window_size,
617 );
618
619 (expired_ssts, time_window_size, active_window, windows)
620 })
621 .await
622 .context(JoinSnafu)?;
623
624 let outputs = self
625 .build_output_with_time_range(region_id, windows, active_window, Some(time_window_size))
626 .await?;
627
628 if outputs.is_empty() && expired_ssts.is_empty() {
629 return Ok(None);
630 }
631
632 let max_file_size = self.max_output_file_size.map(|v| v as usize);
633 Ok(Some(PickerOutput {
634 outputs,
635 expired_ssts,
636 time_window_size,
637 max_file_size,
638 }))
639 }
640}
641
642#[derive(Clone)]
643struct Window {
644 start: Timestamp,
645 end: Timestamp,
646 files: Vec<FileHandle>,
647 time_window: i64,
648 primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>,
649}
650
651impl Window {
652 fn new_with_file(file: FileHandle) -> Self {
654 let (start, end) = file.time_range();
655 let primary_key_range = file.primary_key_range();
656 Self {
657 start,
658 end,
659 files: vec![file],
660 time_window: 0,
661 primary_key_range,
662 }
663 }
664
665 fn range(&self) -> (Timestamp, Timestamp) {
667 (self.start, self.end)
668 }
669
670 fn add_file(&mut self, file: FileHandle) {
672 let (start, end) = file.time_range();
673 self.start = self.start.min(start);
674 self.end = self.end.max(end);
675 self.primary_key_range =
676 merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range());
677 self.files.push(file);
678 }
679
680 fn files(&self) -> impl Iterator<Item = &FileHandle> {
681 self.files.iter()
682 }
683}
684
685fn assign_to_windows<'a>(
687 files: impl Iterator<Item = &'a FileHandle>,
688 time_window_size: i64,
689) -> BTreeMap<i64, Window> {
690 let mut windows: HashMap<i64, Window> = HashMap::new();
691 for f in files {
693 if f.compacting() {
694 continue;
695 }
696 let (_, end) = f.time_range();
697 let time_window = end
698 .convert_to(TimeUnit::Second)
699 .unwrap()
700 .value()
701 .align_to_ceil_by_bucket(time_window_size)
702 .unwrap_or(i64::MIN);
703
704 match windows.entry(time_window) {
705 Entry::Occupied(mut e) => {
706 e.get_mut().add_file(f.clone());
707 }
708 Entry::Vacant(e) => {
709 let mut window = Window::new_with_file(f.clone());
710 window.time_window = time_window;
711 e.insert(window);
712 }
713 }
714 }
715 windows.into_iter().collect()
716}
717
718fn time_window_intersects_range(
719 window_end: i64,
720 time_window_size: i64,
721 time_range: &TimestampRange,
722) -> bool {
723 let first_window = match time_range.start() {
724 None => i64::MIN,
725 Some(start) => {
726 let Some(first_window) = start
727 .convert_to(TimeUnit::Second)
728 .and_then(|timestamp| timestamp.value().align_to_ceil_by_bucket(time_window_size))
729 else {
730 return false;
731 };
732 first_window
733 }
734 };
735 let last_window = match time_range.end() {
736 None => i64::MAX,
737 Some(end) => {
738 let Some(last_window) = end
739 .convert_to_ceil(TimeUnit::Second)
740 .and_then(|timestamp| timestamp.value().checked_sub(1))
741 .and_then(|timestamp| timestamp.align_to_ceil_by_bucket(time_window_size))
742 else {
743 return false;
744 };
745 last_window
746 }
747 };
748 (first_window..=last_window).contains(&window_end)
749}
750
751fn window_has_overlap(this: &Window, windows: &BTreeMap<i64, Window>) -> bool {
752 windows
753 .values()
754 .filter(|that| this.time_window != that.time_window)
755 .any(|that| {
756 overlaps(&this.range(), &that.range()) && {
757 match (&this.primary_key_range, &that.primary_key_range) {
758 (Some(l), Some(r)) => primary_key_ranges_overlap(l, r),
759 _ => true,
760 }
761 }
762 })
763}
764
765fn find_latest_window_in_seconds<'a>(
768 files: impl Iterator<Item = &'a FileHandle>,
769 time_window_size: i64,
770) -> Option<i64> {
771 let mut latest_timestamp = None;
772 for f in files {
773 let (_, end) = f.time_range();
774 if let Some(latest) = latest_timestamp {
775 if end > latest {
776 latest_timestamp = Some(end);
777 }
778 } else {
779 latest_timestamp = Some(end);
780 }
781 }
782 latest_timestamp
783 .and_then(|ts| ts.convert_to_ceil(TimeUnit::Second))
784 .and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size))
785}
786
787#[cfg(test)]
788mod tests {
789 use std::collections::HashSet;
790 use std::num::NonZeroU64;
791 use std::sync::Arc;
792 use std::time::Duration;
793
794 use bytes::Bytes;
795 use common_base::Plugins;
796 use common_time::range::TimestampRange;
797 use store_api::storage::FileId;
798
799 use super::*;
800 use crate::cache::CacheManager;
801 use crate::compaction::compactor::CompactionVersion;
802 use crate::compaction::test_util::{
803 new_file_handle, new_file_handle_with_sequence, new_file_handle_with_size_and_sequence,
804 new_file_handle_with_size_sequence_and_primary_key_range,
805 };
806 use crate::config::MitoConfig;
807 use crate::region::options::RegionOptions;
808 use crate::sst::file::{FileMeta, Level};
809 use crate::sst::version::SstVersion;
810 use crate::test_util::memtable_util::metadata_for_test;
811 use crate::test_util::scheduler_util::SchedulerEnv;
812
813 #[test]
814 fn test_valid_max_input_files_env_overrides_default() {
815 assert_eq!(64, parse_max_input_files(Some("64")));
816 }
817
818 #[test]
819 fn test_invalid_max_input_files_env_falls_back_to_default() {
820 for env_value in [None, Some(""), Some("invalid"), Some("0"), Some("1")] {
821 assert_eq!(32, parse_max_input_files(env_value));
822 }
823 }
824
825 async fn compaction_region_with_expired_sst() -> CompactionRegion {
826 let env = SchedulerEnv::new().await;
827 let metadata = metadata_for_test();
828 let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
829 let mut ssts = SstVersion::new();
830 ssts.add_files(
831 Arc::new(crate::sst::file_purger::NoopFilePurger),
832 (1..=4).map(|sequence| FileMeta {
833 file_id: FileId::random(),
834 time_range: (
835 Timestamp::new_millisecond(0),
836 Timestamp::new_millisecond(10),
837 ),
838 level: 0,
839 sequence: NonZeroU64::new(sequence),
840 ..Default::default()
841 }),
842 );
843
844 CompactionRegion {
845 region_id: metadata.region_id,
846 region_options: RegionOptions::default(),
847 engine_config: Arc::new(MitoConfig::default()),
848 region_metadata: metadata.clone(),
849 cache_manager: Arc::new(CacheManager::default()),
850 access_layer: env.access_layer,
851 manifest_ctx,
852 current_version: CompactionVersion {
853 metadata,
854 options: RegionOptions::default(),
855 ssts: Arc::new(ssts),
856 compaction_time_window: None,
857 },
858 file_purger: None,
859 ttl: Some(Duration::from_millis(1).into()),
860 max_parallelism: 1,
861 plugins: Plugins::new(),
862 }
863 }
864
865 #[tokio::test]
866 async fn test_pick_expired_ssts_without_marking_compacting() {
867 let picker = TwcsPicker {
868 trigger_file_num: 4,
869 time_window_seconds: Some(3),
870 max_output_file_size: None,
871 append_mode: false,
872 max_background_tasks: None,
873 time_range: None,
874 };
875 let compaction_region = compaction_region_with_expired_sst().await;
876
877 let output = picker.pick(&compaction_region).await.unwrap().unwrap();
878
879 assert!(output.outputs.is_empty());
880 assert!(!output.expired_ssts.is_empty());
881 assert!(output.expired_ssts.iter().all(|file| !file.compacting()));
882 }
883
884 #[test]
885 fn test_get_latest_window_in_seconds() {
886 assert_eq!(
887 Some(1),
888 find_latest_window_in_seconds([new_file_handle(FileId::random(), 0, 999, 0)].iter(), 1)
889 );
890 assert_eq!(
891 Some(1),
892 find_latest_window_in_seconds(
893 [new_file_handle(FileId::random(), 0, 1000, 0)].iter(),
894 1
895 )
896 );
897
898 assert_eq!(
899 Some(-9223372036854000),
900 find_latest_window_in_seconds(
901 [new_file_handle(FileId::random(), i64::MIN, i64::MIN + 1, 0)].iter(),
902 3600,
903 )
904 );
905
906 assert_eq!(
907 (i64::MAX / 10000000 + 1) * 10000,
908 find_latest_window_in_seconds(
909 [new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0)].iter(),
910 10000,
911 )
912 .unwrap()
913 );
914
915 assert_eq!(
916 Some((i64::MAX / 3600000 + 1) * 3600),
917 find_latest_window_in_seconds(
918 [
919 new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0),
920 new_file_handle(FileId::random(), 0, 1000, 0)
921 ]
922 .iter(),
923 3600
924 )
925 );
926 }
927
928 #[test]
929 fn test_assign_to_windows() {
930 let windows = assign_to_windows(
931 [
932 new_file_handle(FileId::random(), 0, 999, 0),
933 new_file_handle(FileId::random(), 0, 999, 0),
934 new_file_handle(FileId::random(), 0, 999, 0),
935 new_file_handle(FileId::random(), 0, 999, 0),
936 new_file_handle(FileId::random(), 0, 999, 0),
937 ]
938 .iter(),
939 3,
940 );
941 let fgs = &windows.get(&0).unwrap().files;
942 assert_eq!(5, fgs.len());
943
944 let files = [FileId::random(); 3];
945 let windows = assign_to_windows(
946 [
947 new_file_handle(files[0], -2000, -3, 0),
948 new_file_handle(files[1], 0, 2999, 0),
949 new_file_handle(files[2], 50, 10001, 0),
950 ]
951 .iter(),
952 3,
953 );
954 assert_eq!(
955 files[0],
956 windows
957 .get(&0)
958 .unwrap()
959 .files()
960 .next()
961 .unwrap()
962 .file_id()
963 .file_id()
964 );
965 assert_eq!(
966 files[1],
967 windows
968 .get(&3)
969 .unwrap()
970 .files()
971 .next()
972 .unwrap()
973 .file_id()
974 .file_id()
975 );
976 assert_eq!(
977 files[2],
978 windows
979 .get(&12)
980 .unwrap()
981 .files()
982 .next()
983 .unwrap()
984 .file_id()
985 .file_id()
986 );
987 }
988
989 #[test]
990 fn test_assign_files_to_windows() {
991 let files = [
992 FileId::random(),
993 FileId::random(),
994 FileId::random(),
995 FileId::random(),
996 ];
997 let windows = assign_to_windows(
998 [
999 new_file_handle_with_sequence(files[0], 0, 999, 0, 1),
1000 new_file_handle_with_sequence(files[1], 0, 999, 0, 1),
1001 new_file_handle_with_sequence(files[2], 0, 999, 0, 2),
1002 new_file_handle_with_sequence(files[3], 0, 999, 0, 2),
1003 ]
1004 .iter(),
1005 3,
1006 );
1007 assert_eq!(windows.len(), 1);
1008 let window_files = &windows.get(&0).unwrap().files;
1009 assert_eq!(4, window_files.len());
1010 assert_eq!(
1011 window_files
1012 .iter()
1013 .map(|f| f.file_id().file_id())
1014 .collect::<HashSet<_>>(),
1015 files.into_iter().collect()
1016 );
1017 }
1018
1019 #[test]
1020 fn test_assign_compacting_to_windows() {
1021 let files = [
1022 new_file_handle(FileId::random(), 0, 999, 0),
1023 new_file_handle(FileId::random(), 0, 999, 0),
1024 new_file_handle(FileId::random(), 0, 999, 0),
1025 new_file_handle(FileId::random(), 0, 999, 0),
1026 new_file_handle(FileId::random(), 0, 999, 0),
1027 ];
1028 files[0].set_compacting(true);
1029 files[2].set_compacting(true);
1030 let mut windows = assign_to_windows(files.iter(), 3);
1031 let window0 = windows.remove(&0).unwrap();
1032 assert_eq!(3, window0.files.len());
1033 let candidates = window0
1034 .files
1035 .iter()
1036 .map(|f| f.file_id().file_id())
1037 .collect::<HashSet<_>>();
1038 assert_eq!(candidates.len(), 3);
1039 assert_eq!(
1040 candidates,
1041 [
1042 files[1].file_id().file_id(),
1043 files[3].file_id().file_id(),
1044 files[4].file_id().file_id()
1045 ]
1046 .into_iter()
1047 .collect::<HashSet<_>>()
1048 );
1049 }
1050
1051 type ExpectedWindowSpec = (i64, bool, Vec<(i64, i64)>);
1053
1054 fn pk_range(min: &'static [u8], max: &'static [u8]) -> Option<(Bytes, Bytes)> {
1055 Some((Bytes::from_static(min), Bytes::from_static(max)))
1056 }
1057
1058 fn check_assign_to_windows_with_overlapping(
1059 file_time_ranges: &[(i64, i64)],
1060 time_window: i64,
1061 expected_files: &[ExpectedWindowSpec],
1062 ) {
1063 let files: Vec<_> = (0..file_time_ranges.len())
1064 .map(|_| FileId::random())
1065 .collect();
1066
1067 let file_handles = files
1068 .iter()
1069 .zip(file_time_ranges.iter())
1070 .map(|(file_id, range)| new_file_handle(*file_id, range.0, range.1, 0))
1071 .collect::<Vec<_>>();
1072
1073 let windows = assign_to_windows(file_handles.iter(), time_window);
1074
1075 for (expected_window, overlapping, window_files) in expected_files {
1076 let actual_window = windows.get(expected_window).unwrap();
1077 let actual_overlapping = window_has_overlap(actual_window, &windows);
1078 assert_eq!(*overlapping, actual_overlapping);
1079 let mut file_ranges = actual_window
1080 .files
1081 .iter()
1082 .map(|f| {
1083 let (s, e) = f.time_range();
1084 (s.value(), e.value())
1085 })
1086 .collect::<Vec<_>>();
1087 file_ranges.sort_unstable_by(|l, r| l.0.cmp(&r.0).then(l.1.cmp(&r.1)));
1088 assert_eq!(window_files, &file_ranges);
1089 }
1090 }
1091
1092 #[test]
1093 fn test_assign_to_windows_with_overlapping() {
1094 check_assign_to_windows_with_overlapping(
1095 &[(0, 999), (1000, 1999), (2000, 2999)],
1096 2,
1097 &[
1098 (0, false, vec![(0, 999)]),
1099 (2, false, vec![(1000, 1999), (2000, 2999)]),
1100 ],
1101 );
1102
1103 check_assign_to_windows_with_overlapping(
1104 &[(0, 1), (0, 999), (100, 2999)],
1105 2,
1106 &[
1107 (0, true, vec![(0, 1), (0, 999)]),
1108 (2, true, vec![(100, 2999)]),
1109 ],
1110 );
1111
1112 check_assign_to_windows_with_overlapping(
1113 &[(0, 999), (1000, 1999), (2000, 2999), (3000, 3999)],
1114 2,
1115 &[
1116 (0, false, vec![(0, 999)]),
1117 (2, false, vec![(1000, 1999), (2000, 2999)]),
1118 (4, false, vec![(3000, 3999)]),
1119 ],
1120 );
1121
1122 check_assign_to_windows_with_overlapping(
1123 &[
1124 (0, 999),
1125 (1000, 1999),
1126 (2000, 2999),
1127 (3000, 3999),
1128 (0, 3999),
1129 ],
1130 2,
1131 &[
1132 (0, true, vec![(0, 999)]),
1133 (2, true, vec![(1000, 1999), (2000, 2999)]),
1134 (4, true, vec![(0, 3999), (3000, 3999)]),
1135 ],
1136 );
1137
1138 check_assign_to_windows_with_overlapping(
1139 &[
1140 (0, 999),
1141 (1000, 1999),
1142 (2000, 2999),
1143 (3000, 3999),
1144 (1999, 3999),
1145 ],
1146 2,
1147 &[
1148 (0, false, vec![(0, 999)]),
1149 (2, true, vec![(1000, 1999), (2000, 2999)]),
1150 (4, true, vec![(1999, 3999), (3000, 3999)]),
1151 ],
1152 );
1153
1154 check_assign_to_windows_with_overlapping(
1155 &[
1156 (0, 999), (1000, 1999), (2000, 2999), (3000, 3999), (2999, 3999), ],
1162 2,
1163 &[
1164 (0, false, vec![(0, 999)]),
1166 (2, true, vec![(1000, 1999), (2000, 2999)]),
1167 (4, true, vec![(2999, 3999), (3000, 3999)]),
1168 ],
1169 );
1170
1171 check_assign_to_windows_with_overlapping(
1172 &[
1173 (0, 999), (1000, 1999), (2000, 2999), (3000, 3999), (0, 1000), ],
1179 2,
1180 &[
1181 (0, true, vec![(0, 999)]),
1183 (2, true, vec![(0, 1000), (1000, 1999), (2000, 2999)]),
1184 (4, false, vec![(3000, 3999)]),
1185 ],
1186 );
1187 }
1188
1189 #[test]
1190 fn test_assign_to_windows_not_overlapping_when_pk_disjoint() {
1191 let files = [
1192 new_file_handle_with_size_sequence_and_primary_key_range(
1193 FileId::random(),
1194 0,
1195 1000,
1196 0,
1197 1,
1198 10,
1199 pk_range(b"a", b"f"),
1200 ),
1201 new_file_handle_with_size_sequence_and_primary_key_range(
1202 FileId::random(),
1203 500,
1204 1999,
1205 0,
1206 2,
1207 10,
1208 pk_range(b"x", b"z"),
1209 ),
1210 ];
1211
1212 let windows = assign_to_windows(files.iter(), 2);
1213
1214 let overlapping = window_has_overlap(windows.get(&2).unwrap(), &windows);
1215 assert!(!overlapping);
1216 }
1217
1218 #[test]
1219 fn test_assign_to_windows_pk_unknown_in_earlier_window_does_not_poison_later_windows() {
1220 let files = [
1221 new_file_handle(FileId::random(), 0, 1999, 0),
1222 new_file_handle_with_size_sequence_and_primary_key_range(
1223 FileId::random(),
1224 2000,
1225 3999,
1226 0,
1227 1,
1228 10,
1229 pk_range(b"a", b"f"),
1230 ),
1231 new_file_handle_with_size_sequence_and_primary_key_range(
1232 FileId::random(),
1233 3000,
1234 4999,
1235 0,
1236 2,
1237 10,
1238 pk_range(b"x", b"z"),
1239 ),
1240 ];
1241
1242 let windows = assign_to_windows(files.iter(), 2);
1243
1244 let overlapping = window_has_overlap(windows.get(&4).unwrap(), &windows);
1245 assert!(!overlapping);
1246 }
1247
1248 struct CompactionPickerTestCase {
1249 window_size: i64,
1250 input_files: Vec<FileHandle>,
1251 expected_outputs: Vec<ExpectedOutput>,
1252 }
1253
1254 impl CompactionPickerTestCase {
1255 async fn check(&self) {
1256 let file_id_to_idx = self
1257 .input_files
1258 .iter()
1259 .enumerate()
1260 .map(|(idx, file)| (file.file_id(), idx))
1261 .collect::<HashMap<_, _>>();
1262 let windows = assign_to_windows(self.input_files.iter(), self.window_size);
1263 let active_window =
1264 find_latest_window_in_seconds(self.input_files.iter(), self.window_size);
1265 let output = TwcsPicker {
1266 trigger_file_num: 2,
1267 time_window_seconds: None,
1268 max_output_file_size: None,
1269 append_mode: false,
1270 max_background_tasks: None,
1271 time_range: None,
1272 }
1273 .build_output_with_time_range(RegionId::from_u64(0), windows, active_window, None)
1274 .await
1275 .unwrap();
1276
1277 let output = output
1278 .iter()
1279 .map(|o| {
1280 let input_file_ids = o
1281 .inputs
1282 .iter()
1283 .map(|f| file_id_to_idx.get(&f.file_id()).copied().unwrap())
1284 .collect::<HashSet<_>>();
1285 (input_file_ids, o.output_level)
1286 })
1287 .collect::<Vec<_>>();
1288
1289 let expected = self
1290 .expected_outputs
1291 .iter()
1292 .map(|o| {
1293 let input_file_ids = o.input_files.iter().copied().collect::<HashSet<_>>();
1294 (input_file_ids, o.output_level)
1295 })
1296 .collect::<Vec<_>>();
1297 assert_eq!(expected, output);
1298 }
1299 }
1300
1301 struct ExpectedOutput {
1302 input_files: Vec<usize>,
1303 output_level: Level,
1304 }
1305
1306 #[tokio::test]
1307 async fn test_build_twcs_output() {
1308 let file_ids = (0..4).map(|_| FileId::random()).collect::<Vec<_>>();
1309
1310 CompactionPickerTestCase {
1312 window_size: 3,
1313 input_files: [
1314 new_file_handle_with_sequence(file_ids[0], -2000, -3, 0, 1),
1315 new_file_handle_with_sequence(file_ids[1], -3000, -100, 0, 2),
1316 new_file_handle_with_sequence(file_ids[2], 0, 2999, 0, 3), new_file_handle_with_sequence(file_ids[3], 50, 2998, 0, 4), ]
1319 .to_vec(),
1320 expected_outputs: vec![
1321 ExpectedOutput {
1322 input_files: vec![2, 3],
1323 output_level: 1,
1324 },
1325 ExpectedOutput {
1326 input_files: vec![0, 1],
1327 output_level: 1,
1328 },
1329 ],
1330 }
1331 .check()
1332 .await;
1333
1334 let file_ids = (0..6).map(|_| FileId::random()).collect::<Vec<_>>();
1341 CompactionPickerTestCase {
1342 window_size: 3,
1343 input_files: [
1344 new_file_handle_with_sequence(file_ids[0], -2000, -3, 0, 1),
1345 new_file_handle_with_sequence(file_ids[1], -3000, -100, 0, 2),
1346 new_file_handle_with_sequence(file_ids[2], 0, 2999, 0, 3),
1347 new_file_handle_with_sequence(file_ids[3], 50, 2998, 0, 4),
1348 new_file_handle_with_sequence(file_ids[4], 11, 2990, 0, 5),
1349 ]
1350 .to_vec(),
1351 expected_outputs: vec![
1352 ExpectedOutput {
1353 input_files: vec![2, 3, 4],
1354 output_level: 1,
1355 },
1356 ExpectedOutput {
1357 input_files: vec![0, 1],
1358 output_level: 1,
1359 },
1360 ],
1361 }
1362 .check()
1363 .await;
1364
1365 let file_ids = (0..6).map(|_| FileId::random()).collect::<Vec<_>>();
1370 CompactionPickerTestCase {
1371 window_size: 3,
1372 input_files: [
1373 new_file_handle_with_sequence(file_ids[0], 0, 2999, 1, 1),
1374 new_file_handle_with_sequence(file_ids[1], 0, 2998, 1, 1),
1375 new_file_handle_with_sequence(file_ids[2], 3000, 5999, 1, 2),
1376 new_file_handle_with_sequence(file_ids[3], 3000, 5000, 1, 2),
1377 new_file_handle_with_sequence(file_ids[4], 11, 2990, 0, 3),
1378 ]
1379 .to_vec(),
1380 expected_outputs: vec![
1381 ExpectedOutput {
1382 input_files: vec![2, 3],
1383 output_level: 1,
1384 },
1385 ExpectedOutput {
1386 input_files: vec![0, 1],
1389 output_level: 1,
1390 },
1391 ],
1392 }
1393 .check()
1394 .await;
1395 }
1396
1397 #[tokio::test]
1398 async fn test_build_output_skips_pk_disjoint_files() {
1399 let files = [
1400 new_file_handle_with_size_sequence_and_primary_key_range(
1401 FileId::random(),
1402 0,
1403 2999,
1404 0,
1405 1,
1406 10,
1407 pk_range(b"a", b"f"),
1408 ),
1409 new_file_handle_with_size_sequence_and_primary_key_range(
1410 FileId::random(),
1411 50,
1412 2998,
1413 0,
1414 2,
1415 10,
1416 pk_range(b"x", b"z"),
1417 ),
1418 ];
1419 let windows = assign_to_windows(files.iter(), 3);
1420 let active_window = find_latest_window_in_seconds(files.iter(), 3);
1421 let output = TwcsPicker {
1422 trigger_file_num: 4,
1423 time_window_seconds: None,
1424 max_output_file_size: None,
1425 append_mode: false,
1426 max_background_tasks: None,
1427 time_range: None,
1428 }
1429 .build_output_with_time_range(RegionId::from_u64(0), windows, active_window, None)
1430 .await
1431 .unwrap();
1432
1433 assert!(output.is_empty());
1434 }
1435
1436 #[test]
1437 fn test_append_mode_filter_large_files() {
1438 let file_ids = (0..4).map(|_| FileId::random()).collect::<Vec<_>>();
1439 let max_output_file_size = 1000u64;
1440
1441 let small_file_1 = new_file_handle_with_size_and_sequence(file_ids[0], 0, 999, 0, 1, 500);
1443 let large_file_1 = new_file_handle_with_size_and_sequence(file_ids[1], 0, 999, 0, 2, 1500);
1444 let small_file_2 = new_file_handle_with_size_and_sequence(file_ids[2], 0, 999, 0, 3, 800);
1445 let large_file_2 = new_file_handle_with_size_and_sequence(file_ids[3], 0, 999, 0, 4, 2000);
1446
1447 let mut files_to_merge = vec![small_file_1, large_file_1, small_file_2, large_file_2];
1448
1449 let original_count = files_to_merge.len();
1451
1452 files_to_merge.retain(|file| file.size() <= max_output_file_size);
1454
1455 assert_eq!(files_to_merge.len(), 2);
1457 assert_eq!(original_count, 4);
1458
1459 for file in &files_to_merge {
1461 assert!(
1462 file.size() <= max_output_file_size,
1463 "File size {} should be <= {}",
1464 file.size(),
1465 max_output_file_size
1466 );
1467 }
1468 }
1469
1470 #[tokio::test]
1471 async fn test_build_output_multiple_windows_with_zero_runs() {
1472 let file_ids = (0..7).map(|_| FileId::random()).collect::<Vec<_>>();
1473
1474 let files = [
1475 new_file_handle_with_sequence(file_ids[0], 0, 999, 0, 1),
1477 new_file_handle_with_sequence(file_ids[1], 0, 999, 0, 2),
1478 new_file_handle_with_sequence(file_ids[2], 0, 999, 0, 3),
1479 new_file_handle_with_sequence(file_ids[3], 3000, 3999, 0, 4),
1481 new_file_handle_with_sequence(file_ids[4], 3000, 3999, 0, 5),
1482 new_file_handle_with_sequence(file_ids[5], 3000, 3999, 0, 6),
1483 new_file_handle_with_sequence(file_ids[6], 3000, 3999, 0, 7),
1484 ];
1485
1486 let windows = assign_to_windows(files.iter(), 3);
1487
1488 let picker = TwcsPicker {
1490 trigger_file_num: 4, time_window_seconds: Some(3),
1492 max_output_file_size: None,
1493 append_mode: false,
1494 max_background_tasks: None,
1495 time_range: None,
1496 };
1497
1498 let active_window = find_latest_window_in_seconds(files.iter(), 3);
1499 let output = picker
1500 .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None)
1501 .await
1502 .unwrap();
1503
1504 assert!(
1505 !output.is_empty(),
1506 "Should have output from windows with runs, even when one window has 0 runs"
1507 );
1508
1509 let all_output_files: Vec<_> = output
1510 .iter()
1511 .flat_map(|o| o.inputs.iter())
1512 .map(|f| f.file_id().file_id())
1513 .collect();
1514
1515 assert!(
1516 all_output_files.contains(&file_ids[3])
1517 || all_output_files.contains(&file_ids[4])
1518 || all_output_files.contains(&file_ids[5]),
1519 "Output should contain files from the window with runs"
1520 );
1521 }
1522
1523 #[tokio::test]
1524 async fn test_build_output_single_window_zero_runs() {
1525 let file_ids = (0..2).map(|_| FileId::random()).collect::<Vec<_>>();
1526
1527 let large_file_1 = new_file_handle_with_size_and_sequence(file_ids[0], 0, 999, 0, 1, 2000); let large_file_2 = new_file_handle_with_size_and_sequence(file_ids[1], 0, 999, 0, 2, 2500); let files = [large_file_1, large_file_2];
1531
1532 let windows = assign_to_windows(files.iter(), 3);
1533
1534 let picker = TwcsPicker {
1535 trigger_file_num: 2,
1536 time_window_seconds: Some(3),
1537 max_output_file_size: Some(1000),
1538 append_mode: true,
1539 max_background_tasks: None,
1540 time_range: None,
1541 };
1542
1543 let active_window = find_latest_window_in_seconds(files.iter(), 3);
1544 let output = picker
1545 .build_output_with_time_range(RegionId::from_u64(456), windows, active_window, None)
1546 .await
1547 .unwrap();
1548
1549 assert!(
1551 output.is_empty(),
1552 "Should return empty output when no runs are found after filtering"
1553 );
1554 }
1555
1556 #[tokio::test]
1557 async fn test_append_mode_can_pick_remaining_single_level_files() {
1558 let files = [
1559 new_file_handle_with_size_and_sequence(FileId::random(), 0, 9, 0, 1, 100),
1560 new_file_handle_with_size_and_sequence(FileId::random(), 20, 29, 0, 2, 100),
1561 new_file_handle_with_size_and_sequence(FileId::random(), 40, 49, 0, 3, 100),
1562 new_file_handle_with_size_and_sequence(FileId::random(), 60, 69, 0, 4, 2_000),
1563 ];
1564 let windows = assign_to_windows(files.iter(), 1);
1565 let picker = TwcsPicker {
1566 trigger_file_num: 4,
1567 time_window_seconds: Some(1),
1568 max_output_file_size: Some(1_000),
1569 append_mode: true,
1570 max_background_tasks: None,
1571 time_range: None,
1572 };
1573
1574 let output = picker
1575 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(1), None)
1576 .await
1577 .unwrap();
1578
1579 assert_eq!(1, output.len());
1580 assert_eq!(3, output[0].inputs.len());
1581 }
1582
1583 #[tokio::test]
1584 async fn test_max_background_tasks_truncation() {
1585 let file_ids = (0..10).map(|_| FileId::random()).collect::<Vec<_>>();
1586 let max_background_tasks = 3;
1587
1588 let files = [
1590 new_file_handle_with_sequence(file_ids[0], 0, 999, 0, 1),
1592 new_file_handle_with_sequence(file_ids[1], 0, 999, 0, 2),
1593 new_file_handle_with_sequence(file_ids[2], 0, 999, 0, 3),
1594 new_file_handle_with_sequence(file_ids[3], 0, 999, 0, 4),
1595 new_file_handle_with_sequence(file_ids[4], 3000, 3999, 0, 5),
1597 new_file_handle_with_sequence(file_ids[5], 3000, 3999, 0, 6),
1598 new_file_handle_with_sequence(file_ids[6], 3000, 3999, 0, 7),
1599 new_file_handle_with_sequence(file_ids[7], 3000, 3999, 0, 8),
1600 new_file_handle_with_sequence(file_ids[8], 6000, 6999, 0, 9),
1602 new_file_handle_with_sequence(file_ids[9], 6000, 6999, 0, 10),
1603 ];
1604
1605 let windows = assign_to_windows(files.iter(), 3);
1606
1607 let picker = TwcsPicker {
1608 trigger_file_num: 4,
1609 time_window_seconds: Some(3),
1610 max_output_file_size: None,
1611 append_mode: false,
1612 max_background_tasks: Some(max_background_tasks),
1613 time_range: None,
1614 };
1615
1616 let active_window = find_latest_window_in_seconds(files.iter(), 3);
1617 let output = picker
1618 .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None)
1619 .await
1620 .unwrap();
1621
1622 assert!(
1624 output.len() <= max_background_tasks,
1625 "Output should be truncated to max_background_tasks: expected <= {}, got {}",
1626 max_background_tasks,
1627 output.len()
1628 );
1629
1630 let picker_no_limit = TwcsPicker {
1632 trigger_file_num: 4,
1633 time_window_seconds: Some(3),
1634 max_output_file_size: None,
1635 append_mode: false,
1636 max_background_tasks: None,
1637 time_range: None,
1638 };
1639
1640 let windows_no_limit = assign_to_windows(files.iter(), 3);
1641 let output_no_limit = picker_no_limit
1642 .build_output_with_time_range(
1643 RegionId::from_u64(123),
1644 windows_no_limit,
1645 active_window,
1646 None,
1647 )
1648 .await
1649 .unwrap();
1650
1651 if output_no_limit.len() > max_background_tasks {
1653 assert!(
1654 output_no_limit.len() > output.len(),
1655 "Without limit should have more outputs than with limit"
1656 );
1657 }
1658 }
1659
1660 #[tokio::test]
1661 async fn test_max_background_tasks_no_truncation_when_under_limit() {
1662 let file_ids = (0..4).map(|_| FileId::random()).collect::<Vec<_>>();
1663 let max_background_tasks = 10; let files = [
1667 new_file_handle_with_sequence(file_ids[0], 0, 999, 0, 1),
1668 new_file_handle_with_sequence(file_ids[1], 0, 999, 0, 2),
1669 new_file_handle_with_sequence(file_ids[2], 0, 999, 0, 3),
1670 new_file_handle_with_sequence(file_ids[3], 0, 999, 0, 4),
1671 ];
1672
1673 let windows = assign_to_windows(files.iter(), 3);
1674
1675 let picker = TwcsPicker {
1676 trigger_file_num: 4,
1677 time_window_seconds: Some(3),
1678 max_output_file_size: None,
1679 append_mode: false,
1680 max_background_tasks: Some(max_background_tasks),
1681 time_range: None,
1682 };
1683
1684 let active_window = find_latest_window_in_seconds(files.iter(), 3);
1685 let output = picker
1686 .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None)
1687 .await
1688 .unwrap();
1689
1690 assert!(
1692 output.len() <= max_background_tasks,
1693 "Output should be within limit"
1694 );
1695 assert!(!output.is_empty(), "Should have at least one output");
1697 }
1698
1699 #[tokio::test]
1700 async fn test_pick_multiple_runs() {
1701 common_telemetry::init_default_ut_logging();
1702
1703 let num_files = 8;
1704 let file_ids = (0..num_files).map(|_| FileId::random()).collect::<Vec<_>>();
1705
1706 let files: Vec<_> = file_ids
1708 .iter()
1709 .enumerate()
1710 .map(|(idx, file_id)| {
1711 new_file_handle_with_size_and_sequence(
1712 *file_id,
1713 0,
1714 999,
1715 0,
1716 (idx + 1) as u64,
1717 1024 * 1024,
1718 )
1719 })
1720 .collect();
1721
1722 let windows = assign_to_windows(files.iter(), 3);
1723
1724 let picker = TwcsPicker {
1725 trigger_file_num: 4,
1726 time_window_seconds: Some(3),
1727 max_output_file_size: None,
1728 append_mode: false,
1729 max_background_tasks: None,
1730 time_range: None,
1731 };
1732
1733 let active_window = find_latest_window_in_seconds(files.iter(), 3);
1734 let output = picker
1735 .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None)
1736 .await
1737 .unwrap();
1738
1739 assert_eq!(1, output.len());
1740 assert_eq!(output[0].inputs.len(), num_files);
1741 }
1742
1743 #[tokio::test]
1744 async fn test_window_trigger_can_exceed_input_limit() {
1745 common_telemetry::init_default_ut_logging();
1746
1747 let num_files = 50;
1748 let file_ids = (0..num_files).map(|_| FileId::random()).collect::<Vec<_>>();
1749
1750 let files: Vec<_> = file_ids
1752 .iter()
1753 .enumerate()
1754 .map(|(idx, file_id)| {
1755 new_file_handle_with_size_and_sequence(
1756 *file_id,
1757 (idx / 2 * 10) as i64,
1758 (idx / 2 * 10 + 5) as i64,
1759 0,
1760 (idx + 1) as u64,
1761 1024 * 1024,
1762 )
1763 })
1764 .collect();
1765
1766 let windows = assign_to_windows(files.iter(), 3);
1767
1768 let picker = TwcsPicker {
1769 trigger_file_num: num_files,
1770 time_window_seconds: Some(3),
1771 max_output_file_size: None,
1772 append_mode: false,
1773 max_background_tasks: None,
1774 time_range: None,
1775 };
1776
1777 let active_window = find_latest_window_in_seconds(files.iter(), 3);
1778 let output = picker
1779 .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None)
1780 .await
1781 .unwrap();
1782
1783 assert_eq!(1, output.len());
1784 assert_eq!(output[0].inputs.len(), num_files.min(*MAX_INPUT_FILES));
1785 }
1786
1787 #[tokio::test]
1788 async fn test_limit_max_input_files_keeps_deletion_markers() {
1789 common_telemetry::init_default_ut_logging();
1790
1791 let mut files = vec![new_file_handle_with_size_and_sequence(
1794 FileId::random(),
1795 0,
1796 3_000_000,
1797 0,
1798 1,
1799 1024 * 1024 * 1024,
1800 )];
1801 files.extend((0..32).map(|idx: i64| {
1802 new_file_handle_with_size_and_sequence(
1803 FileId::random(),
1804 (idx + 1) * 10_000,
1805 (idx + 1) * 10_000 + 1_000,
1806 0,
1807 (idx + 2) as u64,
1808 1024,
1809 )
1810 }));
1811
1812 let windows = assign_to_windows(files.iter(), 3600);
1813
1814 let picker = TwcsPicker {
1815 trigger_file_num: 4,
1816 time_window_seconds: Some(3600),
1817 max_output_file_size: None,
1818 append_mode: false,
1819 max_background_tasks: None,
1820 time_range: None,
1821 };
1822
1823 let active_window = find_latest_window_in_seconds(files.iter(), 3600);
1824 let output = picker
1825 .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None)
1826 .await
1827 .unwrap();
1828
1829 assert_eq!(1, output.len());
1830 assert_eq!(32, output[0].inputs.len());
1833 assert!(
1834 !output[0].filter_deleted,
1835 "deletion markers must be kept once the file num limit drops files they may mask"
1836 );
1837 }
1838
1839 #[tokio::test]
1840 async fn test_limit_max_input_files_still_filters_without_overlap() {
1841 common_telemetry::init_default_ut_logging();
1842
1843 let files: Vec<_> = (0..40i64)
1846 .map(|idx| {
1847 new_file_handle_with_size_and_sequence(
1848 FileId::random(),
1849 (idx + 1) * 10_000,
1850 (idx + 1) * 10_000 + 1_000,
1851 0,
1852 (idx + 1) as u64,
1853 1024,
1854 )
1855 })
1856 .collect();
1857
1858 let windows = assign_to_windows(files.iter(), 3600);
1859
1860 let picker = TwcsPicker {
1861 trigger_file_num: 4,
1862 time_window_seconds: Some(3600),
1863 max_output_file_size: Some(1024 * 1024 * 1024),
1864 append_mode: false,
1865 max_background_tasks: None,
1866 time_range: None,
1867 };
1868
1869 let active_window = find_latest_window_in_seconds(files.iter(), 3600);
1870 let output = picker
1871 .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None)
1872 .await
1873 .unwrap();
1874
1875 assert_eq!(1, output.len());
1876 assert_eq!(32, output[0].inputs.len());
1877 assert!(output[0].filter_deleted);
1878 }
1879
1880 #[tokio::test]
1881 async fn test_newer_windows_have_priority() {
1882 let older_file_ids = [FileId::random(), FileId::random()];
1883 let newer_file_ids = [FileId::random(), FileId::random()];
1884 let files = [
1885 new_file_handle_with_sequence(older_file_ids[0], 1_000, 1_999, 0, 1),
1886 new_file_handle_with_sequence(older_file_ids[1], 1_000, 1_999, 0, 2),
1887 new_file_handle_with_sequence(newer_file_ids[0], 7_000, 7_999, 0, 3),
1888 new_file_handle_with_sequence(newer_file_ids[1], 7_000, 7_999, 0, 4),
1889 ];
1890 let windows = assign_to_windows(files.iter(), 3);
1891 let picker = TwcsPicker {
1892 trigger_file_num: 2,
1893 time_window_seconds: Some(3),
1894 max_output_file_size: None,
1895 append_mode: false,
1896 max_background_tasks: Some(1),
1897 time_range: None,
1898 };
1899
1900 let output = picker
1901 .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), None)
1902 .await
1903 .unwrap();
1904
1905 assert_eq!(1, output.len());
1906 assert_eq!(
1907 newer_file_ids.into_iter().collect::<HashSet<_>>(),
1908 output[0]
1909 .inputs
1910 .iter()
1911 .map(|file| file.file_id().file_id())
1912 .collect::<HashSet<_>>()
1913 );
1914 }
1915
1916 #[test]
1917 fn test_filter_time_windows_by_time_range() {
1918 let time_range = TimestampRange::new(
1919 Timestamp::new_millisecond(1_200),
1920 Timestamp::new_millisecond(1_800),
1921 )
1922 .unwrap();
1923
1924 assert!(time_window_intersects_range(3, 3, &time_range));
1925 assert!(!time_window_intersects_range(9, 3, &time_range));
1926
1927 let boundary_range =
1928 TimestampRange::new(Timestamp::new_second(0), Timestamp::new_second(3)).unwrap();
1929 assert!(time_window_intersects_range(0, 3, &boundary_range));
1930 assert!(time_window_intersects_range(3, 3, &boundary_range));
1931 assert!(!time_window_intersects_range(6, 3, &boundary_range));
1932
1933 let overflowing_range = TimestampRange::new(
1934 Timestamp::new_second(i64::MAX - 1),
1935 Timestamp::new_second(i64::MAX),
1936 )
1937 .unwrap();
1938 assert!(!time_window_intersects_range(0, 4, &overflowing_range));
1939 }
1940
1941 #[tokio::test]
1942 async fn test_time_range_filter_precedes_background_task_limit() {
1943 let early_file_ids = [FileId::random(), FileId::random()];
1944 let selected_file_ids = [FileId::random(), FileId::random()];
1945 let files = [
1946 new_file_handle_with_sequence(early_file_ids[0], 1_000, 1_999, 0, 1),
1947 new_file_handle_with_sequence(early_file_ids[1], 1_000, 1_999, 0, 2),
1948 new_file_handle_with_sequence(selected_file_ids[0], 7_000, 7_999, 0, 3),
1949 new_file_handle_with_sequence(selected_file_ids[1], 7_000, 7_999, 0, 4),
1950 ];
1951 let windows = assign_to_windows(files.iter(), 3);
1952 let picker = TwcsPicker {
1953 trigger_file_num: 2,
1954 time_window_seconds: Some(3),
1955 max_output_file_size: None,
1956 append_mode: false,
1957 max_background_tasks: Some(1),
1958 time_range: TimestampRange::new(
1959 Timestamp::new_millisecond(7_200),
1960 Timestamp::new_millisecond(7_800),
1961 ),
1962 };
1963
1964 let output = picker
1965 .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), Some(3))
1966 .await
1967 .unwrap();
1968
1969 assert_eq!(1, output.len());
1970 assert_eq!(
1971 selected_file_ids.into_iter().collect::<HashSet<_>>(),
1972 output[0]
1973 .inputs
1974 .iter()
1975 .map(|file| file.file_id().file_id())
1976 .collect::<HashSet<_>>()
1977 );
1978 }
1979
1980 #[tokio::test]
1981 async fn test_count_first_prefers_more_files_over_smaller_overlap() {
1982 let files = [
1983 new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10),
1984 new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10),
1985 new_file_handle_with_size_and_sequence(FileId::random(), 40, 50, 0, 3, 10),
1986 new_file_handle_with_size_and_sequence(FileId::random(), 5, 15, 0, 4, 10),
1987 ];
1988 let windows = assign_to_windows(files.iter(), 100);
1989 let picker = TwcsPicker {
1990 trigger_file_num: 2,
1991 time_window_seconds: Some(100),
1992 max_output_file_size: None,
1993 append_mode: false,
1994 max_background_tasks: None,
1995 time_range: None,
1996 };
1997
1998 let output = picker
1999 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None)
2000 .await
2001 .unwrap();
2002
2003 assert_eq!(1, output.len());
2004 assert_eq!(4, output[0].inputs.len());
2005 }
2006
2007 #[tokio::test]
2008 async fn test_count_first_trigger_counts_physical_ssts() {
2009 let files = [
2010 new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10),
2011 new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10),
2012 new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10),
2013 ];
2014 let windows = assign_to_windows(files.iter(), 100);
2015 let picker = TwcsPicker {
2016 trigger_file_num: 3,
2017 time_window_seconds: Some(100),
2018 max_output_file_size: None,
2019 append_mode: false,
2020 max_background_tasks: None,
2021 time_range: None,
2022 };
2023
2024 let output = picker
2025 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None)
2026 .await
2027 .unwrap();
2028
2029 assert_eq!(1, output.len());
2030 assert_eq!(3, output[0].inputs.len());
2031 }
2032
2033 #[tokio::test]
2034 async fn test_count_first_does_not_compact_overlap_below_trigger() {
2035 let files = [
2036 new_file_handle_with_size_and_sequence(FileId::random(), 0, 20, 0, 1, 10),
2037 new_file_handle_with_size_and_sequence(FileId::random(), 10, 30, 0, 2, 10),
2038 ];
2039 let windows = assign_to_windows(files.iter(), 100);
2040 let picker = TwcsPicker {
2041 trigger_file_num: 3,
2042 time_window_seconds: Some(100),
2043 max_output_file_size: None,
2044 append_mode: false,
2045 max_background_tasks: None,
2046 time_range: None,
2047 };
2048
2049 let output = picker
2050 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None)
2051 .await
2052 .unwrap();
2053
2054 assert!(output.is_empty());
2055 }
2056
2057 #[tokio::test]
2058 async fn test_filter_deleted_is_false_when_selected_files_overlap_unselected_file() {
2059 let mut files = (0..32)
2060 .map(|idx| {
2061 new_file_handle_with_size_and_sequence(
2062 FileId::random(),
2063 idx * 10,
2064 idx * 10 + 9,
2065 0,
2066 idx as u64 + 1,
2067 10,
2068 )
2069 })
2070 .collect::<Vec<_>>();
2071 files.push(new_file_handle_with_size_and_sequence(
2072 FileId::random(),
2073 0,
2074 320,
2075 0,
2076 33,
2077 10,
2078 ));
2079 let windows = assign_to_windows(files.iter(), 1000);
2080 let picker = TwcsPicker {
2081 trigger_file_num: 2,
2082 time_window_seconds: Some(1000),
2083 max_output_file_size: None,
2084 append_mode: false,
2085 max_background_tasks: None,
2086 time_range: None,
2087 };
2088
2089 let output = picker
2090 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None)
2091 .await
2092 .unwrap();
2093
2094 assert_eq!(1, output.len());
2095 assert_eq!(DEFAULT_MAX_INPUT_FILES, output[0].inputs.len());
2096 assert!(!output[0].filter_deleted);
2097 }
2098
2099 fn new_file(start: i64, end: i64, sequence: u64, file_size: u64) -> FileHandle {
2100 new_file_handle_with_size_and_sequence(FileId::random(), start, end, 0, sequence, file_size)
2101 }
2102
2103 fn new_file_with_level_and_rows(
2104 start: i64,
2105 end: i64,
2106 level: Level,
2107 sequence: u64,
2108 file_size: u64,
2109 num_rows: u64,
2110 ) -> FileHandle {
2111 let file = new_file_handle_with_size_and_sequence(
2112 FileId::random(),
2113 start,
2114 end,
2115 level,
2116 sequence,
2117 file_size,
2118 );
2119 let mut meta = file.meta_ref().clone();
2120 meta.num_rows = num_rows;
2121 FileHandle::new(meta, crate::test_util::new_noop_file_purger())
2122 }
2123
2124 fn picked_ranges(files: &[FileHandle]) -> Vec<(i64, i64)> {
2125 files
2126 .iter()
2127 .map(|file| {
2128 let (start, end) = file.range();
2129 (start.value(), end.value())
2130 })
2131 .collect()
2132 }
2133
2134 #[test]
2135 fn test_count_first_rejects_dominant_historical_file() {
2136 let picked = pick_count_first(
2137 vec![SortedRun::from(vec![
2138 new_file(0, 9, 1, 400),
2139 new_file(20, 29, 2, 100),
2140 ])],
2141 None,
2142 );
2143
2144 assert!(picked.is_empty());
2145 }
2146
2147 #[test]
2148 fn test_count_first_accepts_balanced_historical_file() {
2149 let picked = pick_count_first(
2150 vec![SortedRun::from(vec![
2151 new_file(0, 9, 1, 400),
2152 new_file(20, 29, 2, 100),
2153 new_file(40, 49, 3, 100),
2154 new_file(60, 69, 4, 100),
2155 new_file(80, 89, 5, 100),
2156 ])],
2157 None,
2158 );
2159
2160 assert_eq!(
2161 vec![(0, 9), (20, 29), (40, 49), (60, 69), (80, 89)],
2162 picked_ranges(&picked)
2163 );
2164 }
2165
2166 #[test]
2167 fn test_count_first_finds_smaller_balanced_interval_when_larger_one_is_unbalanced() {
2168 let picked = pick_count_first(
2169 vec![SortedRun::from(vec![
2170 new_file(0, 9, 1, 1000),
2171 new_file(20, 29, 2, 100),
2172 new_file(40, 49, 3, 100),
2173 new_file(60, 69, 4, 100),
2174 new_file(80, 89, 5, 100),
2175 ])],
2176 None,
2177 );
2178
2179 assert_eq!(
2180 vec![(20, 29), (40, 49), (60, 69), (80, 89)],
2181 picked_ranges(&picked)
2182 );
2183 }
2184
2185 #[test]
2186 fn test_count_first_prefers_overlap_participants_when_file_counts_match() {
2187 let first_run = (0..DEFAULT_MAX_INPUT_FILES)
2188 .map(|idx| {
2189 let start = idx as i64 * 20;
2190 let end = if idx + 1 == DEFAULT_MAX_INPUT_FILES {
2191 700
2192 } else {
2193 start + 9
2194 };
2195 new_file(start, end, idx as u64 + 1, 10)
2196 })
2197 .collect::<Vec<_>>();
2198 let overlapping = new_file(690, 710, 100, 10);
2199
2200 let picked = pick_count_first(
2201 vec![
2202 SortedRun::from(first_run),
2203 SortedRun::from(vec![overlapping]),
2204 ],
2205 None,
2206 );
2207 let ranges = picked_ranges(&picked);
2208
2209 assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len());
2210 assert!(!ranges.contains(&(0, 9)));
2211 assert!(ranges.contains(&(690, 710)));
2212 }
2213
2214 #[tokio::test]
2215 async fn test_picker_avoids_chained_l1_rewrites_by_compacting_levels_separately() {
2216 let mut enough_l0 = (0..DEFAULT_MAX_INPUT_FILES)
2217 .map(|idx| {
2218 let start = idx as i64 * 20;
2219 new_file_handle_with_size_and_sequence(
2220 FileId::random(),
2221 start,
2222 start + 9,
2223 0,
2224 idx as u64 + 1,
2225 10,
2226 )
2227 })
2228 .collect::<Vec<_>>();
2229 enough_l0.push(new_file_handle_with_size_and_sequence(
2230 FileId::random(),
2231 0,
2232 700,
2233 1,
2234 100,
2235 100,
2236 ));
2237 let mut enough_l1 = (0..4)
2238 .map(|idx| {
2239 new_file_handle_with_size_and_sequence(
2240 FileId::random(),
2241 idx * 100,
2242 idx * 100 + 99,
2243 1,
2244 idx as u64 + 1,
2245 100,
2246 )
2247 })
2248 .collect::<Vec<_>>();
2249 enough_l1.extend((0..3).map(|idx| {
2250 new_file_handle_with_size_and_sequence(
2251 FileId::random(),
2252 idx * 20,
2253 idx * 20 + 9,
2254 0,
2255 idx as u64 + 10,
2256 10,
2257 )
2258 }));
2259 let picker = TwcsPicker {
2260 trigger_file_num: 4,
2261 time_window_seconds: Some(1),
2262 max_output_file_size: None,
2263 append_mode: false,
2264 max_background_tasks: None,
2265 time_range: None,
2266 };
2267
2268 for (case, files, expected_level, expected_len) in [
2269 ("L0 reaches trigger", enough_l0, 0, DEFAULT_MAX_INPUT_FILES),
2270 ("L1 reaches trigger", enough_l1, 1, 4),
2271 ] {
2272 let windows = assign_to_windows(files.iter(), 1);
2273 let output = picker
2274 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(1), None)
2275 .await
2276 .unwrap();
2277
2278 assert_eq!(1, output.len(), "{case}");
2279 assert_eq!(expected_len, output[0].inputs.len(), "{case}");
2280 assert!(
2281 output[0]
2282 .inputs
2283 .iter()
2284 .all(|file| file.level() == expected_level),
2285 "{case}"
2286 );
2287 }
2288 }
2289
2290 #[tokio::test]
2291 async fn test_picker_falls_back_to_l1_when_triggered_l0_cannot_make_progress() {
2292 let mut files = (0..4)
2293 .map(|idx| {
2294 let start = idx * 20;
2295 new_file_handle_with_size_and_sequence(
2296 FileId::random(),
2297 start,
2298 start + 9,
2299 0,
2300 idx as u64 + 1,
2301 600,
2302 )
2303 })
2304 .collect::<Vec<_>>();
2305 files.extend((0..4).map(|idx| {
2306 let start = idx * 20 + 100;
2307 new_file_handle_with_size_and_sequence(
2308 FileId::random(),
2309 start,
2310 start + 9,
2311 1,
2312 idx as u64 + 10,
2313 100,
2314 )
2315 }));
2316 let windows = assign_to_windows(files.iter(), 1);
2317 let picker = TwcsPicker {
2318 trigger_file_num: 4,
2319 time_window_seconds: Some(1),
2320 max_output_file_size: Some(512),
2321 append_mode: false,
2322 max_background_tasks: None,
2323 time_range: None,
2324 };
2325
2326 let output = picker
2327 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(1), None)
2328 .await
2329 .unwrap();
2330
2331 assert_eq!(1, output.len());
2332 assert_eq!(4, output[0].inputs.len());
2333 assert!(output[0].inputs.iter().all(|file| file.level() == 1));
2334 }
2335
2336 #[test]
2337 fn test_mixed_candidate_row_balance() {
2338 for (case, l1_rows, l0_rows, expected_len) in [
2339 ("L1 rows dominate", 1_000_000, 10_000, 0),
2340 ("L1 rows meet ratio", 60_000, 10_000, 4),
2341 ("L0 rows are unknown", 1_000_000, 0, 4),
2342 ] {
2343 let files = vec![
2344 new_file_with_level_and_rows(0, 99, 1, 1, 100, l1_rows),
2345 new_file_with_level_and_rows(0, 9, 0, 2, 100, l0_rows),
2346 new_file_with_level_and_rows(20, 29, 0, 3, 100, l0_rows),
2347 new_file_with_level_and_rows(40, 49, 0, 4, 100, l0_rows),
2348 ];
2349
2350 let picked = pick_mixed_count_first(vec![SortedRun::from(files)], None);
2351
2352 assert_eq!(expected_len, picked.len(), "{case}");
2353 }
2354 }
2355
2356 #[test]
2357 fn test_count_first_prefers_smaller_bytes_when_file_counts_match() {
2358 let files = (0..=DEFAULT_MAX_INPUT_FILES)
2359 .map(|idx| {
2360 let start = idx as i64 * 20;
2361 let size = if idx == 0 {
2362 100
2363 } else if idx == DEFAULT_MAX_INPUT_FILES {
2364 1
2365 } else {
2366 10
2367 };
2368 new_file(start, start + 9, idx as u64 + 1, size)
2369 })
2370 .collect::<Vec<_>>();
2371
2372 let picked = pick_count_first(vec![SortedRun::from(files)], None);
2373 let ranges = picked_ranges(&picked);
2374
2375 assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len());
2376 assert_eq!(Some(&(20, 29)), ranges.first());
2377 assert_eq!(Some(&(640, 649)), ranges.last());
2378 }
2379
2380 #[test]
2381 fn test_count_first_skips_pure_rewrite_without_progress() {
2382 let picked = pick_count_first(
2386 vec![SortedRun::from(vec![
2387 new_file(0, 9, 1, 600),
2388 new_file(20, 29, 2, 600),
2389 new_file(40, 49, 3, 600),
2390 new_file(60, 69, 4, 600),
2391 ])],
2392 Some(512),
2393 );
2394
2395 assert!(picked.is_empty());
2396 }
2397
2398 #[test]
2399 fn test_count_first_allows_overlap_resolution_without_file_reduction() {
2400 let picked = pick_count_first(
2404 vec![
2405 SortedRun::from(vec![new_file(0, 19, 1, 600)]),
2406 SortedRun::from(vec![new_file(10, 29, 2, 600)]),
2407 ],
2408 Some(512),
2409 );
2410
2411 assert_eq!(vec![(0, 19), (10, 29)], picked_ranges(&picked));
2412 }
2413
2414 #[test]
2415 fn test_count_first_prefers_guaranteed_reduction_over_pure_rewrite() {
2416 let picked = pick_count_first(
2420 vec![SortedRun::from(vec![
2421 new_file(0, 9, 1, 600),
2422 new_file(10, 19, 2, 600),
2423 new_file(20, 29, 3, 600),
2424 new_file(30, 39, 4, 10),
2425 new_file(40, 49, 5, 10),
2426 new_file(50, 59, 6, 10),
2427 ])],
2428 Some(512),
2429 );
2430
2431 assert_eq!(vec![(30, 39), (40, 49), (50, 59)], picked_ranges(&picked));
2432 }
2433
2434 #[test]
2435 fn test_count_first_prefers_earlier_time_on_exact_tie() {
2436 let files = (0..=DEFAULT_MAX_INPUT_FILES)
2437 .map(|idx| {
2438 let start = idx as i64 * 20;
2439 new_file(start, start + 9, idx as u64 + 1, 10)
2440 })
2441 .collect::<Vec<_>>();
2442
2443 let picked = pick_count_first(vec![SortedRun::from(files)], None);
2444 let ranges = picked_ranges(&picked);
2445
2446 assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len());
2447 assert_eq!(Some(&(0, 9)), ranges.first());
2448 assert_eq!(Some(&(620, 629)), ranges.last());
2449 }
2450
2451 #[test]
2452 fn test_count_first_keeps_each_interleaved_run_contiguous() {
2453 let picked = pick_count_first(
2454 vec![
2455 SortedRun::from(vec![
2456 new_file(0, 9, 1, 1),
2457 new_file(20, 29, 2, 1),
2458 new_file(40, 49, 3, 1),
2459 ]),
2460 SortedRun::from(vec![new_file(10, 19, 4, 1), new_file(30, 39, 5, 1)]),
2461 ],
2462 None,
2463 );
2464
2465 assert_eq!(
2466 vec![(0, 9), (10, 19), (20, 29), (30, 39), (40, 49)],
2467 picked_ranges(&picked)
2468 );
2469 }
2470
2471 #[test]
2478 fn test_count_first_candidate_spans_same_sequence_files() {
2479 let picked = pick_count_first(
2480 vec![SortedRun::from(vec![
2481 new_file(0, 9, 1, 10),
2482 new_file(10, 19, 2, 10),
2484 new_file(10, 19, 2, 10),
2485 new_file(10, 19, 2, 10),
2486 new_file(20, 29, 3, 10),
2487 ])],
2488 None,
2489 );
2490
2491 assert_eq!(
2492 vec![(0, 9), (10, 19), (10, 19), (10, 19), (20, 29)],
2493 picked_ranges(&picked)
2494 );
2495 }
2496
2497 #[tokio::test]
2503 async fn test_count_first_merges_window_with_interleaved_flush_groups() {
2504 let files = [
2505 new_file_handle_with_size_and_sequence(FileId::random(), 0, 9, 0, 1, 10),
2507 new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10),
2509 new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10),
2510 new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10),
2511 new_file_handle_with_size_and_sequence(FileId::random(), 20, 29, 0, 3, 10),
2513 ];
2514 let windows = assign_to_windows(files.iter(), 100);
2515 let picker = TwcsPicker {
2516 trigger_file_num: 3,
2517 time_window_seconds: Some(100),
2518 max_output_file_size: None,
2519 append_mode: false,
2520 max_background_tasks: None,
2521 time_range: None,
2522 };
2523
2524 let output = picker
2525 .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None)
2526 .await
2527 .unwrap();
2528
2529 assert_eq!(1, output.len());
2530 assert_eq!(5, output[0].inputs.len());
2532 }
2533}