From eb17e71ac0aa4e2a7c6924594a31d4a77b338377 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 08:59:25 -0600 Subject: [PATCH 01/10] feat: carry Parquet file modification time to the native scan --- native/Cargo.lock | 1 + native/core/Cargo.toml | 1 + native/core/src/execution/planner.rs | 30 +++++++++++++ native/proto/src/proto/operator.proto | 3 ++ .../apache/comet/serde/operator/package.scala | 1 + .../serde/CometFilePartitionSerdeSuite.scala | 43 +++++++++++++++++++ 6 files changed, 79 insertions(+) create mode 100644 spark/src/test/scala/org/apache/comet/serde/CometFilePartitionSerdeSuite.scala diff --git a/native/Cargo.lock b/native/Cargo.lock index fc04499b51..9e60b76646 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1933,6 +1933,7 @@ dependencies = [ "async-trait", "aws-config", "aws-credential-types", + "chrono", "criterion", "datafusion", "datafusion-comet-common", diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 7be91683b8..2ff8915403 100644 --- a/native/core/Cargo.toml +++ b/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 } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a23a7b1f67..b78e22136c 100644 --- a/native/core/src/execution/planner.rs +++ b/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 + ); + } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 00158644d1..1f49a2c2f8 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -64,6 +64,9 @@ message SparkPartitionedFile { int64 length = 3; int64 file_size = 4; repeated spark.spark_expression.Expr partition_values = 5; + // Milliseconds since the epoch, from Spark's PartitionedFile.modificationTime. + // Used with file_size to validate entries in the shared Parquet metadata cache. + int64 modification_time = 6; } // This name and the one above are not great, but they correspond to the (unfortunate) Spark names. diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala index 4bdfc58046..885ab4728b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala @@ -71,6 +71,7 @@ package object operator { .setStart(file.start) .setLength(file.length) .setFileSize(file.fileSize) + .setModificationTime(file.modificationTime) partitionBuilder.addPartitionedFile(fileBuilder.build()) }) partitionBuilder.build() diff --git a/spark/src/test/scala/org/apache/comet/serde/CometFilePartitionSerdeSuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometFilePartitionSerdeSuite.scala new file mode 100644 index 0000000000..f02916f833 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/serde/CometFilePartitionSerdeSuite.scala @@ -0,0 +1,43 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.execution.datasources.FilePartition +import org.apache.spark.sql.types.StructType + +import org.apache.comet.serde.operator.partition2Proto +import org.apache.comet.shims.ShimBatchReader + +class CometFilePartitionSerdeSuite extends AnyFunSuite { + + test("file modification time is serialized to the native side") { + val file = ShimBatchReader + .newPartitionedFile(InternalRow.empty, "file:///tmp/comet/part-0.parquet") + .copy(start = 0, length = 100, fileSize = 100, modificationTime = 1700000000123L) + + val proto = partition2Proto(FilePartition(0, Array(file)), StructType(Seq.empty)) + + assert(proto.getPartitionedFileCount == 1) + assert(proto.getPartitionedFile(0).getModificationTime == 1700000000123L) + } +} From 2e90bb95aacb975c4a388c985046fbc6bbed4391 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 09:08:29 -0600 Subject: [PATCH 02/10] feat: add process-wide Parquet metadata cache registry --- native/core/src/parquet/metadata_cache.rs | 181 ++++++++++++++++++++++ native/core/src/parquet/mod.rs | 1 + 2 files changed, 182 insertions(+) create mode 100644 native/core/src/parquet/metadata_cache.rs diff --git a/native/core/src/parquet/metadata_cache.rs b/native/core/src/parquet/metadata_cache.rs new file mode 100644 index 0000000000..075d63a575 --- /dev/null +++ b/native/core/src/parquet/metadata_cache.rs @@ -0,0 +1,181 @@ +// 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` removes that entirely. 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, + caches: RwLock>>, +} + +#[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. + 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; + drop(config); + + 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. + pub(crate) fn cache_for(&self, url: &ObjectStoreUrl) -> Option> { + 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 = 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> { + 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); + } +} diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index 5644e61c85..dc8c4110ad 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -17,6 +17,7 @@ pub mod encryption_support; +pub mod metadata_cache; pub mod parquet_exec; pub mod parquet_support; pub mod schema_adapter; From 05d9797d83ed98e6cfc0161970a6ad03313fa005 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 09:16:07 -0600 Subject: [PATCH 03/10] fix: serialize metadata cache configure against concurrent calls Hold the config write lock for the whole configure() operation, including the loop that applies the new limit to live caches, so concurrent configure() calls with different limits can no longer race and leave a cache on a stale limit that the idempotency check would then never revisit. Add a test that exercises cache_for() from multiple threads concurrently and asserts they all observe the same cache instance. --- native/core/src/parquet/metadata_cache.rs | 30 ++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/native/core/src/parquet/metadata_cache.rs b/native/core/src/parquet/metadata_cache.rs index 075d63a575..79cc48bab1 100644 --- a/native/core/src/parquet/metadata_cache.rs +++ b/native/core/src/parquet/metadata_cache.rs @@ -76,6 +76,14 @@ impl MetadataCacheRegistry { } /// 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` only ever takes `config` + /// before `caches` and releases `config` first, so acquiring `caches` here while still + /// holding `config` preserves that same lock ordering and cannot deadlock. 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 { @@ -83,7 +91,6 @@ impl MetadataCacheRegistry { } config.enabled = enabled; config.memory_limit = memory_limit; - drop(config); let caches = self.caches.read().unwrap_or_else(PoisonError::into_inner); for cache in caches.values() { @@ -178,4 +185,25 @@ mod tests { 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(®istry); + 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)); + } + } } From d6e0181975cb72aa29f327d08ee8097512095d71 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 09:22:08 -0600 Subject: [PATCH 04/10] fix(native): close race that seeds metadata cache with stale limit cache_for copied config.memory_limit into a local and released the config read guard before acquiring the caches write lock. A concurrent configure() could run entirely in that window, updating config and resyncing every cache already in the map, none of which yet included the cache cache_for was about to insert. The new cache would then be permanently stuck on the old limit, with configure's idempotency short-circuit never revisiting it. Hold the config read guard for the whole of cache_for, including the caches write lock used to insert a new cache, so a concurrent configure blocks until the insert is complete and visible to its cache scan. Acquisition order stays config-before-caches on both paths, so this cannot deadlock with configure. --- native/core/src/parquet/metadata_cache.rs | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/native/core/src/parquet/metadata_cache.rs b/native/core/src/parquet/metadata_cache.rs index 79cc48bab1..a2e546c189 100644 --- a/native/core/src/parquet/metadata_cache.rs +++ b/native/core/src/parquet/metadata_cache.rs @@ -81,9 +81,12 @@ impl MetadataCacheRegistry { /// 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` only ever takes `config` - /// before `caches` and releases `config` first, so acquiring `caches` here while still - /// holding `config` preserves that same lock ordering and cannot deadlock. + /// 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 { @@ -100,8 +103,18 @@ impl MetadataCacheRegistry { /// 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> { - let config = *self.config.read().unwrap_or_else(PoisonError::into_inner); + let config = self.config.read().unwrap_or_else(PoisonError::into_inner); if !config.enabled { return None; } From 5eae24a82b43a64e03d7cf653e886fefff857374 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 09:28:23 -0600 Subject: [PATCH 05/10] feat: share Parquet metadata cache across Spark tasks --- native/core/src/parquet/parquet_exec.rs | 168 ++++++++++++++++++++++-- 1 file changed, 155 insertions(+), 13 deletions(-) diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index d8704d80c0..2a61a5e936 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -17,6 +17,7 @@ use crate::execution::operators::ExecutionError; use crate::parquet::encryption_support::{CometEncryptionConfig, ENCRYPTION_FACTORY_ID}; +use crate::parquet::metadata_cache::shared_metadata_cache; use crate::parquet::parquet_support::SparkParquetOptions; use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory; use arrow::datatypes::{Field, SchemaRef}; @@ -142,16 +143,19 @@ pub(crate) fn init_datasource_exec( ); } - // DataFusion's metadata-caching reader factory: loads each file's full metadata (including the - // page index) once into the per-task RuntimeEnv cache (bounded LRU, `metadata_cache_limit`) and - // reuses it across that file's row-group splits, so the opener does not re-fetch the page index. + // DataFusion's metadata-caching reader factory: loads each file's full metadata (including + // the page index) once and reuses it across that file's row-group splits, so the opener does + // not re-fetch the page index. The cache is shared process-wide across Spark tasks (see + // `metadata_cache`), falling back to the per-task `RuntimeEnv` cache when sharing is + // disabled. // // TODO: metadata I/O is invisible in metrics. `fetch_metadata` reads via `ObjectStore::get_ranges`, // bypassing the `get_bytes` path where `bytes_scanned` is counted. A byte-counting ObjectStore // wrapper would surface it. let runtime_env = session_ctx.runtime_env(); let store = runtime_env.object_store(&object_store_url)?; - let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let metadata_cache = shared_metadata_cache(&object_store_url) + .unwrap_or_else(|| runtime_env.cache_manager.get_file_metadata_cache()); parquet_source = parquet_source.with_parquet_file_reader_factory(Arc::new( CachedParquetFileReaderFactory::new(store, metadata_cache), )); @@ -246,6 +250,7 @@ fn get_options( #[cfg(test)] mod tests { use super::*; + use crate::parquet::metadata_cache::shared_metadata_cache; use arrow::array::Int32Array; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; @@ -258,9 +263,9 @@ mod tests { use std::fs::File; // Regression test for #3978: the scan's reader factory must load the full - // Parquet metadata, including the page index, into the per-task RuntimeEnv - // metadata cache. The previous hand-rolled factory cached only the footer, - // so DataFusion re-fetched the page index on every open of the same file. + // Parquet metadata, including the page index, into the metadata cache. The + // previous hand-rolled factory cached only the footer, so DataFusion + // re-fetched the page index on every open of the same file. #[tokio::test] async fn caches_full_metadata_with_page_index() { // Write a file with a page index: page-level statistics and a small data @@ -318,12 +323,10 @@ mod tests { batch.unwrap(); } - // The per-task RuntimeEnv metadata cache must now hold this file's - // metadata with the page index (column + offset index) loaded. - let cache = session_ctx - .runtime_env() - .cache_manager - .get_file_metadata_cache(); + // The shared metadata cache must now hold this file's metadata with the page index + // (column + offset index) loaded. + let cache = shared_metadata_cache(&ObjectStoreUrl::local_filesystem()) + .expect("metadata cache sharing is enabled by default"); let entry = cache .get(&location) .expect("file metadata should be cached"); @@ -338,4 +341,143 @@ mod tests { "cached metadata must include the page index" ); } + + // Two Spark tasks on one executor are two independent SessionContexts. The second must + // reuse the metadata the first parsed, rather than re-fetching the footer. + #[tokio::test] + async fn shares_metadata_between_session_contexts() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from((0..100).collect::>()))], + ) + .unwrap(); + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let partitioned_file = PartitionedFile::from_path(filename.clone()).unwrap(); + let location = partitioned_file.object_meta.location.clone(); + + for _ in 0..2 { + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + Arc::clone(&schema), + None, + None, + ObjectStoreUrl::local_filesystem(), + vec![vec![PartitionedFile::from_path(filename.clone()).unwrap()]], + None, + None, + None, + "UTC", + true, + false, + false, + false, + &session_ctx, + false, + false, + false, + ) + .unwrap(); + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + while let Some(batch) = stream.next().await { + batch.unwrap(); + } + } + + let cache = shared_metadata_cache(&ObjectStoreUrl::local_filesystem()) + .expect("metadata cache sharing is enabled by default"); + let entries = cache.list_entries(); + let entry = entries + .get(&location) + .expect("file metadata should be cached"); + assert_eq!( + entry.hits, 1, + "the second SessionContext must hit the cache populated by the first" + ); + } + + // A file overwritten with the same size must not serve stale metadata. DataFusion validates + // cached entries against (size, last_modified), so the cached entry must be replaced. + #[tokio::test] + async fn newer_modification_time_replaces_the_cached_entry() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from((0..100).collect::>()))], + ) + .unwrap(); + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let mut first = PartitionedFile::from_path(filename.clone()).unwrap(); + first.object_meta.last_modified = chrono::DateTime::from_timestamp_millis(1_000).unwrap(); + let location = first.object_meta.location.clone(); + let size = first.object_meta.size; + + let mut second = PartitionedFile::from_path(filename.clone()).unwrap(); + second.object_meta.last_modified = chrono::DateTime::from_timestamp_millis(2_000).unwrap(); + assert_eq!( + second.object_meta.size, size, + "sizes must match for this test" + ); + + for pf in [first, second] { + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + Arc::clone(&schema), + None, + None, + ObjectStoreUrl::local_filesystem(), + vec![vec![pf]], + None, + None, + None, + "UTC", + true, + false, + false, + false, + &session_ctx, + false, + false, + false, + ) + .unwrap(); + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + while let Some(batch) = stream.next().await { + batch.unwrap(); + } + } + + let cache = shared_metadata_cache(&ObjectStoreUrl::local_filesystem()) + .expect("metadata cache sharing is enabled by default"); + let entry = cache + .get(&location) + .expect("file metadata should be cached"); + assert_eq!( + entry.meta.last_modified.timestamp_millis(), + 2_000, + "the stale entry must have been replaced by a re-fetch" + ); + } } From 6f81d4f2c6f7e22e1864de2bd03d00943a29e32c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 09:42:02 -0600 Subject: [PATCH 06/10] feat: add configs for the shared Parquet metadata cache --- docs/source/user-guide/latest/tuning.md | 15 ++++++++ native/core/src/execution/jni_api.rs | 16 +++++++-- native/core/src/execution/spark_config.rs | 10 ++++++ .../scala/org/apache/comet/CometConf.scala | 23 ++++++++++++ .../org/apache/comet/CometExecIterator.scala | 8 +++++ .../comet/CometMetadataCacheSuite.scala | 36 +++++++++++++++++++ 6 files changed, 106 insertions(+), 2 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 6f0e9a684b..f2c672f490 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -103,6 +103,21 @@ 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 the cache in each executor, +with least-recently-used eviction. Note that this budget is per executor rather than per task, +so it does not scale with the number of task slots. Entries are validated against each file's +size and modification time, so an overwritten file is 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 diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 6837373500..bd0b36406f 100644 --- a/native/core/src/execution/jni_api.rs +++ b/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)); diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5d..ac28a424ae 100644 --- a/native/core/src/execution/spark_config.rs +++ b/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 { .unwrap_or(false) } + fn get_bool_or(&self, name: &str, default_value: bool) -> bool { + self.get(name) + .and_then(|str_val| str_val.parse::().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::().ok()) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 41c99d2724..63eb409ea1 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -132,6 +132,29 @@ object CometConf extends ShimCometConf { .checkValue(v => v > 0, "Data file concurrency limit must be positive") .createWithDefault(1) + // Used on native side. Check spark_config.rs how the config is used + val COMET_METADATA_CACHE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.metadataCache.enabled") + .category(CATEGORY_SCAN) + .doc( + "When enabled, the native Parquet scan shares its cache of file metadata (footer and " + + "page index) across all tasks in an executor, so tasks reading splits of the same " + + "file parse the footer once. When disabled, each task keeps its own cache, which is " + + "discarded when the task ends.") + .booleanConf + .createWithDefault(true) + + // Used on native side. Check spark_config.rs how the config is used + val COMET_METADATA_CACHE_MEMORY_LIMIT: ConfigEntry[Long] = + conf("spark.comet.scan.metadataCache.memoryLimit") + .category(CATEGORY_SCAN) + .doc( + "The maximum amount of memory (in bytes) used by the shared Parquet metadata cache in " + + "each executor. Entries are evicted least-recently-used. This budget is per executor, " + + "not per task.") + .bytesConf(ByteUnit.BYTE) + .createWithDefault(50L * 1024 * 1024) + val COMET_CSV_V2_NATIVE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.scan.csv.v2.enabled") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e59dec7ebf..8885d8163a 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -278,6 +278,14 @@ object CometExecIterator extends Logging { CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) + builder.putEntries( + CometConf.COMET_METADATA_CACHE_ENABLED.key, + CometConf.COMET_METADATA_CACHE_ENABLED.get(SQLConf.get).toString) + + builder.putEntries( + CometConf.COMET_METADATA_CACHE_MEMORY_LIMIT.key, + CometConf.COMET_METADATA_CACHE_MEMORY_LIMIT.get(SQLConf.get).toString) + builder.build().toByteArray } diff --git a/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala new file mode 100644 index 0000000000..ebcbc25496 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala @@ -0,0 +1,36 @@ +/* + * 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. + */ + +package org.apache.comet + +import org.apache.spark.sql.CometTestBase + +import org.apache.comet.serde.Config.ConfigMap + +class CometMetadataCacheSuite extends CometTestBase { + + test("metadata cache configs cross JNI with their default values") { + val confs = ConfigMap + .parseFrom(CometExecIterator.serializeCometSQLConfs()) + .getEntriesMap + + assert(confs.get(CometConf.COMET_METADATA_CACHE_ENABLED.key) == "true") + assert(confs.get(CometConf.COMET_METADATA_CACHE_MEMORY_LIMIT.key) == "52428800") + } +} From 7dec06a146748aacfcb6368e5f431cccd2ca64dc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 10:17:50 -0600 Subject: [PATCH 07/10] fix: correct metadata cache docs and add missing test coverage Fix code review findings on the shared Parquet metadata cache: correct the false claim that keying by object store URL removes all bucket collisions (ABFS containers on one storage account still share a key, since the container is URL userinfo, not host), correct the memory limit doc to state it bounds each object store's cache rather than one process-wide budget, note that the limit only applies when sharing is enabled, note that the enabled flag is process-global and last-writer-wins, note that staleness validation degrades to size-only without a reported modification time, and document the epoch-0 last_modified risk in the dormant single-file JNI path. Add tests covering a read-after-overwrite in the same session and the per-task fallback scan path when cache sharing is disabled. --- docs/source/user-guide/latest/tuning.md | 11 ++++--- native/core/src/parquet/metadata_cache.rs | 11 +++++-- native/core/src/parquet/mod.rs | 9 +++++ .../scala/org/apache/comet/CometConf.scala | 14 +++++--- .../comet/CometMetadataCacheSuite.scala | 33 ++++++++++++++++++- 5 files changed, 65 insertions(+), 13 deletions(-) diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index f2c672f490..0ce8364136 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -110,11 +110,12 @@ different splits of the same file, or later stages re-reading the same table, do 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 the cache in each executor, -with least-recently-used eviction. Note that this budget is per executor rather than per task, -so it does not scale with the number of task slots. Entries are validated against each file's -size and modification time, so an overwritten file is re-read rather than served from the -cache. +`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. diff --git a/native/core/src/parquet/metadata_cache.rs b/native/core/src/parquet/metadata_cache.rs index a2e546c189..1e335d478f 100644 --- a/native/core/src/parquet/metadata_cache.rs +++ b/native/core/src/parquet/metadata_cache.rs @@ -38,9 +38,14 @@ use std::sync::{Arc, OnceLock, PoisonError, RwLock}; /// 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` removes that entirely. 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. +/// 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 /// diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index dc8c4110ad..f82d278dd1 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -85,6 +85,15 @@ fn get_batch_context<'a>(handle: jlong) -> Result<&'a mut BatchContext, CometErr } } +/// Builds `PartitionedFile`s for the `Native.initRecordBatchReader` JNI entry point. There is +/// currently no `org.apache.comet.parquet.Native` Java class in the repo, so this path is +/// unreachable, but it is still exported. The `PartitionedFile`s built here never have +/// `object_meta.last_modified` set, so it defaults to the Unix epoch, and the `SessionContext` +/// built in `Java_org_apache_comet_parquet_Native_initRecordBatchReader` never calls +/// `configure_metadata_cache`. If this entry point is revived and its scan reaches the shared +/// metadata cache registry through `init_datasource_exec`, every file will validate as though +/// unmodified since the epoch, so cache validation degrades to size alone. Plumb the real file +/// modification time through here before relying on that path again. fn get_file_groups_single_file( path: &Path, file_size: u64, diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 63eb409ea1..d59948de81 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -140,7 +140,9 @@ object CometConf extends ShimCometConf { "When enabled, the native Parquet scan shares its cache of file metadata (footer and " + "page index) across all tasks in an executor, so tasks reading splits of the same " + "file parse the footer once. When disabled, each task keeps its own cache, which is " + - "discarded when the task ends.") + "discarded when the task ends. This setting is process-global: it takes effect for " + + "the whole executor process, and if multiple Spark sessions share an executor, the " + + "value from whichever task most recently planned applies.") .booleanConf .createWithDefault(true) @@ -149,9 +151,13 @@ object CometConf extends ShimCometConf { conf("spark.comet.scan.metadataCache.memoryLimit") .category(CATEGORY_SCAN) .doc( - "The maximum amount of memory (in bytes) used by the shared Parquet metadata cache in " + - "each executor. Entries are evicted least-recently-used. This budget is per executor, " + - "not per task.") + "The maximum amount of memory (in bytes) used by the shared Parquet metadata cache " + + "for each object store within an executor. Entries are evicted least-recently-used. " + + "An executor reading from multiple object stores, or multiple buckets, holds a " + + "multiple of this limit; it does not scale with the number of task slots. This " + + "limit only bounds the shared cache: when " + + "spark.comet.scan.metadataCache.enabled=false, each task's own cache uses " + + "DataFusion's default limit instead.") .bytesConf(ByteUnit.BYTE) .createWithDefault(50L * 1024 * 1024) diff --git a/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala index ebcbc25496..42843ec157 100644 --- a/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala @@ -19,11 +19,12 @@ package org.apache.comet -import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.{CometTestBase, Row} import org.apache.comet.serde.Config.ConfigMap class CometMetadataCacheSuite extends CometTestBase { + import testImplicits._ test("metadata cache configs cross JNI with their default values") { val confs = ConfigMap @@ -33,4 +34,34 @@ class CometMetadataCacheSuite extends CometTestBase { assert(confs.get(CometConf.COMET_METADATA_CACHE_ENABLED.key) == "true") assert(confs.get(CometConf.COMET_METADATA_CACHE_MEMORY_LIMIT.key) == "52428800") } + + // This exercises the shared-cache scan path end to end across an overwrite of the same + // path within a single session: it fails if a stale cached footer or page index makes the + // second read return the first read's data (or blow up on the first read's row group + // layout). It does not, on its own, prove that (size, last_modified) validation is what + // catches the change, since the row count change below also changes the file's size, and a + // size change alone would invalidate a stale entry too. + test("reading a table after it is overwritten returns the new data") { + withTempPath { dir => + val path = dir.getCanonicalPath + + Seq(1, 2, 3).toDF("id").write.parquet(path) + checkAnswer(spark.read.parquet(path), Seq(Row(1), Row(2), Row(3))) + + Seq(4, 5).toDF("id").write.mode("overwrite").parquet(path) + checkAnswer(spark.read.parquet(path), Seq(Row(4), Row(5))) + } + } + + // This is the only coverage of the per-task cache fallback in parquet_exec.rs, taken when + // sharing is disabled. + test("scan returns correct results with metadata cache sharing disabled") { + withSQLConf(CometConf.COMET_METADATA_CACHE_ENABLED.key -> "false") { + withTempPath { dir => + val path = dir.getCanonicalPath + Seq((1, "a"), (2, "b"), (3, "c")).toDF("id", "value").write.parquet(path) + checkSparkAnswer(spark.read.parquet(path)) + } + } + } } From cf660a2344b30dde496fe2ed2aa50d1c32245133 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 10:20:04 -0600 Subject: [PATCH 08/10] fix: make the overwrite metadata cache test exercise a cache hit The overwrite test previously asserted results only around a overwrite, but write.mode("overwrite") produces new file names, so those reads never hit the same cache entry and the test exercised nothing the shared cache is involved in. Read the same files twice before the overwrite to cover an actual shared-cache hit, and correct the comment to describe what the overwrite step does and does not prove. Use checkSparkAnswer so each read is also checked against vanilla Spark, consistent with the rest of the suite. --- .../comet/CometMetadataCacheSuite.scala | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala index 42843ec157..2a2f4e762e 100644 --- a/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala @@ -19,7 +19,7 @@ package org.apache.comet -import org.apache.spark.sql.{CometTestBase, Row} +import org.apache.spark.sql.CometTestBase import org.apache.comet.serde.Config.ConfigMap @@ -35,21 +35,22 @@ class CometMetadataCacheSuite extends CometTestBase { assert(confs.get(CometConf.COMET_METADATA_CACHE_MEMORY_LIMIT.key) == "52428800") } - // This exercises the shared-cache scan path end to end across an overwrite of the same - // path within a single session: it fails if a stale cached footer or page index makes the - // second read return the first read's data (or blow up on the first read's row group - // layout). It does not, on its own, prove that (size, last_modified) validation is what - // catches the change, since the row count change below also changes the file's size, and a - // size change alone would invalidate a stale entry too. + // The repeated read of the same files is the part that matters here: it is the only + // JVM-level coverage of a shared-cache hit across queries in one process, since the + // registry outlives the task that populated it. The subsequent overwrite-and-read-again + // step is a cheap guard on the whole scan path, but it does not test (size, last_modified) + // invalidation: `write.mode("overwrite")` writes new file names, so the post-overwrite files + // are distinct cache entries, not a stale hit on the pre-overwrite ones. test("reading a table after it is overwritten returns the new data") { withTempPath { dir => val path = dir.getCanonicalPath Seq(1, 2, 3).toDF("id").write.parquet(path) - checkAnswer(spark.read.parquet(path), Seq(Row(1), Row(2), Row(3))) + checkSparkAnswer(spark.read.parquet(path)) + checkSparkAnswer(spark.read.parquet(path)) Seq(4, 5).toDF("id").write.mode("overwrite").parquet(path) - checkAnswer(spark.read.parquet(path), Seq(Row(4), Row(5))) + checkSparkAnswer(spark.read.parquet(path)) } } From 11411a701c236f6037b8dae4f2cc62920dd600dd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 10:42:11 -0600 Subject: [PATCH 09/10] test: cover same-size in-place file replacement against a stale footer --- .../comet/CometMetadataCacheSuite.scala | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala index 2a2f4e762e..096936bff9 100644 --- a/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala @@ -19,6 +19,9 @@ package org.apache.comet +import java.io.File +import java.nio.file.{Files, StandardCopyOption} + import org.apache.spark.sql.CometTestBase import org.apache.comet.serde.Config.ConfigMap @@ -54,6 +57,78 @@ class CometMetadataCacheSuite extends CometTestBase { } } + // Writes `numRows` rows whose pad column is `padWidth` characters, with dictionary encoding + // off so that file size responds to the pad width. `extra` widens only the first row, which + // gives finer control over the resulting file size than `padWidth` alone. + private def writePaddedFile(dir: File, numRows: Int, padWidth: Int, extra: Int = 0): File = { + (1 to numRows) + .map(i => (i, "p" * (padWidth + (if (i == 1) extra else 0)))) + .toDF("id", "pad") + .coalesce(1) + .write + .option("parquet.enable.dictionary", "false") + .parquet(dir.getCanonicalPath) + dir.listFiles().filter(_.getName.endsWith(".parquet")).head + } + + // The shared cache outlives the task that populated it, so a file replaced in place must not + // be read through the previous file's footer. Size alone does not catch this: the two files + // here are byte-identical in length but hold different row counts, so the stale footer's row + // group offsets do not describe the new file. Without the modification time reaching + // `object_meta.last_modified`, this read fails outright with FAILED_READ_FILE. + test("a same-size in-place replacement is not read through the stale footer") { + withTempPath { root => + root.mkdirs() + val target = new File(root, "target") + target.mkdirs() + + val fileA = writePaddedFile(new File(root, "a"), numRows = 50, padWidth = 40) + val sizeA = fileA.length() + + // Search for a 30-row file of exactly sizeA bytes. Encoding overhead means the size + // change is not exactly the padding added, so iterate to a fixed point. + var matched: Option[File] = None + var padWidth = 20 + while (matched.isEmpty && padWidth < 60) { + var extra = + (sizeA - writePaddedFile(new File(root, s"b_${padWidth}_0"), 30, padWidth) + .length()).toInt + var attempt = 0 + while (matched.isEmpty && attempt < 12 && extra > 0 && extra < 6000) { + val candidate = + writePaddedFile( + new File(root, s"b_${padWidth}_${attempt}_$extra"), + 30, + padWidth, + extra) + val delta = (sizeA - candidate.length()).toInt + if (delta == 0) { + matched = Some(candidate) + } else { + extra += delta + } + attempt += 1 + } + padWidth += 1 + } + + val fileB = matched.getOrElse( + fail(s"could not construct a 30-row Parquet file of exactly $sizeA bytes")) + + val live = new File(target, "part-00000.parquet") + Files.copy(fileA.toPath, live.toPath, StandardCopyOption.REPLACE_EXISTING) + assert(spark.read.parquet(target.getCanonicalPath).count() == 50) + + Files.copy(fileB.toPath, live.toPath, StandardCopyOption.REPLACE_EXISTING) + assert(live.length() == sizeA, "the replacement must not change the file size") + + // count(*) is answered from footer row counts, so a stale entry shows up here even + // though no data page is touched. + assert(spark.read.parquet(target.getCanonicalPath).count() == 30) + assert(spark.read.parquet(target.getCanonicalPath).collect().length == 30) + } + } + // This is the only coverage of the per-task cache fallback in parquet_exec.rs, taken when // sharing is disabled. test("scan returns correct results with metadata cache sharing disabled") { From 547c6ec3a839e4d39a37b3d7b78e2026a8256461 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 10:54:42 -0600 Subject: [PATCH 10/10] ci: register the metadata cache and file partition serde suites --- .github/workflows/pr_build_linux.yml | 2 ++ .github/workflows/pr_build_macos.yml | 2 ++ 2 files changed, 4 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 6aaac53204..d5f64b9e08 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.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 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 36ed5f9405..9cea6c6465 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.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