Cache initial reads in VortexFormat (#1633)

.gitignore

@@ -204,4 +204,7 @@ benchmarks/.out
 # Zed
 .zed/
 
-*.plan
+*.plan
+
+# Default samply profile
+profile.json.gz

Cargo.toml

@@ -95,6 +95,7 @@ jiff = "0.1.8"
 libfuzzer-sys = "0.4"
 log = "0.4.21"
 mimalloc = "0.1.42"
+moka = "0.12"
 num-traits = "0.2.18"
 num_enum = "0.7.2"
 object_store = "0.11.0"

bench-vortex/Cargo.toml

@@ -46,14 +46,14 @@ mimalloc = { workspace = true }
 object_store = { workspace = true, features = ["aws"] }
 parquet = { workspace = true, features = ["async"] }
 prettytable-rs = { workspace = true }
-tabled = { workspace = true, features = ["std"] }
 rand = { workspace = true }
 rayon = { workspace = true }
 regex = { workspace = true }
 reqwest = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 simplelog = { workspace = true }
+tabled = { workspace = true, features = ["std"] }
 tar = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 uuid = { workspace = true, features = ["v4"] }

fuzz/Cargo.toml

@@ -14,7 +14,6 @@ rust-version = { workspace = true }
 categories = { workspace = true }
 readme = "README.md"
 
-
 [package.metadata]
 cargo-fuzz = true
 

vortex-datafusion/Cargo.toml

@@ -21,6 +21,7 @@ path = "src/lib.rs"
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
 async-trait = { workspace = true }
+bytes = { workspace = true }
 chrono = { workspace = true }
 datafusion = { workspace = true }
 datafusion-common = { workspace = true }
@@ -31,8 +32,10 @@ datafusion-physical-plan = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
+moka = { workspace = true, features = ["future", "sync"] }
 object_store = { workspace = true }
 pin-project = { workspace = true }
+tokio = { workspace = true }
 vortex-array = { workspace = true }
 vortex-datetime-dtype = { workspace = true }
 vortex-dtype = { workspace = true }

vortex-datafusion/src/persistent/cache.rs

@@ -0,0 +1,65 @@
+use std::sync::Arc;
+
+use chrono::{DateTime, Utc};
+use moka::future::Cache;
+use object_store::path::Path;
+use object_store::{ObjectMeta, ObjectStore};
+use vortex_error::{vortex_err, VortexError, VortexResult};
+use vortex_file::{read_initial_bytes, InitialRead};
+use vortex_io::ObjectStoreReadAt;
+
+#[derive(Debug, Clone)]
+pub struct InitialReadCache {
+    inner: Cache<Key, InitialRead>,
+}
+
+impl Default for InitialReadCache {
+    fn default() -> Self {
+        let inner = Cache::builder()
+            .weigher(|k: &Key, v: &InitialRead| {
+                (k.location.as_ref().as_bytes().len() + v.buf.len()) as u32
+            })
+            .max_capacity(256 * (2 << 20))
+            .eviction_listener(|k, _v, cause| {
+                log::trace!("Removed {} due to {:?}", k.location, cause);
+            })
+            .build();
+
+        Self { inner }
+    }
+}
+
+#[derive(Hash, Eq, PartialEq, Debug)]
+pub struct Key {
+    location: Path,
+    m_time: DateTime<Utc>,
+}
+
+impl From<&ObjectMeta> for Key {
+    fn from(value: &ObjectMeta) -> Self {
+        Self {
+            location: value.location.clone(),
+            m_time: value.last_modified,
+        }
+    }
+}
+
+impl InitialReadCache {
+    pub async fn try_get(
+        &self,
+        object: &ObjectMeta,
+        store: Arc<dyn ObjectStore>,
+    ) -> VortexResult<InitialRead> {
+        self.inner
+            .try_get_with(Key::from(object), async {
+                let os_read_at = ObjectStoreReadAt::new(store.clone(), object.location.clone());
+                let initial_read = read_initial_bytes(&os_read_at, object.size as u64).await?;
+                VortexResult::Ok(initial_read)
+            })
+            .await
+            .map_err(|e: Arc<VortexError>| match Arc::try_unwrap(e) {
+                Ok(e) => e,
+                Err(e) => vortex_err!("{}", e.to_string()),
+            })
+    }
+}

vortex-datafusion/src/persistent/execution.rs

@@ -14,6 +14,7 @@ use datafusion_physical_plan::{
 use itertools::Itertools;
 use vortex_array::Context;
 
+use super::cache::InitialReadCache;
 use crate::persistent::opener::VortexFileOpener;
 
 #[derive(Debug, Clone)]
@@ -24,6 +25,7 @@ pub struct VortexExec {
     plan_properties: PlanProperties,
     projected_statistics: Statistics,
     ctx: Arc<Context>,
+    initial_read_cache: InitialReadCache,
 }
 
 impl VortexExec {
@@ -32,6 +34,7 @@ impl VortexExec {
         metrics: ExecutionPlanMetricsSet,
         predicate: Option<Arc<dyn PhysicalExpr>>,
         ctx: Arc<Context>,
+        initial_read_cache: InitialReadCache,
     ) -> DFResult<Self> {
         let projected_schema = project_schema(
             &file_scan_config.file_schema,
@@ -60,6 +63,7 @@ impl VortexExec {
             plan_properties,
             projected_statistics,
             ctx,
+            initial_read_cache,
         })
     }
 
@@ -118,6 +122,7 @@ impl ExecutionPlan for VortexExec {
             object_store,
             projection: self.file_scan_config.projection.clone(),
             predicate: self.predicate.clone(),
+            initial_read_cache: self.initial_read_cache.clone(),
             arrow_schema,
         };
         let stream = FileStream::new(&self.file_scan_config, partition, opener, &self.metrics)?;