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 diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 6f0e9a684b..0ce8364136 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/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 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/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/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/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/native/core/src/parquet/metadata_cache.rs b/native/core/src/parquet/metadata_cache.rs new file mode 100644 index 0000000000..1e335d478f --- /dev/null +++ b/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, + 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. + /// + /// 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> { + 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); + } + + #[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)); + } + } +} diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index 5644e61c85..f82d278dd1 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; @@ -84,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/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" + ); + } } 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/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 41c99d2724..d59948de81 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -132,6 +132,35 @@ 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. 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) + + // 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 " + + "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) + 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/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/CometMetadataCacheSuite.scala b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala new file mode 100644 index 0000000000..096936bff9 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometMetadataCacheSuite.scala @@ -0,0 +1,143 @@ +/* + * 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 java.io.File +import java.nio.file.{Files, StandardCopyOption} + +import org.apache.spark.sql.CometTestBase + +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 + .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") + } + + // 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) + checkSparkAnswer(spark.read.parquet(path)) + checkSparkAnswer(spark.read.parquet(path)) + + Seq(4, 5).toDF("id").write.mode("overwrite").parquet(path) + checkSparkAnswer(spark.read.parquet(path)) + } + } + + // 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") { + 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)) + } + } + } +} 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) + } +}