diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index e811d7f2f51..e95050342cb 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -323,6 +323,30 @@ impl WriteCache { .await } + /// Downloads the target file into write cache only when it is not cached. + /// + /// Returns `Ok(true)` if this call performs a download, or `Ok(false)` if the + /// file is already present in write cache and download is skipped. + pub(crate) async fn download_if_absent( + &self, + index_key: IndexKey, + remote_path: &str, + remote_store: &ObjectStore, + file_size: u64, + ) -> Result { + if self.file_cache.contains_key(&index_key) { + debug!( + "Skip downloading file already in write cache, region: {}, file: {}", + index_key.region_id, index_key.file_id + ); + return Ok(false); + } + + self.download(index_key, remote_path, remote_store, file_size) + .await?; + Ok(true) + } + /// Uploads a Parquet file or a Puffin file to the remote object store. pub(crate) async fn upload( &self, @@ -472,13 +496,16 @@ impl UploadTracker { #[cfg(test)] mod tests { + use bytes::Bytes; use common_test_util::temp_dir::create_temp_dir; use object_store::ATOMIC_WRITE_DIR; use parquet::file::metadata::PageIndexPolicy; use store_api::region_request::PathType; + use store_api::storage::FileId; use super::*; use crate::access_layer::OperationType; + use crate::cache::file_cache::IndexValue; use crate::cache::test_util::{assert_parquet_metadata_equal, new_fs_store}; use crate::cache::{CacheManager, CacheStrategy}; use crate::error::InvalidBatchSnafu; @@ -745,4 +772,67 @@ mod tests { assert!(!has_files); } + + #[tokio::test] + async fn test_download_if_absent_skips_when_cached() { + let mut env = TestEnv::new().await; + let remote_store = env.init_object_store_manager(); + + let local_dir = create_temp_dir(""); + let local_store = new_fs_store(local_dir.path().to_str().unwrap()); + let write_cache = env + .create_write_cache(local_store.clone(), ReadableSize::mb(10)) + .await; + + let region_id = RegionId::new(1024, 1); + let file_id = FileId::random(); + let key = IndexKey::new(region_id, file_id, FileType::Parquet); + write_cache + .file_cache() + .put(key, IndexValue { file_size: 1 }) + .await; + + let downloaded = write_cache + .download_if_absent(key, "missing/path.parquet", &remote_store, 1) + .await + .unwrap(); + + assert!(!downloaded); + } + + #[tokio::test] + async fn test_download_if_absent_downloads_when_missing() { + let mut env = TestEnv::new().await; + let remote_store = env.init_object_store_manager(); + + let local_dir = create_temp_dir(""); + let local_store = new_fs_store(local_dir.path().to_str().unwrap()); + let write_cache = env + .create_write_cache(local_store.clone(), ReadableSize::mb(10)) + .await; + + let region_id = RegionId::new(1024, 2); + let file_id = FileId::random(); + let key = IndexKey::new(region_id, file_id, FileType::Parquet); + let remote_path = format!("download-if-absent/{file_id}.parquet"); + let remote_data = Bytes::from_static(b"download-if-absent-test"); + remote_store + .write(&remote_path, remote_data.clone()) + .await + .unwrap(); + + let downloaded = write_cache + .download_if_absent(key, &remote_path, &remote_store, remote_data.len() as u64) + .await + .unwrap(); + + assert!(downloaded); + assert!(write_cache.file_cache().contains_key(&key)); + + let cached_data = local_store + .read(&write_cache.file_cache().cache_file_path(key)) + .await + .unwrap(); + assert_eq!(cached_data.to_vec(), remote_data.to_vec()); + } } diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index e7708f7b651..1be65864307 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -514,11 +514,11 @@ async fn edit_region( common_runtime::spawn_global(async move { WRITE_CACHE_INFLIGHT_DOWNLOAD.add(1); - if write_cache - .download(index_key, &remote_path, layer.object_store(), file_size) - .await - .is_ok() - { + let parquet_cached = write_cache + .download_if_absent(index_key, &remote_path, layer.object_store(), file_size) + .await; + + if parquet_cached.is_ok() { // Triggers the filling of the parquet metadata cache. // The parquet file is already downloaded. let mut cache_metrics = Default::default(); @@ -531,7 +531,9 @@ async fn edit_region( ) .await; - listener.on_file_cache_filled(index_key.file_id); + if matches!(parquet_cached, Ok(true)) { + listener.on_file_cache_filled(index_key.file_id); + } } if is_index_exist { // also download puffin file