feat: share Parquet metadata cache across tasks in an executor

.github/workflows/pr_build_linux.yml

@@ -303,6 +303,8 @@ jobs:
           - name: "scans"
             value: |
               org.apache.comet.parquet.CometParquetWriterSuite
+              org.apache.comet.CometMetadataCacheSuite
+              org.apache.comet.serde.CometFilePartitionSerdeSuite
               org.apache.comet.parquet.ParquetReadV1Suite
               org.apache.comet.parquet.ParquetReadV2Suite
               org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite

.github/workflows/pr_build_macos.yml

@@ -119,6 +119,8 @@ jobs:
           - name: "scans"
             value: |
               org.apache.comet.parquet.CometParquetWriterSuite
+              org.apache.comet.CometMetadataCacheSuite
+              org.apache.comet.serde.CometFilePartitionSerdeSuite
               org.apache.comet.parquet.ParquetReadV1Suite
               org.apache.comet.parquet.ParquetReadV2Suite
               org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite

docs/source/user-guide/latest/tuning.md

@@ -103,6 +103,22 @@ Comet Performance
 
 It may be possible to reduce Comet's memory overhead by reducing batch sizes or increasing number of partitions.
 
+## Parquet Metadata Cache
+
+Comet's native Parquet scan caches each file's footer and page index so that tasks reading
+different splits of the same file, or later stages re-reading the same table, do not re-fetch
+and re-parse it. The cache is shared by all tasks in an executor process and is enabled by
+default.
+
+`spark.comet.scan.metadataCache.memoryLimit` (default 50MB) bounds, with least-recently-used
+eviction, the cache for each distinct object store an executor reads from, so an executor
+reading from several object stores, or several buckets, holds a multiple of this limit. It does
+not scale with the number of task slots, however. Entries are validated against each file's size
+and modification time (validation falls back to size alone if the source reports no modification
+time), so an overwritten file is normally re-read rather than served from the cache.
+
+Set `spark.comet.scan.metadataCache.enabled=false` to give each task its own cache instead.
+
 ## Optimizing Sorting on Floating-Point Values
 
 Sorting on floating-point data types (or complex types containing floating-point values) is not compatible with

native/core/Cargo.toml

@@ -36,6 +36,7 @@ publish = false
 
 [dependencies]
 arrow = { workspace = true }
+chrono = { workspace = true }
 parquet = { workspace = true, default-features = false, features = ["experimental", "arrow", "snap", "lz4", "zstd", "flate2-zlib-rs"] }
 futures = { workspace = true }
 mimalloc = { version = "*", default-features = false, optional = true }

native/core/src/execution/jni_api.rs

@@ -109,10 +109,12 @@ use crate::execution::tracing::{
 use crate::execution::memory_pools::logging_pool::LoggingMemoryPool;
 use crate::execution::spark_config::{
     SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED,
-    COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED,
-    COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES,
+    COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_METADATA_CACHE_ENABLED, COMET_METADATA_CACHE_MEMORY_LIMIT,
+    COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES,
 };
 use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID};
+use crate::parquet::metadata_cache::configure_metadata_cache;
+use datafusion::execution::cache::cache_manager::DEFAULT_METADATA_CACHE_LIMIT;
 use datafusion_comet_proto::spark_operator::operator::OpStruct;
 use log::info;
 use std::sync::OnceLock;
@@ -595,6 +597,16 @@ fn prepare_datafusion_session_context(
             .set_str("datafusion.execution.parquet.reorder_filters", "true");
     }
 
+    // The metadata cache is process-wide rather than per-RuntimeEnv, so apply Spark's settings
+    // to the registry on every task. `configure` is idempotent.
+    configure_metadata_cache(
+        spark_config.get_bool_or(COMET_METADATA_CACHE_ENABLED, true),
+        spark_config.get_usize(
+            COMET_METADATA_CACHE_MEMORY_LIMIT,
+            DEFAULT_METADATA_CACHE_LIMIT,
+        ),
+    );
+
     let runtime = rt_config.build()?;
 
     let mut session_ctx = SessionContext::new_with_config_rt(session_config, Arc::new(runtime));

native/core/src/execution/planner.rs

@@ -261,6 +261,12 @@ impl PhysicalPlanner {
             // Convert that to a Path object to use in the PartitionedFile.
             let path = Path::from_url_path(url.path()).map_err(|e| GeneralError(e.to_string()))?;
             partitioned_file.object_meta.location = path;
+            // Spark sends the file's modification time in milliseconds. DataFusion validates
+            // cached Parquet metadata with (size, last_modified), so this must be real for the
+            // shared metadata cache to detect an overwritten file.
+            partitioned_file.object_meta.last_modified =
+                chrono::DateTime::from_timestamp_millis(file.modification_time)
+                    .unwrap_or_else(|| chrono::DateTime::from_timestamp_nanos(0));
 
             // Process partition values
             // Create an empty input schema for partition values because they are all literals.
@@ -5442,4 +5448,28 @@ mod tests {
             }
         });
     }
+
+    #[test]
+    fn test_partitioned_file_carries_modification_time() {
+        let session_ctx = SessionContext::new();
+        let planner = PhysicalPlanner::new(Arc::from(session_ctx), 0);
+
+        let partition = spark_operator::SparkFilePartition {
+            partitioned_file: vec![spark_operator::SparkPartitionedFile {
+                file_path: "file:///tmp/comet-test/part-0.parquet".to_string(),
+                start: 0,
+                length: 100,
+                file_size: 100,
+                partition_values: vec![],
+                modification_time: 1_700_000_000_123,
+            }],
+        };
+
+        let files = planner.get_partitioned_files(&partition).unwrap();
+        assert_eq!(files.len(), 1);
+        assert_eq!(
+            files[0].object_meta.last_modified.timestamp_millis(),
+            1_700_000_000_123
+        );
+    }
 }

native/core/src/execution/spark_config.rs

@@ -25,9 +25,13 @@ pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory";
 pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str =
     "spark.comet.parquet.rowFilterPushdown.enabled";
 pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores";
+pub(crate) const COMET_METADATA_CACHE_ENABLED: &str = "spark.comet.scan.metadataCache.enabled";
+pub(crate) const COMET_METADATA_CACHE_MEMORY_LIMIT: &str =
+    "spark.comet.scan.metadataCache.memoryLimit";
 
 pub(crate) trait SparkConfig {
     fn get_bool(&self, name: &str) -> bool;
+    fn get_bool_or(&self, name: &str, default_value: bool) -> bool;
     fn get_u64(&self, name: &str, default_value: u64) -> u64;
     fn get_usize(&self, name: &str, default_value: usize) -> usize;
 }
@@ -39,6 +43,12 @@ impl SparkConfig for HashMap<String, String> {
             .unwrap_or(false)
     }
 
+    fn get_bool_or(&self, name: &str, default_value: bool) -> bool {
+        self.get(name)
+            .and_then(|str_val| str_val.parse::<bool>().ok())
+            .unwrap_or(default_value)
+    }
+
     fn get_u64(&self, name: &str, default_value: u64) -> u64 {
         self.get(name)
             .and_then(|str_val| str_val.parse::<u64>().ok())

native/core/src/parquet/metadata_cache.rs

@@ -0,0 +1,227 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datafusion::execution::cache::cache_manager::{
+    FileMetadataCache, DEFAULT_METADATA_CACHE_LIMIT,
+};
+use datafusion::execution::cache::DefaultFilesMetadataCache;
+use datafusion::execution::object_store::ObjectStoreUrl;
+use std::collections::HashMap;
+use std::sync::{Arc, OnceLock, PoisonError, RwLock};
+
+/// Process-wide cache of Parquet file metadata (footer + page index), keyed by object store.
+///
+/// ## Why process lifetime?
+///
+/// Comet builds a fresh DataFusion `RuntimeEnv` for every Spark task, so the `FileMetadataCache`
+/// that `RuntimeEnv` owns is discarded when the task ends and two tasks on the same executor
+/// re-fetch and re-parse the same footer. There is no executor-scoped Rust object with a longer
+/// lifetime to own this cache, so the process is the natural scope, as it already is for
+/// `parquet_support::object_store_cache`.
+///
+/// ## Why keyed by object store URL?
+///
+/// The cache key inside `DefaultFilesMetadataCache` is a bare `object_store::path::Path`, and
+/// `PhysicalPlanner::get_partitioned_files` builds that path from `url.path()`, which drops the
+/// bucket for `s3://bucket/key`. A single shared instance would therefore let two buckets that
+/// share a relative path collide. One cache per `scheme://host` separates buckets on S3/GCS and
+/// separates hosts on HDFS. It does not separate ABFS containers on the same storage account:
+/// `abfss://container@account.dfs.core.windows.net/...` carries the container as URL userinfo,
+/// not host, so `prepare_object_store_with_configs` derives the same registry key, and the same
+/// `Path`, for every container on that account. Two such containers with colliding relative
+/// paths rely entirely on `(size, last_modified)` validation, below, to keep their entries
+/// apart. Unlike the object store cache, the key deliberately omits the credential config hash:
+/// a footer's identity does not depend on which credentials read it.
+///
+/// ## Staleness
+///
+/// Entries are validated by DataFusion against `(size, last_modified)` on every lookup, and
+/// Comet populates both from Spark's `PartitionedFile`. Eviction is the LRU policy of
+/// `DefaultFilesMetadataCache`, bounded by the configured memory limit.
+pub(crate) struct MetadataCacheRegistry {
+    config: RwLock<MetadataCacheConfig>,
+    caches: RwLock<HashMap<String, Arc<dyn FileMetadataCache>>>,
+}
+
+#[derive(Clone, Copy)]
+struct MetadataCacheConfig {
+    enabled: bool,
+    memory_limit: usize,
+}
+
+impl Default for MetadataCacheConfig {
+    fn default() -> Self {
+        Self {
+            enabled: true,
+            memory_limit: DEFAULT_METADATA_CACHE_LIMIT,
+        }
+    }
+}
+
+impl MetadataCacheRegistry {
+    pub(crate) fn new() -> Self {
+        Self {
+            config: RwLock::new(MetadataCacheConfig::default()),
+            caches: RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Applies the Spark-provided settings. Called once per task, so it must be idempotent.
+    ///
+    /// The `config` write guard is held for the whole operation, including the loop that applies
+    /// the new limit to live caches, so that concurrent `configure` calls serialize instead of
+    /// racing: without that, two calls with different limits could write `config` in one order
+    /// and update the live caches in the other, leaving a cache on a stale limit that the
+    /// idempotency check above would then never revisit. `cache_for` also only ever takes
+    /// `config` before `caches`, and now holds `config` until after it has inserted a new cache,
+    /// so acquiring `caches` here while still holding `config` preserves that same lock ordering
+    /// and cannot deadlock. It does mean a concurrent `cache_for` that is still inserting a new
+    /// cache will block this call until that insert is visible, which is what stops a fresh cache
+    /// from being seeded with a limit this call is about to make stale.
+    pub(crate) fn configure(&self, enabled: bool, memory_limit: usize) {
+        let mut config = self.config.write().unwrap_or_else(PoisonError::into_inner);
+        if config.enabled == enabled && config.memory_limit == memory_limit {
+            return;
+        }
+        config.enabled = enabled;
+        config.memory_limit = memory_limit;
+
+        let caches = self.caches.read().unwrap_or_else(PoisonError::into_inner);
+        for cache in caches.values() {
+            cache.update_cache_limit(memory_limit);
+        }
+    }
+
+    /// Returns the shared cache for `url`, or `None` when sharing is disabled, in which case the
+    /// caller should fall back to the per-task `RuntimeEnv` cache.
+    ///
+    /// The `config` read guard is held for the whole method, including the `caches` write lock
+    /// used to insert a newly created cache, so that a concurrent `configure` cannot land between
+    /// reading `config.memory_limit` here and the insert becoming visible to `configure`'s cache
+    /// scan. Without that, `configure` could apply a new limit to every cache already in the map
+    /// and return, while this method then inserts a cache built from the limit it read before
+    /// `configure` ran, leaving that cache permanently on the stale limit and the registry's
+    /// config permanently disagreeing with it. Taking `config` first and only then `caches`
+    /// (read, then write) keeps the same `config`-before-`caches` ordering `configure` relies on,
+    /// so this cannot deadlock with it.
+    pub(crate) fn cache_for(&self, url: &ObjectStoreUrl) -> Option<Arc<dyn FileMetadataCache>> {
+        let config = self.config.read().unwrap_or_else(PoisonError::into_inner);
+        if !config.enabled {
+            return None;
+        }
+
+        let key = url.as_str();
+        {
+            let caches = self.caches.read().unwrap_or_else(PoisonError::into_inner);
+            if let Some(cache) = caches.get(key) {
+                return Some(Arc::clone(cache));
+            }
+        }
+
+        let mut caches = self.caches.write().unwrap_or_else(PoisonError::into_inner);
+        let cache = caches
+            .entry(key.to_string())
+            .or_insert_with(|| Arc::new(DefaultFilesMetadataCache::new(config.memory_limit)));
+        Some(Arc::clone(cache))
+    }
+}
+
+fn registry() -> &'static MetadataCacheRegistry {
+    static REGISTRY: OnceLock<MetadataCacheRegistry> = OnceLock::new();
+    REGISTRY.get_or_init(MetadataCacheRegistry::new)
+}
+
+/// Returns the process-wide metadata cache for `url`, or `None` when sharing is disabled.
+pub(crate) fn shared_metadata_cache(url: &ObjectStoreUrl) -> Option<Arc<dyn FileMetadataCache>> {
+    registry().cache_for(url)
+}
+
+/// Applies Spark's metadata cache settings to the process-wide registry.
+pub(crate) fn configure_metadata_cache(enabled: bool, memory_limit: usize) {
+    registry().configure(enabled, memory_limit);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn url(s: &str) -> ObjectStoreUrl {
+        ObjectStoreUrl::parse(s).unwrap()
+    }
+
+    #[test]
+    fn same_url_returns_the_same_cache() {
+        let registry = MetadataCacheRegistry::new();
+        let first = registry.cache_for(&url("s3://bucket-a")).unwrap();
+        let second = registry.cache_for(&url("s3://bucket-a")).unwrap();
+        assert!(Arc::ptr_eq(&first, &second));
+    }
+
+    #[test]
+    fn different_buckets_get_separate_caches() {
+        let registry = MetadataCacheRegistry::new();
+        let a = registry.cache_for(&url("s3://bucket-a")).unwrap();
+        let b = registry.cache_for(&url("s3://bucket-b")).unwrap();
+        assert!(!Arc::ptr_eq(&a, &b));
+    }
+
+    #[test]
+    fn disabled_registry_returns_none() {
+        let registry = MetadataCacheRegistry::new();
+        registry.configure(false, DEFAULT_METADATA_CACHE_LIMIT);
+        assert!(registry.cache_for(&url("s3://bucket-a")).is_none());
+    }
+
+    #[test]
+    fn default_limit_applies_to_new_caches() {
+        let registry = MetadataCacheRegistry::new();
+        let cache = registry.cache_for(&url("s3://bucket-a")).unwrap();
+        assert_eq!(cache.cache_limit(), DEFAULT_METADATA_CACHE_LIMIT);
+    }
+
+    #[test]
+    fn configure_updates_the_limit_of_live_caches() {
+        let registry = MetadataCacheRegistry::new();
+        let cache = registry.cache_for(&url("s3://bucket-a")).unwrap();
+        registry.configure(true, 1024);
+        assert_eq!(cache.cache_limit(), 1024);
+        // and newly created caches pick up the configured limit
+        let other = registry.cache_for(&url("s3://bucket-b")).unwrap();
+        assert_eq!(other.cache_limit(), 1024);
+    }
+
+    #[test]
+    fn concurrent_cache_for_returns_the_same_instance() {
+        let registry = Arc::new(MetadataCacheRegistry::new());
+        let handles: Vec<_> = (0..8)
+            .map(|_| {
+                let registry = Arc::clone(&registry);
+                std::thread::spawn(move || registry.cache_for(&url("s3://bucket-a")).unwrap())
+            })
+            .collect();
+
+        let caches: Vec<_> = handles
+            .into_iter()
+            .map(|handle| handle.join().unwrap())
+            .collect();
+
+        let first = &caches[0];
+        for cache in &caches[1..] {
+            assert!(Arc::ptr_eq(first, cache));
+        }
+    }
+}