Implement object store that caches to disk

Signed-off-by: Peter Ke <[email protected]>

Author: PeterKeDer

crates/core/src/logstore/mod.rs

@@ -51,6 +51,7 @@ use std::cmp::{max, min};
 use std::collections::HashMap;
 use std::io::{BufRead, BufReader, Cursor};
 use std::sync::{Arc, LazyLock};
+use std::time::Duration;
 
 use bytes::Bytes;
 #[cfg(feature = "datafusion")]
@@ -61,11 +62,13 @@ use delta_kernel::engine::default::executor::tokio::{
 use delta_kernel::engine::default::DefaultEngine;
 use delta_kernel::{AsAny, Engine};
 use futures::{StreamExt, TryStreamExt};
+use humantime::parse_duration;
 use object_store::{path::Path, Error as ObjectStoreError, ObjectStore};
 use regex::Regex;
 use serde::de::{Error, SeqAccess, Visitor};
 use serde::ser::SerializeSeq;
 use serde::{Deserialize, Serialize};
+use storage::file_cache::{FileCachePolicy, FileCacheStorageBackend};
 use tokio::runtime::RuntimeFlavor;
 use tracing::{debug, error, warn};
 use url::Url;
@@ -75,6 +78,7 @@ use crate::kernel::log_segment::PathExt;
 use crate::kernel::transaction::TransactionError;
 use crate::kernel::Action;
 use crate::protocol::{get_last_checkpoint, ProtocolError};
+use crate::table::builder::ensure_table_uri;
 use crate::{DeltaResult, DeltaTableError};
 
 pub use self::config::StorageConfig;
@@ -191,20 +195,80 @@ pub fn logstore_for(location: Url, storage_config: StorageConfig) -> DeltaResult
     Err(DeltaTableError::InvalidTableLocation(location.into()))
 }
 
+#[derive(Debug)]
+struct DeltaFileCachePolicy {
+    last_checkpoint_valid_duration: Duration,
+}
+
+impl DeltaFileCachePolicy {
+    pub fn new(last_checkpoint_valid_duration: Duration) -> Self {
+        Self {
+            last_checkpoint_valid_duration,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl FileCachePolicy for DeltaFileCachePolicy {
+    async fn should_update_cache(
+        &self,
+        cache_meta: &object_store::ObjectMeta,
+        object_store: ObjectStoreRef,
+    ) -> object_store::Result<bool> {
+        let location = &cache_meta.location;
+        if location.filename() == Some("_last_checkpoint") {
+            let meta = object_store.head(location).await?;
+            Ok(cache_meta.last_modified + self.last_checkpoint_valid_duration < meta.last_modified)
+        } else {
+            // Other delta files are always immutable, so no need to update cache
+            Ok(false)
+        }
+    }
+}
+
 /// Return the [LogStoreRef] using the given [ObjectStoreRef]
 pub fn logstore_with(
     root_store: ObjectStoreRef,
     location: Url,
     storage_config: StorageConfig,
 ) -> DeltaResult<LogStoreRef> {
+    let store = if let Some(location) = storage_config.raw.get("file_cache_path") {
+        let path = ensure_table_uri(location)?.to_file_path().map_err(|_| {
+            DeltaTableError::generic(format!(
+                "Expected file_cache_path to be a valid file path: {location}",
+            ))
+        })?;
+
+        let last_checkpoint_valid_duration = if let Some(duration_str) = storage_config
+            .raw
+            .get("file_cache_last_checkpoint_valid_duration")
+        {
+            parse_duration(duration_str).map_err(|e| {
+                DeltaTableError::generic(
+                    format!("Failed to parse {duration_str} as Duration: {e}",),
+                )
+            })?
+        } else {
+            Duration::ZERO
+        };
+
+        Arc::new(FileCacheStorageBackend::try_new(
+            root_store,
+            path,
+            Arc::new(DeltaFileCachePolicy::new(last_checkpoint_valid_duration)),
+        )?)
+    } else {
+        root_store
+    };
+
     let scheme = Url::parse(&format!("{}://", location.scheme()))
         .map_err(|_| DeltaTableError::InvalidTableLocation(location.clone().into()))?;
 
     if let Some(factory) = logstore_factories().get(&scheme) {
         debug!("Found a logstore provider for {scheme}");
         return factory
             .value()
-            .with_options_internal(root_store, &location, &storage_config);
+            .with_options_internal(store, &location, &storage_config);
     }
 
     error!("Could not find a logstore for the scheme {scheme}");
@@ -913,6 +977,42 @@ pub(crate) mod tests {
             .try_collect::<Vec<Path>>()
             .await
     }
+
+    #[tokio::test]
+    async fn test_file_cache_enable_by_storage_option() {
+        use object_store::memory::InMemory;
+        use tempfile::tempdir;
+        let memory_store = Arc::new(InMemory::new());
+        let log_path = Path::from("delta-table/_delta_log/00000000000000000001.json");
+
+        let payload = "test";
+        memory_store
+            .put(&log_path, payload.into())
+            .await
+            .expect("Failed to write log file");
+
+        // Enable file cache by setting the storage option
+        let cache_dir = tempdir().unwrap();
+        let storge_options = HashMap::from([(
+            "file_cache_path".to_string(),
+            cache_dir.path().to_str().unwrap().to_string(),
+        )]);
+
+        let table_uri = "memory:///delta-table";
+        let table = crate::DeltaTableBuilder::from_valid_uri(table_uri)
+            .unwrap()
+            .with_storage_backend(memory_store, Url::parse(table_uri).unwrap())
+            .with_storage_options(storge_options)
+            .build()
+            .unwrap();
+
+        // Intentially discarding result. We just want to verify cache is enabled
+        let _ = table.log_store().peek_next_commit(0).await;
+
+        // File should have been cached
+        let cache_path = cache_dir.path().join(log_path.as_ref());
+        assert!(cache_path.exists());
+    }
 }
 
 #[cfg(all(test, feature = "datafusion"))]