chore: use CacheReader on remotes

tonbo-io · Oct 23, 2024 · cd33a8f · cd33a8f
1 parent e16065c
commit cd33a8f
Show file tree

Hide file tree

Showing 8 changed files with 160 additions and 185 deletions.
diff --git a/src/compaction/mod.rs b/src/compaction/mod.rs
@@ -136,7 +136,7 @@ where
     ) -> Result<Option<Scope<R::Key>>, CompactionError<R>> {
         if !batches.is_empty() {
             let level_0_path = option.level_fs_path(0).unwrap_or(&option.base_path);
-            let level_0_fs = manager.get_fs(level_0_path);
+            let (level_0_fs, _) = manager.get_fs(level_0_path);
 
             let mut min = None;
             let mut max = None;
@@ -207,7 +207,7 @@ where
                 Self::next_level_scopes(version, &mut min, &mut max, level, &meet_scopes_l)?;
 
             let level_path = option.level_fs_path(level).unwrap_or(&option.base_path);
-            let level_fs = manager.get_fs(level_path);
+            let (level_fs, is_local) = manager.get_fs(level_path);
             let mut streams = Vec::with_capacity(meet_scopes_l.len() + meet_scopes_ll.len());
             // This Level
             if level == 0 {
@@ -218,7 +218,7 @@ where
                         .await?;
 
                     streams.push(ScanStream::SsTable {
-                        inner: SsTable::open(option, file, path)
+                        inner: SsTable::open(option, file, path, is_local)
                             .await?
                             .scan(
                                 (Bound::Unbounded, Bound::Unbounded),
@@ -240,7 +240,7 @@ where
                     u32::MAX.into(),
                     None,
                     ProjectionMask::all(),
-                    level_fs.clone(),
+                    (level_fs.clone(), is_local),
                 )
                 .ok_or(CompactionError::EmptyLevel)?;
 
@@ -260,7 +260,7 @@ where
                     u32::MAX.into(),
                     None,
                     ProjectionMask::all(),
-                    level_fs.clone(),
+                    (level_fs.clone(), is_local),
                 )
                 .ok_or(CompactionError::EmptyLevel)?;
 
@@ -846,14 +846,14 @@ pub(crate) mod tests {
         option: &Arc<DbOption<Test>>,
         manager: &StoreManager,
     ) -> ((FileId, FileId, FileId, FileId, FileId), Version<Test>) {
-        let level_0_fs = option
+        let (level_0_fs, _) = option
             .level_fs_path(0)
             .map(|path| manager.get_fs(path))
-            .unwrap_or(manager.base_fs());
-        let level_1_fs = option
+            .unwrap_or((manager.base_fs(), true));
+        let (level_1_fs, _) = option
             .level_fs_path(1)
             .map(|path| manager.get_fs(path))
-            .unwrap_or(manager.base_fs());
+            .unwrap_or((manager.base_fs(), true));
 
         // level 0
         let table_gen_1 = FileId::new();
@@ -1121,14 +1121,14 @@ pub(crate) mod tests {
             .await
             .unwrap();
 
-        let level_0_fs = option
+        let (level_0_fs, _) = option
             .level_fs_path(0)
             .map(|path| manager.get_fs(path))
-            .unwrap_or(manager.base_fs());
-        let level_1_fs = option
+            .unwrap_or((manager.base_fs(), true));
+        let (level_1_fs, _) = option
             .level_fs_path(1)
             .map(|path| manager.get_fs(path))
-            .unwrap_or(manager.base_fs());
+            .unwrap_or((manager.base_fs(), true));
 
         let table_gen0 = FileId::new();
         let table_gen1 = FileId::new();

diff --git a/src/fs/cache_reader.rs b/src/fs/cache_reader.rs
@@ -25,7 +25,8 @@ impl<R> CacheReader<R> {
         inner: R,
         meta_path: Path,
     ) -> Result<CacheReader<R>, CacheError> {
-        let path = path_to_local(&option.path)?;
+        // SAFETY: `meta_path` must be the path of a parquet file
+        let path = path_to_local(&option.path.child(meta_path.filename().unwrap()))?;
         let cache: HybridCache<Range<usize>, Bytes> = HybridCacheBuilder::new()
             .memory(option.memory)
             .storage(Engine::Large) // use large object disk cache engine only

diff --git a/src/fs/manager.rs b/src/fs/manager.rs
@@ -5,7 +5,7 @@ use fusio_dispatch::FsOptions;
 
 pub struct StoreManager {
     base_fs: Arc<dyn DynFs>,
-    fs_map: HashMap<Path, Arc<dyn DynFs>>,
+    fs_map: HashMap<Path, (Arc<dyn DynFs>, bool)>,
 }
 
 impl StoreManager {
@@ -16,7 +16,10 @@ impl StoreManager {
         let mut fs_map = HashMap::with_capacity(levels_fs.len());
 
         for (path, fs_options) in levels_fs.into_iter().flatten() {
-            fs_map.entry(path).or_insert(fs_options.parse()?);
+            let is_local = matches!(fs_options, FsOptions::Local);
+            fs_map
+                .entry(path)
+                .or_insert((fs_options.parse()?, is_local));
         }
         let base_fs = base_options.parse()?;
 
@@ -27,8 +30,11 @@ impl StoreManager {
         &self.base_fs
     }
 
-    pub fn get_fs(&self, path: &Path) -> &Arc<dyn DynFs> {
-        self.fs_map.get(path).unwrap_or(&self.base_fs)
+    pub fn get_fs(&self, path: &Path) -> (&Arc<dyn DynFs>, bool) {
+        self.fs_map
+            .get(path)
+            .map(|(fs, is_local)| (fs, *is_local))
+            .unwrap_or((&self.base_fs, false))
     }
 }
 

diff --git a/src/ondisk/scan.rs b/src/ondisk/scan.rs
@@ -6,13 +6,14 @@ use std::{
 };
 
 use arrow::datatypes::Schema;
-use fusio_parquet::reader::AsyncReader;
 use futures_core::{ready, Stream};
-use parquet::arrow::{async_reader::ParquetRecordBatchStream, ProjectionMask};
+use parquet::arrow::{
+    async_reader::{AsyncFileReader, ParquetRecordBatchStream},
+    ProjectionMask,
+};
 use pin_project_lite::pin_project;
 
 use crate::{
-    fs::cache_reader::CacheReader,
     record::Record,
     stream::record_batch::{RecordBatchEntry, RecordBatchIterator},
 };
@@ -21,7 +22,7 @@ pin_project! {
     #[derive(Debug)]
     pub struct SsTableScan<'scan, R>{
         #[pin]
-        stream: ParquetRecordBatchStream<CacheReader<AsyncReader>>,
+        stream: ParquetRecordBatchStream<Box<dyn AsyncFileReader>>,
         iter: Option<RecordBatchIterator<R>>,
         projection_mask: ProjectionMask,
         full_schema: Arc<Schema>,
@@ -31,7 +32,7 @@ pin_project! {
 
 impl<R> SsTableScan<'_, R> {
     pub(crate) fn new(
-        stream: ParquetRecordBatchStream<CacheReader<AsyncReader>>,
+        stream: ParquetRecordBatchStream<Box<dyn AsyncFileReader>>,
         projection_mask: ProjectionMask,
         full_schema: Arc<Schema>,
     ) -> Self {