Add DataFusion parquet range cache

alexanderbianchi · alexanderbianchi · commit 6287142b6d3a · 2026-05-09T16:29:23.000-04:00
diff --git a/quickwit/quickwit-datafusion/src/object_store_registry.rs b/quickwit/quickwit-datafusion/src/object_store_registry.rs
@@ -51,7 +51,7 @@ use datafusion::common::{DataFusionError, Result as DFResult};
 use datafusion::execution::object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry};
 use object_store::ObjectStore;
 use quickwit_common::uri::Uri;
-use quickwit_storage::StorageResolver;
+use quickwit_storage::{StorageCache, StorageResolver};
 use url::Url;
 
 use crate::storage_bridge::QuickwitObjectStore;
@@ -69,6 +69,7 @@ pub struct QuickwitObjectStoreRegistry {
     /// fallback.
     default: DefaultObjectStoreRegistry,
     storage_resolver: StorageResolver,
+    storage_cache: Option<Arc<dyn StorageCache>>,
     /// Lazy wrappers constructed by `get_store` on demand, keyed by
     /// `scheme://authority`. Plain `RwLock<HashMap>` is fine — contention
     /// is negligible because the write lock is only taken once per unique
@@ -92,10 +93,16 @@ impl QuickwitObjectStoreRegistry {
         Self {
             default: DefaultObjectStoreRegistry::new(),
             storage_resolver,
+            storage_cache: None,
             lazy_stores: RwLock::new(HashMap::new()),
         }
     }
 
+    pub fn with_storage_cache(mut self, storage_cache: Arc<dyn StorageCache>) -> Self {
+        self.storage_cache = Some(storage_cache);
+        self
+    }
+
     /// Canonical cache key mirroring DataFusion's `DefaultObjectStoreRegistry`:
     /// `scheme://authority`. Preserves the authority so indexes in
     /// different buckets stay distinct; paths within an authority share
@@ -138,8 +145,11 @@ impl ObjectStoreRegistry for QuickwitObjectStoreRegistry {
                 "failed to build Quickwit URI from `{key}`: {err}"
             ))))
         })?;
-        let store: Arc<dyn ObjectStore> =
-            Arc::new(QuickwitObjectStore::new(uri, self.storage_resolver.clone()));
+        let mut quickwit_store = QuickwitObjectStore::new(uri, self.storage_resolver.clone());
+        if let Some(storage_cache) = &self.storage_cache {
+            quickwit_store = quickwit_store.with_storage_cache(Arc::clone(storage_cache));
+        }
+        let store: Arc<dyn ObjectStore> = Arc::new(quickwit_store);
         let mut write = self
             .lazy_stores
             .write()
diff --git a/quickwit/quickwit-datafusion/src/storage_bridge.rs b/quickwit/quickwit-datafusion/src/storage_bridge.rs
@@ -30,6 +30,8 @@
 //! operations return `NotSupported` — DataFusion only reads parquet files
 //! through this store.
 
+use std::ops::Range;
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -43,7 +45,9 @@ use object_store::{
     Result as ObjectStoreResult,
 };
 use quickwit_common::uri::Uri;
-use quickwit_storage::{Storage, StorageResolver};
+use quickwit_storage::{
+    OwnedBytes, Storage, StorageCache, StorageResolver, wrap_storage_with_cache,
+};
 use tokio::sync::OnceCell;
 
 /// Adapts Quickwit's `Storage` trait to DataFusion's `ObjectStore` interface.
@@ -54,6 +58,7 @@ use tokio::sync::OnceCell;
 pub struct QuickwitObjectStore {
     index_uri: Uri,
     storage_resolver: StorageResolver,
+    storage_cache: Option<Arc<dyn StorageCache>>,
     storage: OnceCell<Arc<dyn Storage>>,
 }
 
@@ -62,27 +67,80 @@ impl QuickwitObjectStore {
         Self {
             index_uri,
             storage_resolver,
+            storage_cache: None,
             storage: OnceCell::new(),
         }
     }
 
+    pub fn with_storage_cache(mut self, storage_cache: Arc<dyn StorageCache>) -> Self {
+        self.storage_cache = Some(storage_cache);
+        self
+    }
+
     /// Returns the handle to the underlying `Storage`, resolving it via the
     /// `StorageResolver` if this is the first call.
     async fn storage(&self) -> ObjectStoreResult<&Arc<dyn Storage>> {
         self.storage
             .get_or_try_init(|| async {
-                self.storage_resolver
+                let storage = self
+                    .storage_resolver
                     .resolve(&self.index_uri)
                     .await
                     .map_err(|err| object_store::Error::Generic {
                         store: "QuickwitObjectStore",
                         source: Box::new(err),
-                    })
+                    })?;
+                Ok(match &self.storage_cache {
+                    Some(cache) => wrap_storage_with_cache(
+                        Arc::new(ScopedStorageCache::new(
+                            self.index_uri.as_str().to_string(),
+                            Arc::clone(cache),
+                        )),
+                        storage,
+                    ),
+                    None => storage,
+                })
             })
             .await
     }
 }
 
+struct ScopedStorageCache {
+    scope: String,
+    inner: Arc<dyn StorageCache>,
+}
+
+impl ScopedStorageCache {
+    fn new(scope: String, inner: Arc<dyn StorageCache>) -> Self {
+        Self { scope, inner }
+    }
+
+    fn scoped_path(&self, path: &Path) -> PathBuf {
+        PathBuf::from(format!("{}#{}", self.scope, path.to_string_lossy()))
+    }
+}
+
+#[async_trait]
+impl StorageCache for ScopedStorageCache {
+    async fn get(&self, path: &Path, byte_range: Range<usize>) -> Option<OwnedBytes> {
+        self.inner.get(&self.scoped_path(path), byte_range).await
+    }
+
+    async fn get_all(&self, path: &Path) -> Option<OwnedBytes> {
+        self.inner.get_all(&self.scoped_path(path)).await
+    }
+
+    async fn put(&self, path: PathBuf, byte_range: Range<usize>, bytes: OwnedBytes) {
+        self.inner
+            .put(self.scoped_path(&path), byte_range, bytes)
+            .await
+    }
+
+    async fn put_all(&self, path: PathBuf, bytes: OwnedBytes) {
+        self.inner.put_all(self.scoped_path(&path), bytes).await
+    }
+}
+
 impl std::fmt::Debug for QuickwitObjectStore {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("QuickwitObjectStore")
@@ -264,3 +322,73 @@ impl ObjectStore for QuickwitObjectStore {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Mutex;
+
+    use super::*;
+
+    #[derive(Default)]
+    struct RecordingStorageCache {
+        paths: Mutex<Vec<PathBuf>>,
+    }
+
+    impl RecordingStorageCache {
+        fn paths(&self) -> Vec<PathBuf> {
+            self.paths.lock().unwrap().clone()
+        }
+
+        fn record(&self, path: &Path) {
+            self.paths.lock().unwrap().push(path.to_path_buf());
+        }
+    }
+
+    #[async_trait]
+    impl StorageCache for RecordingStorageCache {
+        async fn get(&self, path: &Path, _byte_range: Range<usize>) -> Option<OwnedBytes> {
+            self.record(path);
+            None
+        }
+
+        async fn get_all(&self, path: &Path) -> Option<OwnedBytes> {
+            self.record(path);
+            None
+        }
+
+        async fn put(&self, path: PathBuf, _byte_range: Range<usize>, _bytes: OwnedBytes) {
+            self.record(&path);
+        }
+
+        async fn put_all(&self, path: PathBuf, _bytes: OwnedBytes) {
+            self.record(&path);
+        }
+    }
+
+    #[tokio::test]
+    async fn scoped_storage_cache_prefixes_cache_keys() {
+        let inner = Arc::new(RecordingStorageCache::default());
+        let inner_cache: Arc<dyn StorageCache> = inner.clone();
+        let cache = ScopedStorageCache::new("s3://metrics-bucket".to_string(), inner_cache);
+
+        cache
+            .get(Path::new("indexes/metrics.parquet"), 10..20)
+            .await;
+        cache
+            .put(
+                Path::new("indexes/metrics.parquet").to_path_buf(),
+                10..20,
+                OwnedBytes::new(&b"bytes"[..]),
+            )
+            .await;
+
+        let paths = inner.paths();
+        assert_eq!(
+            paths,
+            vec![
+                PathBuf::from("s3://metrics-bucket#indexes/metrics.parquet"),
+                PathBuf::from("s3://metrics-bucket#indexes/metrics.parquet"),
+            ]
+        );
+    }
+}
diff --git a/quickwit/quickwit-serve/src/datafusion_api/setup.rs b/quickwit/quickwit-serve/src/datafusion_api/setup.rs
@@ -20,16 +20,19 @@
 
 use std::collections::BTreeSet;
 use std::net::SocketAddr;
+use std::ops::Range;
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::Duration;
 
 use anyhow::Context;
+use async_trait::async_trait;
 use bytesize::ByteSize;
 use futures::{StreamExt, stream};
 use quickwit_cluster::{ClusterChange, ClusterChangeStream, ClusterNode};
 use quickwit_common::tower::Change;
-use quickwit_config::NodeConfig;
 use quickwit_config::service::QuickwitService;
+use quickwit_config::{CacheConfig, NodeConfig};
 use quickwit_datafusion::grpc::DataFusionServiceGrpcImpl;
 use quickwit_datafusion::proto::data_fusion_service_server::{
     DataFusionServiceServer, SERVICE_NAME as DATAFUSION_SERVICE_NAME,
@@ -41,7 +44,7 @@ use quickwit_datafusion::{
 };
 use quickwit_proto::metastore::MetastoreServiceClient;
 use quickwit_search::{SearchServiceClient, SearcherPool, create_search_client_from_grpc_addr};
-use quickwit_storage::StorageResolver;
+use quickwit_storage::{MemorySizedCache, OwnedBytes, StorageCache, StorageResolver};
 use tokio::time::timeout;
 use tonic::transport::server::Router;
 use tonic_reflection::pb::v1::ServerReflectionRequest;
@@ -51,6 +54,10 @@ use tonic_reflection::pb::v1::server_reflection_response::MessageResponse;
 
 use crate::QuickwitServices;
 
+const DATAFUSION_PARQUET_RANGE_CACHE_CAPACITY_ENV: &str =
+    "QW_DATAFUSION_PARQUET_RANGE_CACHE_CAPACITY";
+const FULL_SLICE: Range<usize> = 0..usize::MAX;
+
 /// Build the generic DataFusion session builder for this node.
 ///
 /// Returns `None` if the searcher role is disabled or the endpoint toggle is
@@ -84,7 +91,13 @@ pub(crate) fn build_datafusion_session_builder(
     );
     let worker_resolver = QuickwitWorkerResolver::new(datafusion_worker_pool)
         .with_tls(node_config.grpc_config.tls.is_some());
-    let registry = Arc::new(QuickwitObjectStoreRegistry::new(storage_resolver));
+    let datafusion_parquet_range_cache = datafusion_parquet_range_cache_config();
+    let datafusion_storage_cache =
+        DataFusionParquetRangeCache::from_config(&datafusion_parquet_range_cache);
+    let registry = Arc::new(
+        QuickwitObjectStoreRegistry::new(storage_resolver)
+            .with_storage_cache(Arc::new(datafusion_storage_cache)),
+    );
     let builder = DataFusionSessionBuilder::new()
         .with_object_store_registry(registry)
         .context("failed to install DataFusion object store registry")?
@@ -97,6 +110,49 @@ pub(crate) fn build_datafusion_session_builder(
     Ok(Some(Arc::new(builder)))
 }
 
+fn datafusion_parquet_range_cache_config() -> CacheConfig {
+    let capacity = quickwit_common::get_from_env(
+        DATAFUSION_PARQUET_RANGE_CACHE_CAPACITY_ENV,
+        ByteSize::gb(4),
+        false,
+    );
+    CacheConfig::default_with_capacity(capacity)
+}
+
+struct DataFusionParquetRangeCache {
+    inner: MemorySizedCache,
+}
+
+impl DataFusionParquetRangeCache {
+    fn from_config(cache_config: &CacheConfig) -> Self {
+        Self {
+            inner: MemorySizedCache::from_config(
+                cache_config,
+                &quickwit_storage::STORAGE_METRICS.datafusion_parquet_range_cache,
+            ),
+        }
+    }
+}
+
+#[async_trait]
+impl StorageCache for DataFusionParquetRangeCache {
+    async fn get(&self, path: &Path, byte_range: Range<usize>) -> Option<OwnedBytes> {
+        self.inner.get_slice(path, byte_range)
+    }
+
+    async fn get_all(&self, path: &Path) -> Option<OwnedBytes> {
+        self.inner.get_slice(path, FULL_SLICE)
+    }
+
+    async fn put(&self, path: PathBuf, byte_range: Range<usize>, bytes: OwnedBytes) {
+        self.inner.put_slice(path, byte_range, bytes);
+    }
+
+    async fn put_all(&self, path: PathBuf, bytes: OwnedBytes) {
+        self.inner.put_slice(path, FULL_SLICE, bytes);
+    }
+}
+
 fn setup_datafusion_worker_pool(
     cluster_change_stream: ClusterChangeStream,
     max_message_size: ByteSize,
diff --git a/quickwit/quickwit-storage/src/metrics.rs b/quickwit/quickwit-storage/src/metrics.rs
@@ -30,6 +30,7 @@ pub struct StorageMetrics {
     pub predicate_cache: CacheMetrics,
     pub fd_cache_metrics: CacheMetrics,
     pub fast_field_cache: CacheMetrics,
+    pub datafusion_parquet_range_cache: CacheMetrics,
     pub split_footer_cache: CacheMetrics,
     pub searcher_split_cache: CacheMetrics,
     pub get_slice_timeout_successes: [IntCounter; 3],
@@ -94,6 +95,7 @@ impl Default for StorageMetrics {
 
         StorageMetrics {
             fast_field_cache: CacheMetrics::for_component("fastfields"),
+            datafusion_parquet_range_cache: CacheMetrics::for_component("datafusion_parquet_range"),
             fd_cache_metrics: CacheMetrics::for_component("fd"),
             partial_request_cache: CacheMetrics::for_component("partial_request"),
             predicate_cache: CacheMetrics::for_component("predicate"),
