From e375060b73b00c55db66820ea800c6c2c49515d3 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 31 Jan 2024 15:21:30 +0800 Subject: [PATCH] refactor: add same SST files (#3270) * Make adding same SST file multiple times possible, instead of panic there. * Update src/mito2/src/sst/version.rs Co-authored-by: Yingwen --------- Co-authored-by: Yingwen --- src/mito2/src/sst/version.rs | 37 ++++++++++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/src/mito2/src/sst/version.rs b/src/mito2/src/sst/version.rs index c690e8e0860a..dc388c742d73 100644 --- a/src/mito2/src/sst/version.rs +++ b/src/mito2/src/sst/version.rs @@ -55,10 +55,10 @@ impl SstVersion { ) { for file in files_to_add { let level = file.level; - let handle = FileHandle::new(file, file_purger.clone()); - let file_id = handle.file_id(); - let old = self.levels[level as usize].files.insert(file_id, handle); - assert!(old.is_none(), "Adds an existing file: {file_id}"); + self.levels[level as usize] + .files + .entry(file.file_id) + .or_insert_with(|| FileHandle::new(file, file_purger.clone())); } } @@ -157,3 +157,32 @@ fn new_level_meta_vec() -> LevelMetaArray { .try_into() .unwrap() // safety: LevelMetaArray is a fixed length array with length MAX_LEVEL } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_util::new_noop_file_purger; + + #[test] + fn test_add_files() { + let purger = new_noop_file_purger(); + + let files = (1..=3) + .map(|_| FileMeta { + file_id: FileId::random(), + ..Default::default() + }) + .collect::>(); + + let mut version = SstVersion::new(); + // files[1] is added multiple times, and that's ok. + version.add_files(purger.clone(), files[..=1].iter().cloned()); + version.add_files(purger, files[1..].iter().cloned()); + + let added_files = &version.levels()[0].files; + assert_eq!(added_files.len(), 3); + files.iter().for_each(|f| { + assert!(added_files.contains_key(&f.file_id)); + }); + } +}