refactor: use parquet lru reader as trait object

Author: ethe

bindings/python/src/db.rs

@@ -1,6 +1,5 @@
 use std::sync::Arc;
 
-use parquet_lru::NoopCache;
 use pyo3::{
     prelude::*,
     pyclass, pymethods,
@@ -10,7 +9,6 @@ use pyo3::{
 use pyo3_asyncio::tokio::{future_into_py, get_runtime};
 use tonbo::{
     executor::tokio::TokioExecutor,
-    fs::FileId,
     record::{ColumnDesc, DynRecord},
     DB,
 };
@@ -30,7 +28,7 @@ type PyExecutor = TokioExecutor;
 pub struct TonboDB {
     desc: Arc<Vec<Column>>,
     primary_key_index: usize,
-    db: Arc<DB<DynRecord, PyExecutor, NoopCache<FileId>>>,
+    db: Arc<DB<DynRecord, PyExecutor>>,
 }
 
 #[pymethods]

bindings/python/src/transaction.rs

@@ -1,13 +1,12 @@
 use std::{mem::transmute, sync::Arc};
 
-use parquet_lru::NoopCache;
 use pyo3::{
     pyclass, pymethods,
     types::{PyAnyMethods, PyMapping, PyMappingMethods, PySequenceMethods, PyTuple},
     Bound, IntoPy, Py, PyAny, PyResult, Python,
 };
 use pyo3_asyncio::tokio::future_into_py;
-use tonbo::{fs::FileId, record::DynRecord, transaction, Projection};
+use tonbo::{record::DynRecord, transaction, Projection};
 
 use crate::{
     column::Column,
@@ -19,14 +18,14 @@ use crate::{
 
 #[pyclass]
 pub struct Transaction {
-    txn: Option<transaction::Transaction<'static, DynRecord, NoopCache<FileId>>>,
+    txn: Option<transaction::Transaction<'static, DynRecord>>,
     desc: Arc<Vec<Column>>,
     primary_key_index: usize,
 }
 
 impl Transaction {
     pub(crate) fn new<'txn>(
-        txn: transaction::Transaction<'txn, DynRecord, NoopCache<FileId>>,
+        txn: transaction::Transaction<'txn, DynRecord>,
         desc: Arc<Vec<Column>>,
     ) -> Self {
         let primary_key_index = desc
@@ -38,8 +37,8 @@ impl Transaction {
         Transaction {
             txn: Some(unsafe {
                 transmute::<
-                    transaction::Transaction<'txn, DynRecord, NoopCache<FileId>>,
-                    transaction::Transaction<'static, DynRecord, NoopCache<FileId>>,
+                    transaction::Transaction<'txn, DynRecord>,
+                    transaction::Transaction<'static, DynRecord>,
                 >(txn)
             }),
             desc,
@@ -85,8 +84,8 @@ impl Transaction {
         let txn = self.txn.as_ref().unwrap();
         let txn = unsafe {
             transmute::<
-                &transaction::Transaction<'_, DynRecord, NoopCache<FileId>>,
-                &'static transaction::Transaction<'_, DynRecord, NoopCache<FileId>>,
+                &transaction::Transaction<'_, DynRecord>,
+                &'static transaction::Transaction<'_, DynRecord>,
             >(txn)
         };
 
@@ -170,8 +169,8 @@ impl Transaction {
         let txn = self.txn.as_ref().unwrap();
         let txn = unsafe {
             transmute::<
-                &transaction::Transaction<'_, DynRecord, NoopCache<FileId>>,
-                &'static transaction::Transaction<'_, DynRecord, NoopCache<FileId>>,
+                &transaction::Transaction<'_, DynRecord>,
+                &'static transaction::Transaction<'_, DynRecord>,
             >(txn)
         };
         let col_desc = self.desc.get(self.primary_key_index).unwrap();

examples/datafusion.rs

@@ -26,7 +26,7 @@ use datafusion::{
 use fusio::path::Path;
 use futures_core::Stream;
 use futures_util::StreamExt;
-use parquet_lru::NoopCache;
+use parquet_lru::NoCache;
 use tokio::fs;
 use tonbo::{
     executor::tokio::TokioExecutor, fs::FileId, inmem::immutable::ArrowArrays, record::Record,
@@ -43,12 +43,12 @@ pub struct Music {
 }
 
 struct MusicProvider {
-    db: Arc<DB<Music, TokioExecutor, NoopCache<FileId>>>,
+    db: Arc<DB<Music, TokioExecutor, NoCache<FileId>>>,
 }
 
 struct MusicExec {
     cache: PlanProperties,
-    db: Arc<DB<Music, TokioExecutor, NoopCache<FileId>>>,
+    db: Arc<DB<Music, TokioExecutor, NoCache<FileId>>>,
     projection: Option<Vec<usize>>,
     limit: Option<usize>,
     range: (Bound<<Music as Record>::Key>, Bound<<Music as Record>::Key>),
@@ -98,7 +98,7 @@ impl TableProvider for MusicProvider {
 
 impl MusicExec {
     fn new(
-        db: Arc<DB<Music, TokioExecutor, NoopCache<FileId>>>,
+        db: Arc<DB<Music, TokioExecutor, NoCache<FileId>>>,
         projection: Option<&Vec<usize>>,
     ) -> Self {
         let schema = Music::arrow_schema();

parquet-lru/Cargo.toml

@@ -7,8 +7,8 @@ name = "parquet-lru"
 version = "0.1.0"
 
 [features]
+default = ["foyer"]
 foyer = ["dep:foyer", "dep:serde"]
-full = ["foyer"]
 
 [dependencies]
 bytes = { version = "1.8.0", features = ["serde"] }
@@ -17,4 +17,3 @@ futures-core = "0.3.31"
 futures-util = "0.3.31"
 parquet = { version = "53.2.0", features = ["async"] }
 serde = { version = "1.0.214", optional = true }
-thiserror = "2.0.3"

parquet-lru/src/foyer.rs

@@ -10,7 +10,7 @@ use parquet::{
 };
 use serde::{Deserialize, Serialize};
 
-use crate::{Error, LruCache, Options};
+use crate::LruCache;
 
 #[derive(Clone)]
 pub struct FoyerCache<K>
@@ -32,23 +32,14 @@ impl<K> LruCache<K> for FoyerCache<K>
 where
     for<'a> K: Send + Sync + Hash + Eq + Serialize + Deserialize<'a> + Clone + 'static,
 {
-    type LruReader<R: AsyncFileReader + 'static> = FoyerReader<K, R>;
+    type LruReader<R> = FoyerReader<K, R>
+    where
+        R: AsyncFileReader + 'static;
 
-    async fn new(options: Options) -> Result<Self, Error> {
-        Ok(Self {
-            inner: Arc::new(FoyerCacheInner {
-                meta: foyer::CacheBuilder::new(options.meta_capacity).build(),
-                data: foyer::HybridCacheBuilder::new()
-                    .memory(options.data_capacity)
-                    .storage(foyer::Engine::Large)
-                    .build()
-                    .await
-                    .map_err(|e| Error::External(e.into()))?,
-            }),
-        })
-    }
-
-    async fn get_reader<R: AsyncFileReader>(&self, key: K, reader: R) -> FoyerReader<K, R> {
+    async fn get_reader<R>(&self, key: K, reader: R) -> FoyerReader<K, R>
+    where
+        R: AsyncFileReader,
+    {
         FoyerReader::new(self.clone(), key, reader)
     }
 }

parquet-lru/src/lib.rs

@@ -1,63 +1,56 @@
+mod r#dyn;
 #[cfg(feature = "foyer")]
 pub mod foyer;
 
 use std::{future::Future, marker::PhantomData};
 
-use parquet::{arrow::async_reader::AsyncFileReader, errors::Result};
-use thiserror::Error;
+use parquet::arrow::async_reader::AsyncFileReader;
 
-#[derive(Default)]
-pub struct Options {
-    meta_capacity: usize,
-    data_capacity: usize,
-}
-
-impl Options {
-    pub fn meta_capacity(mut self, meta_capacity: usize) -> Self {
-        self.meta_capacity = meta_capacity;
-        self
-    }
-
-    pub fn data_capacity(mut self, data_capacity: usize) -> Self {
-        self.data_capacity = data_capacity;
-        self
-    }
-}
+pub use crate::r#dyn::*;
 
-pub trait LruCache<K>: Clone + Send + Sync + 'static {
-    type LruReader<R: AsyncFileReader + 'static>: AsyncFileReader + 'static;
-
-    fn new(options: Options) -> impl Future<Output = Result<Self, Error>> + Send;
+pub trait LruCache<K>
+where
+    K: 'static,
+{
+    type LruReader<R>: AsyncFileReader + 'static
+    where
+        R: AsyncFileReader + 'static;
 
     fn get_reader<R>(&self, key: K, reader: R) -> impl Future<Output = Self::LruReader<R>> + Send
     where
         R: AsyncFileReader + 'static;
 }
 
-#[derive(Clone, Default)]
-pub struct NoopCache<K> {
+#[derive(Default)]
+pub struct NoCache<K> {
     _phantom: PhantomData<K>,
 }
 
-impl<K> LruCache<K> for NoopCache<K>
-where
-    K: Send + Sync + Clone + 'static,
-{
-    type LruReader<R: AsyncFileReader + 'static> = R;
-
-    async fn new(_options: Options) -> Result<Self, Error> {
-        Ok(Self {
+impl<K> Clone for NoCache<K> {
+    fn clone(&self) -> Self {
+        Self {
             _phantom: PhantomData,
-        })
-    }
-
-    async fn get_reader<R: AsyncFileReader>(&self, _key: K, reader: R) -> R {
-        reader
+        }
     }
 }
 
-#[derive(Debug, Error)]
-pub enum Error {
-    #[error("External lru implementation error: {0}")]
-    External(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
+unsafe impl<K> Send for NoCache<K> {}
+
+unsafe impl<K> Sync for NoCache<K> {}
+
+impl<K> LruCache<K> for NoCache<K>
+where
+    K: 'static,
+{
+    type LruReader<R> = R
+    where
+        R: AsyncFileReader + 'static;
+
+    #[allow(clippy::manual_async_fn)]
+    fn get_reader<R>(&self, _key: K, reader: R) -> impl Future<Output = R> + Send
+    where
+        R: AsyncFileReader,
+    {
+        async move { reader }
+    }
 }

src/compaction/mod.rs

@@ -5,10 +5,10 @@ use fusio::DynFs;
 use fusio_parquet::writer::AsyncWriter;
 use futures_util::StreamExt;
 use parquet::arrow::{AsyncArrowWriter, ProjectionMask};
-use parquet_lru::LruCache;
 use thiserror::Error;
 use tokio::sync::oneshot;
 
+use crate::ParquetLru;
 use crate::{
     fs::{manager::StoreManager, FileId, FileType},
     inmem::{
@@ -60,13 +60,10 @@ where
         }
     }
 
-    pub(crate) async fn check_then_compaction<C>(
+    pub(crate) async fn check_then_compaction(
         &mut self,
-        parquet_lru_cache: C,
-    ) -> Result<(), CompactionError<R>>
-    where
-        C: LruCache<FileId> + Unpin,
-    {
+        parquet_lru_cache: ParquetLru,
+    ) -> Result<(), CompactionError<R>> {
         let mut guard = self.schema.write().await;
 
         guard.trigger.reset();
@@ -114,7 +111,7 @@ where
                         &mut delete_gens,
                         &guard.record_instance,
                         &self.manager,
-                        &parquet_lru_cache,
+                        parquet_lru_cache,
                     )
                     .await?;
                 }
@@ -193,7 +190,7 @@ where
     }
 
     #[allow(clippy::too_many_arguments)]
-    pub(crate) async fn major_compaction<C>(
+    pub(crate) async fn major_compaction(
         version: &Version<R>,
         option: &DbOption<R>,
         mut min: &R::Key,
@@ -202,11 +199,8 @@ where
         delete_gens: &mut Vec<(FileId, usize)>,
         instance: &RecordInstance,
         manager: &StoreManager,
-        parquet_cache: &C,
-    ) -> Result<(), CompactionError<R>>
-    where
-        C: LruCache<FileId> + Unpin,
-    {
+        parquet_cache: ParquetLru,
+    ) -> Result<(), CompactionError<R>> {
         let mut level = 0;
 
         while level < MAX_LEVEL - 2 {
@@ -391,18 +385,15 @@ where
         (meet_scopes_l, start_l, end_l - 1)
     }
 
-    async fn build_tables<'scan, C>(
+    async fn build_tables<'scan>(
         option: &DbOption<R>,
         version_edits: &mut Vec<VersionEdit<<R as Record>::Key>>,
         level: usize,
-        streams: Vec<ScanStream<'scan, R, C>>,
+        streams: Vec<ScanStream<'scan, R>>,
         instance: &RecordInstance,
         fs: &Arc<dyn DynFs>,
-    ) -> Result<(), CompactionError<R>>
-    where
-        C: LruCache<FileId> + Unpin,
-    {
-        let mut stream = MergeStream::<R, C>::from_vec(streams, u32::MAX.into()).await?;
+    ) -> Result<(), CompactionError<R>> {
+        let mut stream = MergeStream::<R>::from_vec(streams, u32::MAX.into()).await?;
 
         // Kould: is the capacity parameter necessary?
         let mut builder = R::Columns::builder(&instance.arrow_schema::<R>(), 8192);
@@ -529,7 +520,7 @@ pub(crate) mod tests {
     use fusio_dispatch::FsOptions;
     use fusio_parquet::writer::AsyncWriter;
     use parquet::arrow::AsyncArrowWriter;
-    use parquet_lru::NoopCache;
+    use parquet_lru::NoCache;
     use tempfile::TempDir;
 
     use crate::{
@@ -827,7 +818,7 @@ pub(crate) mod tests {
             &mut vec![],
             &RecordInstance::Normal,
             &manager,
-            &NoopCache::default(),
+            Arc::new(NoCache::default()),
         )
         .await
         .unwrap();
@@ -1219,7 +1210,7 @@ pub(crate) mod tests {
             &mut vec![],
             &RecordInstance::Normal,
             &manager,
-            &NoopCache::default(),
+            Arc::new(NoCache::default()),
         )
         .await
         .unwrap();
@@ -1240,7 +1231,7 @@ pub(crate) mod tests {
         option.major_default_oldest_table_num = 1;
         option.trigger_type = TriggerType::Length(5);
 
-        let db: DB<Test, TokioExecutor, _> = DB::new(option, TokioExecutor::new()).await.unwrap();
+        let db: DB<Test, TokioExecutor> = DB::new(option, TokioExecutor::new()).await.unwrap();
 
         for i in 5..9 {
             let item = Test {