From 07327e7c25745e62d1e8536afd943174ad18998c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 15:46:27 -0600 Subject: [PATCH 1/9] perf: share Parquet footer and page index metadata across the tasks of a stage Closes #4994 for the case of a file split across the tasks of one stage. --- native/core/src/execution/jni_api.rs | 33 +++- native/core/src/execution/spark_config.rs | 5 + native/core/src/parquet/mod.rs | 1 + .../core/src/parquet/shared_metadata_cache.rs | 164 ++++++++++++++++++ .../scala/org/apache/comet/CometConf.scala | 13 ++ .../org/apache/comet/CometExecIterator.scala | 10 ++ .../apache/comet/exec/CometExecSuite.scala | 70 ++++++++ 7 files changed, 294 insertions(+), 2 deletions(-) create mode 100644 native/core/src/parquet/shared_metadata_cache.rs diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index cabbb58d2bc..0cf7cafe758 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -44,6 +44,9 @@ use arrow::array::{Array, RecordBatch}; use arrow::datatypes::DataType as ArrowDataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::config::SpillCompression; +use datafusion::execution::cache::cache_manager::{ + CacheManagerConfig, DEFAULT_METADATA_CACHE_LIMIT, +}; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::MemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -135,11 +138,13 @@ use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, - COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS, - COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED, + COMET_STAGE_ATTEMPT, COMET_STAGE_ID, COMET_STRICT_TIMESTAMP_LABELS, COMET_TRACING_ENABLED, + SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_support::CometObjectStoreRegistry; +use crate::parquet::shared_metadata_cache::{shared_metadata_caches, StageKey}; use datafusion_comet_proto::spark_operator::operator::OpStruct; use log::{info, warn}; use std::sync::OnceLock; @@ -1009,6 +1014,30 @@ fn prepare_datafusion_session_context( .with_object_store_registry(Arc::new(CometObjectStoreRegistry::default())); rt_config = rt_config.with_memory_pool(memory_pool); + // Let the tasks of one stage on this executor share the Parquet footer and page index cache + // instead of each fetching and parsing the footers of the files it reads. + if spark_config.get_bool(COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED) { + if let (Some(stage_id), Some(stage_attempt)) = ( + spark_config + .get(COMET_STAGE_ID) + .and_then(|v| v.parse::().ok()), + spark_config + .get(COMET_STAGE_ATTEMPT) + .and_then(|v| v.parse::().ok()), + ) { + let cache = shared_metadata_caches().cache_for( + StageKey { + stage_id, + stage_attempt, + }, + DEFAULT_METADATA_CACHE_LIMIT, + ); + rt_config = rt_config.with_cache_manager( + CacheManagerConfig::default().with_file_metadata_cache(Some(cache)), + ); + } + } + let mut session_config = SessionConfig::new() .with_target_partitions(task_cpus) // This DataFusion context is within the scope of an executing Spark Task. We want to set diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 7755f6d5dc4..06207f3667d 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -30,6 +30,11 @@ pub(crate) const COMET_EXEC_SPILL_COMPRESSION_CODEC: &str = "spark.comet.exec.spill.compression.codec"; pub(crate) const COMET_STRICT_TIMESTAMP_LABELS: &str = "spark.comet.testing.strictTimestampLabels.enabled"; +pub(crate) const COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED: &str = + "spark.comet.parquet.metadataCache.sharedWithinStage.enabled"; +/// Not a user config: `CometExecIterator.serializeCometSQLConfs` injects the stage of the task. +pub(crate) const COMET_STAGE_ID: &str = "spark.comet.internal.stageId"; +pub(crate) const COMET_STAGE_ATTEMPT: &str = "spark.comet.internal.stageAttemptNumber"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; /// Comet configs read through this trait must be resolved by the JVM first: diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index 9c63b34bff2..4cbccea83e7 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -27,3 +27,4 @@ mod cast_column; mod datetime_rebase; mod name_fold; pub(crate) mod objectstore; +pub(crate) mod shared_metadata_cache; diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs new file mode 100644 index 00000000000..6d445089674 --- /dev/null +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -0,0 +1,164 @@ +// 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. + +//! A registry that lets the tasks of one Spark stage share a Parquet `FileMetadataCache`. +//! +//! Comet builds a `RuntimeEnv`, and with it a metadata cache, for every Spark task, so two tasks +//! that read splits of the same file each fetch and parse its footer and page index. This +//! registry hands every task of a stage on this executor the same cache instead. +//! +//! The scope is the stage, not the executor, on purpose. A cache that outlives the query would +//! answer a later query from metadata of a file that has since been deleted or rewritten, and +//! `(size, modification time)` validation cannot catch that: Spark's cached file index lists the +//! same size and time as the cached footer. Spark's `HiveMetadataCacheSuite` reads a table, +//! deletes one data file, reads again, and expects `FILE_NOT_EXIST`. Every new query is a new +//! stage, so it opens its files and sees the deletion, as it does without the registry. +//! +//! A stage attempt is part of the key because a retried stage may read a changed file. +//! +//! The registry keeps the cache of the stages that still have a running task, plus the last +//! [`MAX_IDLE_CACHES`] that do not. The idle ones matter when an executor runs one task at a +//! time: the next task of the stage starts after the previous one released its plan. + +use datafusion::execution::cache::cache_manager::FileMetadataCache; +use datafusion::execution::cache::default_cache::DefaultCache; +use std::sync::{Arc, Mutex, OnceLock}; + +/// How many caches of stages with no running task the registry holds on to. +const MAX_IDLE_CACHES: usize = 2; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct StageKey { + pub stage_id: i32, + pub stage_attempt: i32, +} + +struct Entry { + key: StageKey, + cache: Arc, +} + +#[derive(Default)] +pub(crate) struct SharedMetadataCaches { + /// Ordered from least to most recently requested. + entries: Mutex>, +} + +impl SharedMetadataCaches { + /// The cache shared by the tasks of `key`, created with `limit` bytes of budget if no task of + /// that stage has asked for one yet. + pub(crate) fn cache_for(&self, key: StageKey, limit: usize) -> Arc { + let mut entries = self.entries.lock().unwrap(); + let entry = match entries.iter().position(|e| e.key == key) { + Some(pos) => entries.remove(pos), + None => Entry { + key, + cache: Arc::new(DefaultCache::new(limit).with_name("SharedFileMetadataCache")), + }, + }; + let cache = Arc::clone(&entry.cache); + entries.push(entry); + + // The registry holds one reference; a task holds another for as long as its + // `RuntimeEnv` lives. The newest entry also has the reference returned from here, which + // the task takes over, so it counts as idle at two. The newest entry is never evicted. + let newest = entries.len() - 1; + let is_idle = |index: usize, entry: &Entry| { + Arc::strong_count(&entry.cache) == if index == newest { 2 } else { 1 } + }; + let mut idle = entries + .iter() + .enumerate() + .filter(|(i, e)| is_idle(*i, e)) + .count(); + let mut index = 0; + entries.retain(|e| { + let evict = index != newest && idle > MAX_IDLE_CACHES && is_idle(index, e); + index += 1; + if evict { + idle -= 1; + } + !evict + }); + cache + } + + #[cfg(test)] + fn len(&self) -> usize { + self.entries.lock().unwrap().len() + } +} + +/// The registry of this process. +pub(crate) fn shared_metadata_caches() -> &'static SharedMetadataCaches { + static CACHES: OnceLock = OnceLock::new(); + CACHES.get_or_init(SharedMetadataCaches::default) +} + +#[cfg(test)] +mod tests { + use super::*; + + const LIMIT: usize = 1024 * 1024; + + fn key(stage_id: i32, stage_attempt: i32) -> StageKey { + StageKey { + stage_id, + stage_attempt, + } + } + + #[test] + fn tasks_of_one_stage_share_a_cache() { + let caches = SharedMetadataCaches::default(); + let a = caches.cache_for(key(1, 0), LIMIT); + let b = caches.cache_for(key(1, 0), LIMIT); + assert!(Arc::ptr_eq(&a, &b)); + } + + #[test] + fn stages_and_attempts_do_not_share() { + let caches = SharedMetadataCaches::default(); + let first = caches.cache_for(key(1, 0), LIMIT); + let other_stage = caches.cache_for(key(2, 0), LIMIT); + let retry = caches.cache_for(key(1, 1), LIMIT); + assert!(!Arc::ptr_eq(&first, &other_stage)); + assert!(!Arc::ptr_eq(&first, &retry)); + } + + #[test] + fn an_idle_cache_is_kept_for_the_next_task_of_its_stage() { + let caches = SharedMetadataCaches::default(); + let first = caches.cache_for(key(1, 0), LIMIT); + let address = Arc::as_ptr(&first) as *const () as usize; + drop(first); + let next = caches.cache_for(key(1, 0), LIMIT); + assert_eq!(Arc::as_ptr(&next) as *const () as usize, address); + } + + #[test] + fn idle_caches_are_bounded_but_running_ones_are_not_dropped() { + let caches = SharedMetadataCaches::default(); + let running = caches.cache_for(key(0, 0), LIMIT); + for stage in 1..=10 { + drop(caches.cache_for(key(stage, 0), LIMIT)); + } + // The running stage plus the idle ones that fit. + assert_eq!(caches.len(), 1 + MAX_IDLE_CACHES); + assert!(Arc::ptr_eq(&running, &caches.cache_for(key(0, 0), LIMIT))); + } +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 4cf228e47ee..2371ed59059 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -184,6 +184,19 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.parquet.metadataCache.sharedWithinStage.enabled") + .category(CATEGORY_PARQUET) + .doc( + "When enabled, the tasks of one Spark stage that run on the same executor share the " + + "native Parquet scan's cache of file footers and page indexes, so a file split " + + "across many tasks has its footer fetched and parsed once per executor instead of " + + "once per task. The cache is never shared between stages or stage attempts, so a " + + "query always reads the footers of its files afresh. When disabled, each task " + + "keeps its own cache.") + .booleanConf + .createWithDefault(true) + val COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.parquet.rowFilterPushdown.enabled") .category(CATEGORY_PARQUET) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 8997be2d7e9..38bcf6413da 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -877,6 +877,15 @@ object CometExecIterator extends Logging { val executorCores = numDriverOrExecutorCores(SparkEnv.get.conf) builder.putEntries("spark.executor.cores", executorCores.toString) + // The stage of this task, which scopes the Parquet metadata cache the native scan shares + // between the tasks of one stage. + Option(TaskContext.get()).foreach { context => + builder.putEntries("spark.comet.internal.stageId", context.stageId().toString) + builder.putEntries( + "spark.comet.internal.stageAttemptNumber", + context.stageAttemptNumber().toString) + } + // Any Comet config that the native side reads must be added here manually, resolved. // `cometSqlConfs` only carries values that were explicitly set, exactly as they were // written, so defaults from `createWithDefault(...)` would otherwise not cross JNI, and @@ -890,6 +899,7 @@ object CometExecIterator extends Logging { CometConf.COMET_EXPLAIN_NATIVE_ENABLED, CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, + CometConf.COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED, CometConf.COMET_STRICT_TIMESTAMP_LABELS, CometConf.COMET_TRACING_ENABLED).foreach { entry => builder.putEntries(entry.key, entry.get(SQLConf.get).toString) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 8ec76483027..d22ae89de01 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -29,6 +29,7 @@ import org.scalactic.source.Position import org.scalatest.Tag import org.apache.hadoop.fs.Path +import org.apache.spark.SparkException import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier} import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStatistics, CatalogTable} @@ -3360,6 +3361,75 @@ class CometExecSuite extends CometTestBase { } } + /** The footer metadata cache hits and misses of the native scan that ran `df`. */ + private def footerCacheCounts(df: DataFrame): (Long, Long) = { + val scan = find(df.queryExecution.executedPlan)(s => + s.isInstanceOf[CometScanExec] || s.isInstanceOf[CometNativeScanExec]) + assert(scan.isDefined, "Expected to find a Comet scan node") + val metrics = scan.get.metrics + (metrics("scan_io_metadata_cache_hits").value, metrics("scan_io_metadata_cache_misses").value) + } + + test("Parquet footers are shared by the tasks of a stage") { + withTempDir { dir => + val path = new Path(dir.toURI.toString, "one-file").toString + spark + .range(0, 200000) + .selectExpr("id", "cast(id * 7 as string) as s") + .coalesce(1) + .write + .parquet(path) + + // One file split into many tasks, so every task needs the same footer. + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.FILES_MAX_PARTITION_BYTES.key -> "64k", + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") { + val query = s"SELECT sum(id), count(s) FROM parquet.`$path`" + val tasks = spark.read.parquet(path).rdd.getNumPartitions + assert(tasks > 8, s"Expected the file to be split into many tasks, got $tasks") + + withSQLConf(CometConf.COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED.key -> "true") { + val df = sql(query) + df.collect() + val (hits, misses) = footerCacheCounts(df) + // Tasks that start together can each miss, but not all of them do. + assert(hits > 0, s"hits=$hits misses=$misses") + assert(misses < tasks, s"hits=$hits misses=$misses tasks=$tasks") + } + + withSQLConf(CometConf.COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED.key -> "false") { + val df = sql(query) + df.collect() + val (hits, misses) = footerCacheCounts(df) + assert(hits == 0, s"hits=$hits misses=$misses") + assert(misses >= tasks, s"hits=$hits misses=$misses tasks=$tasks") + } + } + } + } + + test("Parquet footers are not shared between queries") { + withTempDir { dir => + val path = new Path(dir.toURI.toString, "table").toString + spark.range(0, 1000).repartition(4).write.parquet(path) + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED.key -> "true") { + // The relation keeps its file listing, as a cached table does, so the second query is + // given a file that no longer exists. It must find that out by opening the file. + val table = spark.read.parquet(path) + assert(table.count() == 1000) + val deleted = new java.io.File(dir, "table") + .listFiles() + .filter(_.getName.endsWith(".parquet")) + .head + assert(deleted.delete()) + intercept[SparkException](table.count()) + } + } + } + test("Comet native metrics: project and filter") { withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "true") { withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl") { From 8f0a5cf6dc595f943a520f3f49c8d3969e4835ec Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 16:24:21 -0600 Subject: [PATCH 2/9] refactor: move shared metadata cache lookup into one helper --- native/core/src/execution/jni_api.rs | 35 +++-------- .../core/src/parquet/shared_metadata_cache.rs | 63 +++++++++++++------ 2 files changed, 51 insertions(+), 47 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 0cf7cafe758..ac269dbf9be 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -44,9 +44,7 @@ use arrow::array::{Array, RecordBatch}; use arrow::datatypes::DataType as ArrowDataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::config::SpillCompression; -use datafusion::execution::cache::cache_manager::{ - CacheManagerConfig, DEFAULT_METADATA_CACHE_LIMIT, -}; +use datafusion::execution::cache::cache_manager::CacheManagerConfig; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::MemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -138,13 +136,12 @@ use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, - COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED, - COMET_STAGE_ATTEMPT, COMET_STAGE_ID, COMET_STRICT_TIMESTAMP_LABELS, COMET_TRACING_ENABLED, - SPARK_EXECUTOR_CORES, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS, + COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_support::CometObjectStoreRegistry; -use crate::parquet::shared_metadata_cache::{shared_metadata_caches, StageKey}; +use crate::parquet::shared_metadata_cache::shared_metadata_cache_for_task; use datafusion_comet_proto::spark_operator::operator::OpStruct; use log::{info, warn}; use std::sync::OnceLock; @@ -1016,26 +1013,10 @@ fn prepare_datafusion_session_context( // Let the tasks of one stage on this executor share the Parquet footer and page index cache // instead of each fetching and parsing the footers of the files it reads. - if spark_config.get_bool(COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED) { - if let (Some(stage_id), Some(stage_attempt)) = ( - spark_config - .get(COMET_STAGE_ID) - .and_then(|v| v.parse::().ok()), - spark_config - .get(COMET_STAGE_ATTEMPT) - .and_then(|v| v.parse::().ok()), - ) { - let cache = shared_metadata_caches().cache_for( - StageKey { - stage_id, - stage_attempt, - }, - DEFAULT_METADATA_CACHE_LIMIT, - ); - rt_config = rt_config.with_cache_manager( - CacheManagerConfig::default().with_file_metadata_cache(Some(cache)), - ); - } + if let Some(cache) = shared_metadata_cache_for_task(spark_config) { + rt_config = rt_config.with_cache_manager( + CacheManagerConfig::default().with_file_metadata_cache(Some(cache)), + ); } let mut session_config = SessionConfig::new() diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs index 6d445089674..3066ad8f5f3 100644 --- a/native/core/src/parquet/shared_metadata_cache.rs +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -34,8 +34,14 @@ //! [`MAX_IDLE_CACHES`] that do not. The idle ones matter when an executor runs one task at a //! time: the next task of the stage starts after the previous one released its plan. -use datafusion::execution::cache::cache_manager::FileMetadataCache; +use crate::execution::spark_config::{ + SparkConfig, COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED, COMET_STAGE_ATTEMPT, COMET_STAGE_ID, +}; +use datafusion::execution::cache::cache_manager::{ + FileMetadataCache, DEFAULT_METADATA_CACHE_LIMIT, +}; use datafusion::execution::cache::default_cache::DefaultCache; +use std::collections::HashMap; use std::sync::{Arc, Mutex, OnceLock}; /// How many caches of stages with no running task the registry holds on to. @@ -59,15 +65,18 @@ pub(crate) struct SharedMetadataCaches { } impl SharedMetadataCaches { - /// The cache shared by the tasks of `key`, created with `limit` bytes of budget if no task of - /// that stage has asked for one yet. - pub(crate) fn cache_for(&self, key: StageKey, limit: usize) -> Arc { + /// The cache shared by the tasks of `key`, created if no task of that stage has asked for one + /// yet. + fn cache_for(&self, key: StageKey) -> Arc { let mut entries = self.entries.lock().unwrap(); let entry = match entries.iter().position(|e| e.key == key) { Some(pos) => entries.remove(pos), None => Entry { key, - cache: Arc::new(DefaultCache::new(limit).with_name("SharedFileMetadataCache")), + cache: Arc::new( + DefaultCache::new(DEFAULT_METADATA_CACHE_LIMIT) + .with_name("SharedFileMetadataCache"), + ), }, }; let cache = Arc::clone(&entry.cache); @@ -103,18 +112,32 @@ impl SharedMetadataCaches { } } -/// The registry of this process. -pub(crate) fn shared_metadata_caches() -> &'static SharedMetadataCaches { +/// The cache the task with these Spark configs should use, or `None` when sharing is off or the +/// task's stage is unknown, in which case DataFusion's per-task cache applies. +pub(crate) fn shared_metadata_cache_for_task( + spark_config: &HashMap, +) -> Option> { static CACHES: OnceLock = OnceLock::new(); - CACHES.get_or_init(SharedMetadataCaches::default) + + if !spark_config.get_bool(COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED) { + return None; + } + let parse = |name: &str| spark_config.get(name)?.parse::().ok(); + let key = StageKey { + stage_id: parse(COMET_STAGE_ID)?, + stage_attempt: parse(COMET_STAGE_ATTEMPT)?, + }; + Some( + CACHES + .get_or_init(SharedMetadataCaches::default) + .cache_for(key), + ) } #[cfg(test)] mod tests { use super::*; - const LIMIT: usize = 1024 * 1024; - fn key(stage_id: i32, stage_attempt: i32) -> StageKey { StageKey { stage_id, @@ -125,17 +148,17 @@ mod tests { #[test] fn tasks_of_one_stage_share_a_cache() { let caches = SharedMetadataCaches::default(); - let a = caches.cache_for(key(1, 0), LIMIT); - let b = caches.cache_for(key(1, 0), LIMIT); + let a = caches.cache_for(key(1, 0)); + let b = caches.cache_for(key(1, 0)); assert!(Arc::ptr_eq(&a, &b)); } #[test] fn stages_and_attempts_do_not_share() { let caches = SharedMetadataCaches::default(); - let first = caches.cache_for(key(1, 0), LIMIT); - let other_stage = caches.cache_for(key(2, 0), LIMIT); - let retry = caches.cache_for(key(1, 1), LIMIT); + let first = caches.cache_for(key(1, 0)); + let other_stage = caches.cache_for(key(2, 0)); + let retry = caches.cache_for(key(1, 1)); assert!(!Arc::ptr_eq(&first, &other_stage)); assert!(!Arc::ptr_eq(&first, &retry)); } @@ -143,22 +166,22 @@ mod tests { #[test] fn an_idle_cache_is_kept_for_the_next_task_of_its_stage() { let caches = SharedMetadataCaches::default(); - let first = caches.cache_for(key(1, 0), LIMIT); + let first = caches.cache_for(key(1, 0)); let address = Arc::as_ptr(&first) as *const () as usize; drop(first); - let next = caches.cache_for(key(1, 0), LIMIT); + let next = caches.cache_for(key(1, 0)); assert_eq!(Arc::as_ptr(&next) as *const () as usize, address); } #[test] fn idle_caches_are_bounded_but_running_ones_are_not_dropped() { let caches = SharedMetadataCaches::default(); - let running = caches.cache_for(key(0, 0), LIMIT); + let running = caches.cache_for(key(0, 0)); for stage in 1..=10 { - drop(caches.cache_for(key(stage, 0), LIMIT)); + drop(caches.cache_for(key(stage, 0))); } // The running stage plus the idle ones that fit. assert_eq!(caches.len(), 1 + MAX_IDLE_CACHES); - assert!(Arc::ptr_eq(&running, &caches.cache_for(key(0, 0), LIMIT))); + assert!(Arc::ptr_eq(&running, &caches.cache_for(key(0, 0)))); } } From e6faa9d3f1f5d38f2f78960913b2c3936051dfbe Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 18 Oct 2026 18:50:56 +0100 Subject: [PATCH 3/9] fix: scope the shared Parquet metadata cache by object store and application Key each scan's view of the stage cache by its registered object store URL, so equal paths in different buckets do not share a footer, including for the Delta deletion vector footer read. Include the Spark application id in the stage key, since stage ids restart in every SparkContext and the registry outlives one. Hold the registry's caches by weak reference plus a bounded queue of strong ones, so idle retention stays bounded after the last task releases its cache. --- native/core/src/execution/delta_dv.rs | 146 ++++++++--- .../src/execution/planner/delta_spark_scan.rs | 6 +- native/core/src/execution/spark_config.rs | 5 +- native/core/src/parquet/parquet_exec.rs | 146 ++++++++++- .../core/src/parquet/shared_metadata_cache.rs | 235 ++++++++++++++---- .../org/apache/comet/CometExecIterator.scala | 9 +- ...CometSharedMetadataCacheRestartSuite.scala | 110 ++++++++ 7 files changed, 558 insertions(+), 99 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometSharedMetadataCacheRestartSuite.scala diff --git a/native/core/src/execution/delta_dv.rs b/native/core/src/execution/delta_dv.rs index 982c4434207..5a8907514fa 100644 --- a/native/core/src/execution/delta_dv.rs +++ b/native/core/src/execution/delta_dv.rs @@ -33,8 +33,10 @@ use std::mem::size_of; use std::sync::Arc; use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; use datafusion::datasource::physical_plan::parquet::{ParquetAccessPlan, RowGroupAccess}; +use datafusion::execution::cache::cache_manager::FileMetadataCache; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion_datasource::FileRange; @@ -47,6 +49,7 @@ use roaring::{RoaringBitmap, RoaringTreemap}; use crate::execution::operators::ExecutionError; use crate::execution::operators::ExecutionError::GeneralError; +use crate::parquet::shared_metadata_cache::scope_to_store; use datafusion_comet_proto::spark_operator::DeltaSparkDvDescriptor; const NATIVE_MAGIC: i32 = 1681511376; @@ -588,13 +591,24 @@ const DV_FETCH_CONCURRENCY: usize = 8; /// dropped. When no file carries a DV, it is dropped when this call returns. pub async fn attach_access_plans( runtime_env: Arc, + object_store_url: &ObjectStoreUrl, files: Vec, ) -> Result, ExecutionError> { let partition_reservation = MemoryConsumer::new("DeltaDeletionVectorAccessPlan").register(&runtime_env.memory_pool); + // The footers read here are the ones the scan will look up, so they go through the scan's + // view of the metadata cache: `object_store_url` is the URL the scan is built with. + let metadata_cache = scope_to_store( + runtime_env.cache_manager.get_file_metadata_cache(), + object_store_url, + ); futures::stream::iter(files) .map(|scan_file| { - attach_access_plan(Arc::clone(&runtime_env), &partition_reservation, scan_file) + attach_access_plan( + Arc::clone(&metadata_cache), + &partition_reservation, + scan_file, + ) }) .buffered(DV_FETCH_CONCURRENCY) .try_collect() @@ -606,7 +620,7 @@ pub async fn attach_access_plans( /// call's one registered consumer, and a DV'd file takes an empty reservation /// from it rather than registering a consumer of its own. async fn attach_access_plan( - runtime_env: Arc, + metadata_cache: Arc, partition_reservation: &MemoryReservation, scan_file: DvScanFile, ) -> Result { @@ -680,7 +694,6 @@ async fn attach_access_plan( // size hint (mirroring EagerPageIndexReaderFactory): the one fetch here also serves the // subsequent data-file open, so DV files pay no extra footer round-trip. Keyed by // `file.object_meta`, the exact ObjectMeta the scan's reader factory will look up. - let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); let metadata = DFParquetMetadata::new(data_store.as_ref(), &file.object_meta) .with_file_metadata_cache(Some(metadata_cache)) .with_page_index_policy(Some(PageIndexPolicy::Optional)) @@ -1006,9 +1019,13 @@ mod tests { }) .collect(); - let out = attach_access_plans(Arc::clone(&runtime_env), files) - .await - .unwrap(); + let out = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + files, + ) + .await + .unwrap(); assert_eq!(out.len(), names.len()); for (file, name) in out.iter().zip(names) { @@ -1053,7 +1070,11 @@ mod tests { // Footer reads must go through the shared FileMetadataCache so the scan's // subsequent open of the same file is served from cache instead of paying a // second footer round-trip. Files without a DV read no footer at all. - let cache = runtime_env.cache_manager.get_file_metadata_cache(); + // The scan looks the footers up through its store's view of the cache, so this does too. + let cache = scope_to_store( + runtime_env.cache_manager.get_file_metadata_cache(), + &ObjectStoreUrl::local_filesystem(), + ); for (file, name) in out.iter().zip(names) { let cached = cache.get(&file.object_meta.location); match name { @@ -1448,9 +1469,13 @@ mod tests { start: offset, end: offset + 1, }); - let files = attach_access_plans(runtime_env, vec![full, split]) - .await - .unwrap(); + let files = attach_access_plans( + runtime_env, + &ObjectStoreUrl::local_filesystem(), + vec![full, split], + ) + .await + .unwrap(); let full_plan = files[0].extensions.get::().unwrap(); let split_plan = files[1].extensions.get::().unwrap(); assert_eq!(total_selectors(full_plan), 3_000); @@ -1515,9 +1540,13 @@ mod tests { start: offset, end: offset + 1, }); - let files = attach_access_plans(Arc::clone(&runtime_env), vec![file]) - .await - .unwrap(); + let files = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![file], + ) + .await + .unwrap(); let session = Arc::new(SessionContext::new_with_config_rt( SessionConfig::new(), runtime_env, @@ -1619,9 +1648,13 @@ mod tests { ); let scan_file = alternating_dv_scan_file(&runtime_env, tmp.path(), 1024); - let err = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) - .await - .unwrap_err(); + let err = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file], + ) + .await + .unwrap_err(); let msg = err.to_string(); assert!( msg.contains("alternating.parquet"), @@ -1663,9 +1696,13 @@ mod tests { // this file's retained selector count regardless of its real page-index bound. let expected_bytes = reader_peak_bytes(1024, 1).unwrap(); - let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) - .await - .unwrap(); + let out = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file], + ) + .await + .unwrap(); assert_eq!(out.len(), 1); assert_eq!( pool.reserved(), @@ -1709,9 +1746,13 @@ mod tests { let expected_bytes = reader_peak_bytes(1024, 1).unwrap() + reader_peak_bytes(512, 1).unwrap(); - let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file_a, scan_file_b]) - .await - .unwrap(); + let out = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file_a, scan_file_b], + ) + .await + .unwrap(); assert_eq!(out.len(), 2); assert_eq!( pool.reserved(), @@ -1867,6 +1908,7 @@ mod tests { let mut out = attach_access_plans( Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), vec![scan_file_a, scan_file_b, scan_file_c], ) .await @@ -1951,9 +1993,13 @@ mod tests { let scan_file_a = dv_scan_file_for_alternating(&runtime_env, &path_a, size_a, &deleted_a); let scan_file_b = alternating_dv_scan_file(&runtime_env, tmp_b.path(), 512); - let err = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file_a, scan_file_b]) - .await - .unwrap_err(); + let err = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file_a, scan_file_b], + ) + .await + .unwrap_err(); let msg = err.to_string(); assert!( msg.contains("Resources exhausted") || msg.contains("exceeding"), @@ -2014,9 +2060,13 @@ mod tests { ); let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); - let err = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) - .await - .unwrap_err(); + let err = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file], + ) + .await + .unwrap_err(); let msg = err.to_string(); assert!( msg.to_lowercase().contains("construct"), @@ -2048,9 +2098,13 @@ mod tests { let deleted = alternating_deleted(num_rows as u64); let scan_file = alternating_dv_scan_file(&runtime_env, tmp.path(), num_rows); - let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) - .await - .unwrap(); + let out = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file], + ) + .await + .unwrap(); let plan = full_file_access_plan(&[num_rows], &deleted).unwrap(); let expected_bytes = reader_peak_bytes(total_selectors(&plan), 1).unwrap(); @@ -2129,9 +2183,13 @@ mod tests { let (path, size, deleted) = write_alternating_parquet(tmp.path(), 1024); let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); - let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) - .await - .unwrap(); + let out = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file], + ) + .await + .unwrap(); let attached_plan = out[0] .extensions .get::() @@ -2196,9 +2254,13 @@ mod tests { let size = write_two_row_groups(&path, rows_per_group); let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); - let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) - .await - .unwrap(); + let out = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file], + ) + .await + .unwrap(); let attached_plan = out[0] .extensions .get::() @@ -2270,9 +2332,13 @@ mod tests { let (path, size, deleted) = write_alternating_parquet(tmp.path(), num_rows); let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); - let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) - .await - .unwrap(); + let out = attach_access_plans( + Arc::clone(&runtime_env), + &ObjectStoreUrl::local_filesystem(), + vec![scan_file], + ) + .await + .unwrap(); let attached_plan = out[0] .extensions .get::() diff --git a/native/core/src/execution/planner/delta_spark_scan.rs b/native/core/src/execution/planner/delta_spark_scan.rs index a12effbea42..f6c35edf3f8 100644 --- a/native/core/src/execution/planner/delta_spark_scan.rs +++ b/native/core/src/execution/planner/delta_spark_scan.rs @@ -190,7 +190,11 @@ fn plan_delta_spark_scan( "Delta partition resolved {store_references} store references with {memo_hits} memo hits" ); files = crate::execution::jni_api::get_runtime().block_on( - crate::execution::delta_dv::attach_access_plans(runtime_env, dv_files), + crate::execution::delta_dv::attach_access_plans( + runtime_env, + &object_store_url, + dv_files, + ), )?; } diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 06207f3667d..76b64f82151 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -32,7 +32,10 @@ pub(crate) const COMET_STRICT_TIMESTAMP_LABELS: &str = "spark.comet.testing.strictTimestampLabels.enabled"; pub(crate) const COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED: &str = "spark.comet.parquet.metadataCache.sharedWithinStage.enabled"; -/// Not a user config: `CometExecIterator.serializeCometSQLConfs` injects the stage of the task. +/// Not a user config: `CometExecIterator.serializeCometSQLConfs` injects the application and +/// stage of the task. Stage ids restart at zero in every `SparkContext`, so the application +/// identifies which context's stage 0 this is. +pub(crate) const COMET_APP_ID: &str = "spark.comet.internal.appId"; pub(crate) const COMET_STAGE_ID: &str = "spark.comet.internal.stageId"; pub(crate) const COMET_STAGE_ATTEMPT: &str = "spark.comet.internal.stageAttemptNumber"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 660d1073892..47e6482f38d 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -23,6 +23,7 @@ use crate::parquet::parquet_support::{ object_store_authority, ObjectStoreBackend, SparkParquetOptions, }; use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory; +use crate::parquet::shared_metadata_cache::scope_to_store; use arrow::datatypes::{Field, FieldRef, Schema, SchemaRef}; use datafusion::common::tree_node::TreeNode; use datafusion::config::{ParquetOptions, TableParquetOptions}; @@ -227,7 +228,10 @@ pub(crate) fn init_datasource_exec( // scan_efficiency_ratio, which continue to derive from bytes_scanned. 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 = scope_to_store( + runtime_env.cache_manager.get_file_metadata_cache(), + &object_store_url, + ); // // A rebase-enabled scan also has the factory stamp each file's INT96 leaf ordinals into // its footer metadata (see `datetime_rebase.rs`), which is how the expression adapter @@ -439,7 +443,7 @@ mod tests { use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::ExecutionPlan; use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename; - use futures::StreamExt; + use futures::{StreamExt, TryStreamExt}; use object_store::memory::InMemory; use object_store::throttle::{ThrottleConfig, ThrottledStore}; use object_store::{ObjectStore, ObjectStoreExt}; @@ -1307,10 +1311,13 @@ mod tests { // The per-task RuntimeEnv metadata cache holds this file's metadata with the page index // (column + offset index) loaded, even though this scan pushed no pruning predicate. - let cache = session_ctx - .runtime_env() - .cache_manager - .get_file_metadata_cache(); + let cache = scope_to_store( + session_ctx + .runtime_env() + .cache_manager + .get_file_metadata_cache(), + &ObjectStoreUrl::local_filesystem(), + ); let entry = cache .get(&location) .expect("file metadata should be cached"); @@ -2111,4 +2118,131 @@ mod tests { ); } } + + /// A parquet file of one `a` column holding `values`. + fn parquet_of(values: Vec) -> Bytes { + 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(values))], + ) + .unwrap(); + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buffer) + } + + /// Two stores that hold a file at the same path with the same size but different contents + /// must not hand each other's footer to a scan, whether the metadata cache is shared by the + /// tasks of a stage or private to one task. The path DataFusion keys the cache by has no + /// bucket, so the cache a scan uses is keyed by the scan's object store as well. + #[tokio::test] + async fn equal_paths_in_two_stores_do_not_share_footers() { + use crate::parquet::shared_metadata_cache::{SharedMetadataCaches, StageKey}; + use datafusion::execution::cache::cache_manager::CacheManagerConfig; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + + let location = object_store::path::Path::from("table/part-00000.parquet"); + let contents = [vec![1, 2, 3], vec![101, 102, 103]]; + let files: Vec = contents.iter().cloned().map(parquet_of).collect(); + assert_eq!( + files[0].len(), + files[1].len(), + "the files must have equal sizes for the cache to accept one for the other" + ); + let size = files[0].len() as u64; + let urls = [ + "s3+comet-0000000000000001-native://bucket-a", + "s3+comet-0000000000000001-native://bucket-b", + ]; + + for shared in [true, false] { + let registry = SharedMetadataCaches::default(); + let stage_cache = registry.cache_for(StageKey { + app_id: "app-1".to_string(), + stage_id: 0, + stage_attempt: 0, + }); + let mut scanned = Vec::new(); + for (url, file) in urls.iter().zip(&files) { + // One RuntimeEnv per task, as Comet builds them. A shared cache outlives them. + let mut cache_config = CacheManagerConfig::default(); + if shared { + cache_config = + cache_config.with_file_metadata_cache(Some(Arc::clone(&stage_cache))); + } + let runtime_env = RuntimeEnvBuilder::new() + .with_cache_manager(cache_config) + .build_arc() + .unwrap(); + let store: Arc = Arc::new(InMemory::new()); + store.put(&location, file.clone().into()).await.unwrap(); + let object_store_url = ObjectStoreUrl::parse(url).unwrap(); + runtime_env.register_object_store(object_store_url.as_ref(), store); + let session = Arc::new(SessionContext::new_with_config_rt( + datafusion::prelude::SessionConfig::new(), + runtime_env, + )); + + // Only the rows above 100 are wanted. The footer of the first file says it has + // none, so a scan handed that footer prunes every row of the second. + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let filter: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(100)))), + )); + let scan = init_datasource_exec( + Arc::clone(&schema), + Some(schema), + None, + object_store_url, + ObjectStoreBackend::Remote, + vec![vec![PartitionedFile::new(location.to_string(), size)]], + None, + Some(vec![filter]), + None, + "UTC", + true, + false, + false, + false, + &session, + false, + false, + false, + false, + "", + "", + ) + .unwrap(); + let batches: Vec = scan + .execute(0, session.task_ctx()) + .unwrap() + .try_collect() + .await + .unwrap(); + let rows: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + scanned.push(rows); + } + assert_eq!( + scanned, + vec![vec![], vec![101, 102, 103]], + "metadata cache shared within the stage: {shared}" + ); + } + } } diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs index 3066ad8f5f3..481f7d17f07 100644 --- a/native/core/src/parquet/shared_metadata_cache.rs +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -28,90 +28,190 @@ //! deletes one data file, reads again, and expects `FILE_NOT_EXIST`. Every new query is a new //! stage, so it opens its files and sees the deletion, as it does without the registry. //! -//! A stage attempt is part of the key because a retried stage may read a changed file. +//! A stage attempt is part of the key because a retried stage may read a changed file. So is the +//! Spark application: stage ids restart at zero in every `SparkContext`, and this registry is +//! process-wide, so a context started after another one stopped in the same JVM must not match +//! the stages the old one left behind. //! -//! The registry keeps the cache of the stages that still have a running task, plus the last -//! [`MAX_IDLE_CACHES`] that do not. The idle ones matter when an executor runs one task at a -//! time: the next task of the stage starts after the previous one released its plan. +//! The registry holds a weak reference to the cache of every stage with a running task, since the +//! tasks own it through their `RuntimeEnv`. It additionally holds strong references to the last +//! [`MAX_IDLE_CACHES`] caches it handed out. Those matter when an executor runs one task at a +//! time: the next task of the stage starts after the previous one released its plan. Whatever the +//! tasks do, at most [`MAX_IDLE_CACHES`] caches outlive the tasks that used them. +//! +//! The stage scopes the cache; it does not identify the files in it. DataFusion keys an entry by +//! the file's path within its object store, which drops the bucket, so [`scope_to_store`] gives +//! each scan a view of the cache that also keys by the scan's object store. use crate::execution::spark_config::{ - SparkConfig, COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED, COMET_STAGE_ATTEMPT, COMET_STAGE_ID, + SparkConfig, COMET_APP_ID, COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED, COMET_STAGE_ATTEMPT, + COMET_STAGE_ID, }; +use datafusion::common::TableReference; +use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::execution::cache::cache_manager::{ - FileMetadataCache, DEFAULT_METADATA_CACHE_LIMIT, + CachedFileMetadataEntry, FileMetadataCache, DEFAULT_METADATA_CACHE_LIMIT, }; use datafusion::execution::cache::default_cache::DefaultCache; -use std::collections::HashMap; -use std::sync::{Arc, Mutex, OnceLock}; +use datafusion::execution::cache::{Cache, CacheEntryInfo}; +use object_store::path::Path; +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex, OnceLock, Weak}; +use std::time::Duration; /// How many caches of stages with no running task the registry holds on to. const MAX_IDLE_CACHES: usize = 2; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(crate) struct StageKey { + pub app_id: String, pub stage_id: i32, pub stage_attempt: i32, } -struct Entry { - key: StageKey, - cache: Arc, +#[derive(Default)] +struct Registry { + /// Every cache that is still alive, running or retained. + alive: HashMap>, + /// The caches kept alive for the next task of their stage, least recently requested first. + retained: VecDeque<(StageKey, Arc)>, } #[derive(Default)] pub(crate) struct SharedMetadataCaches { - /// Ordered from least to most recently requested. - entries: Mutex>, + registry: Mutex, } impl SharedMetadataCaches { /// The cache shared by the tasks of `key`, created if no task of that stage has asked for one /// yet. - fn cache_for(&self, key: StageKey) -> Arc { - let mut entries = self.entries.lock().unwrap(); - let entry = match entries.iter().position(|e| e.key == key) { - Some(pos) => entries.remove(pos), - None => Entry { - key, - cache: Arc::new( + pub(crate) fn cache_for(&self, key: StageKey) -> Arc { + let mut registry = self.registry.lock().unwrap(); + let cache = registry + .alive + .get(&key) + .and_then(Weak::upgrade) + .unwrap_or_else(|| { + Arc::new( DefaultCache::new(DEFAULT_METADATA_CACHE_LIMIT) .with_name("SharedFileMetadataCache"), - ), - }, - }; - let cache = Arc::clone(&entry.cache); - entries.push(entry); - - // The registry holds one reference; a task holds another for as long as its - // `RuntimeEnv` lives. The newest entry also has the reference returned from here, which - // the task takes over, so it counts as idle at two. The newest entry is never evicted. - let newest = entries.len() - 1; - let is_idle = |index: usize, entry: &Entry| { - Arc::strong_count(&entry.cache) == if index == newest { 2 } else { 1 } - }; - let mut idle = entries - .iter() - .enumerate() - .filter(|(i, e)| is_idle(*i, e)) - .count(); - let mut index = 0; - entries.retain(|e| { - let evict = index != newest && idle > MAX_IDLE_CACHES && is_idle(index, e); - index += 1; - if evict { - idle -= 1; - } - !evict - }); + ) + }); + registry.alive.insert(key.clone(), Arc::downgrade(&cache)); + + registry.retained.retain(|(retained, _)| *retained != key); + registry.retained.push_back((key, Arc::clone(&cache))); + while registry.retained.len() > MAX_IDLE_CACHES { + registry.retained.pop_front(); + } + registry.alive.retain(|_, weak| weak.strong_count() > 0); cache } #[cfg(test)] fn len(&self) -> usize { - self.entries.lock().unwrap().len() + self.registry.lock().unwrap().alive.len() + } +} + +/// A view of a [`FileMetadataCache`] that keys every entry by the object store as well as the +/// path within it. +/// +/// DataFusion keys an entry by `ObjectMeta.location`, and the location of a file in a scan has no +/// bucket or authority, so equal paths in two stores collide. Two files at the same path with the +/// same size then pass the cache's `(size, modification time)` check, and a task would plan its +/// scan with the other file's footer and statistics. Only the key changes; the entries, the memory +/// limit and the eviction are those of the cache underneath, which is why a view can be made per +/// scan from the cache of the whole stage. `len`, `clear` and `list_entries` act on the whole +/// cache, not on this scope. +#[derive(Debug)] +struct StoreScopedCache { + scope: String, + inner: Arc, +} + +impl StoreScopedCache { + fn scoped(&self, key: &Path) -> Path { + std::iter::once(self.scope.clone()) + .chain(key.parts().map(|part| part.as_ref().to_string())) + .collect() } } +impl Cache for StoreScopedCache { + fn get(&self, key: &Path) -> Option { + self.inner.get(&self.scoped(key)) + } + + fn put(&self, key: &Path, value: CachedFileMetadataEntry) -> Option { + self.inner.put(&self.scoped(key), value) + } + + fn remove(&self, key: &Path) -> Option { + self.inner.remove(&self.scoped(key)) + } + + fn contains_key(&self, key: &Path) -> bool { + self.inner.contains_key(&self.scoped(key)) + } + + fn len(&self) -> usize { + self.inner.len() + } + + fn clear(&self) { + self.inner.clear() + } + + fn name(&self) -> String { + self.inner.name() + } + + fn cache_limit(&self) -> usize { + self.inner.cache_limit() + } + + fn update_cache_limit(&self, limit: usize) { + self.inner.update_cache_limit(limit) + } + + fn cache_ttl(&self) -> Option { + self.inner.cache_ttl() + } + + fn update_cache_ttl(&self, ttl: Option) { + self.inner.update_cache_ttl(ttl) + } + + fn drop_table_entries(&self, table_ref: &TableReference) -> datafusion::error::Result<()> { + self.inner.drop_table_entries(table_ref) + } + + fn list_entries( + &self, + ) -> datafusion::common::HashMap> { + self.inner.list_entries() + } +} + +/// The metadata cache a scan of `object_store_url` should use, given the one of its +/// `RuntimeEnv`. +/// +/// `object_store_url` is the URL the scan's store is registered under, which already carries the +/// backend and a hash of the store's configuration (see `object_store_registration_url`), so two +/// stores with equal paths never share an entry. The scan and any code that reads the footers +/// ahead of it, such as the Delta deletion vector planning, must pass the same URL to find each +/// other's entries. +pub(crate) fn scope_to_store( + cache: Arc, + object_store_url: &ObjectStoreUrl, +) -> Arc { + Arc::new(StoreScopedCache { + scope: object_store_url.as_str().to_string(), + inner: cache, + }) +} + /// The cache the task with these Spark configs should use, or `None` when sharing is off or the /// task's stage is unknown, in which case DataFusion's per-task cache applies. pub(crate) fn shared_metadata_cache_for_task( @@ -124,6 +224,7 @@ pub(crate) fn shared_metadata_cache_for_task( } let parse = |name: &str| spark_config.get(name)?.parse::().ok(); let key = StageKey { + app_id: spark_config.get(COMET_APP_ID)?.clone(), stage_id: parse(COMET_STAGE_ID)?, stage_attempt: parse(COMET_STAGE_ATTEMPT)?, }; @@ -140,6 +241,7 @@ mod tests { fn key(stage_id: i32, stage_attempt: i32) -> StageKey { StageKey { + app_id: "app-1".to_string(), stage_id, stage_attempt, } @@ -173,6 +275,41 @@ mod tests { assert_eq!(Arc::as_ptr(&next) as *const () as usize, address); } + #[test] + fn applications_do_not_share() { + // A context started after another one stopped in the same JVM reuses its stage ids. + let caches = SharedMetadataCaches::default(); + let old = caches.cache_for(key(1, 0)); + drop(old); + let mut other = key(1, 0); + other.app_id = "app-2".to_string(); + let new = caches.cache_for(other); + assert_eq!(new.len(), 0); + assert!(!Arc::ptr_eq(&new, &caches.cache_for(key(1, 0)))); + } + + #[test] + fn idle_retention_is_bounded_after_all_tasks_finish() { + let caches = SharedMetadataCaches::default(); + let handles: Vec<_> = (0..16) + .map(|stage| caches.cache_for(key(stage, 0))) + .collect(); + let weak: Vec<_> = handles.iter().map(Arc::downgrade).collect(); + drop(handles); + assert!(weak.iter().filter(|w| w.upgrade().is_some()).count() <= MAX_IDLE_CACHES); + } + + #[test] + fn a_dropped_cache_is_forgotten_by_the_next_request() { + let caches = SharedMetadataCaches::default(); + for stage in 0..16 { + drop(caches.cache_for(key(stage, 0))); + } + let _running = caches.cache_for(key(100, 0)); + // The caches of the first sixteen stages are gone but for the ones retained. + assert_eq!(caches.len(), MAX_IDLE_CACHES); + } + #[test] fn idle_caches_are_bounded_but_running_ones_are_not_dropped() { let caches = SharedMetadataCaches::default(); diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 38bcf6413da..54c7d3bfbb4 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -878,8 +878,13 @@ object CometExecIterator extends Logging { builder.putEntries("spark.executor.cores", executorCores.toString) // The stage of this task, which scopes the Parquet metadata cache the native scan shares - // between the tasks of one stage. - Option(TaskContext.get()).foreach { context => + // between the tasks of one stage. Stage ids restart at zero in every SparkContext and the + // native cache outlives a context stopped in the same JVM, so the application is part of it. + for { + context <- Option(TaskContext.get()) + appId <- SparkEnv.get.conf.getOption("spark.app.id") + } { + builder.putEntries("spark.comet.internal.appId", appId) builder.putEntries("spark.comet.internal.stageId", context.stageId().toString) builder.putEntries( "spark.comet.internal.stageAttemptNumber", diff --git a/spark/src/test/scala/org/apache/comet/exec/CometSharedMetadataCacheRestartSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometSharedMetadataCacheRestartSuite.scala new file mode 100644 index 00000000000..5fee767e8aa --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometSharedMetadataCacheRestartSuite.scala @@ -0,0 +1,110 @@ +/* + * 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.exec + +import java.io.File +import java.nio.file.{Files, Paths, StandardCopyOption} + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.commons.io.FileUtils +import org.apache.spark.sql.SparkSession + +/** + * The native Parquet scan shares a footer cache between the tasks of a stage in a registry that + * lives as long as the JVM. Stage ids restart at zero in every SparkContext, so a context started + * after another one stopped must not be handed the footers the old one left in it. + * + * Runs its own contexts, which is why it is not part of a suite that shares a session. + */ +class CometSharedMetadataCacheRestartSuite extends AnyFunSuite { + + private def newSession(): SparkSession = + SparkSession + .builder() + .master("local[2]") + .config("spark.plugins", "org.apache.spark.CometPlugin") + .config( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .config("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions") + .config("spark.comet.enabled", "true") + .config("spark.comet.exec.enabled", "true") + .config("spark.comet.parquet.metadataCache.sharedWithinStage.enabled", "true") + .config("spark.memory.offHeap.enabled", "true") + .config("spark.memory.offHeap.size", "256m") + // Equal sizes for different values need an encoding whose size does not follow the values. + .config("spark.sql.parquet.compression.codec", "uncompressed") + .config("spark.ui.enabled", "false") + .getOrCreate() + + /** + * Writes the ids in `[from, until)` as a single file named `part.parquet` in `table`, replacing + * the file already there, then counts the rows with an id of at least 1000 in a fresh context. + * The jobs run are the same ones each time, so the scan has the same stage id in every context. + */ + private def countFrom1000(scratch: File, table: File, from: Long, until: Long): (Long, Long) = { + val spark = newSession() + try { + val written = new File(scratch, s"written-$from") + spark.range(from, until).coalesce(1).write.parquet(written.getAbsolutePath) + val source = written + .listFiles() + .filter(_.getName.endsWith(".parquet")) + .head + .toPath + Files.copy( + source, + new File(table, "part.parquet").toPath, + StandardCopyOption.REPLACE_EXISTING) + + val df = spark.read.parquet(table.getAbsolutePath).where("id >= 1000") + val count = df.count() + val plan = df.queryExecution.executedPlan.toString + require(plan.contains("CometNativeScan"), s"Expected a native scan but got:\n$plan") + (count, new File(table, "part.parquet").length()) + } finally { + spark.stop() + } + } + + test("a restarted SparkContext does not read the footers of the previous one") { + assume(SparkSession.getDefaultSession.isEmpty, "another SparkSession is running in this JVM") + // java.io.tmpdir is pinned to target/tmp by the pom and may not exist yet + val tmpDir = Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + val scratch = Files.createTempDirectory(tmpDir, "comet-footer-restart").toFile + try { + val table = new File(scratch, "table") + assert(table.mkdirs()) + + // The first file holds ids 0 to 999, so its footer says nothing is at least 1000. + val (before, sizeBefore) = countFrom1000(scratch, table, 0, 1000) + assert(before == 0) + + // The second holds ids 1000 to 1999 at the same path with the same size. Its footer is + // the one the scan must read, not the cached one that would prune every row. + val (after, sizeAfter) = countFrom1000(scratch, table, 1000, 2000) + assert(sizeAfter == sizeBefore, "the replacement must have the size of the original") + assert(after == 1000) + } finally { + FileUtils.deleteDirectory(scratch) + } + } +} From e4801164964a6760b3f5e03959a956b007c972ea Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 19 Oct 2026 16:11:14 +0100 Subject: [PATCH 4/9] ci: run CometSharedMetadataCacheRestartSuite in the PR builds dev/ci/check-suites.py fails Preflight when a suite is missing from either PR build workflow. List it in the exec group of both, next to CometExecSuite. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index c9fc1cd7f72..4f5d31070b6 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -552,6 +552,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometSharedMetadataCacheRestartSuite org.apache.comet.exec.CometMergeRowsSuite org.apache.comet.exec.CometTopKSuite org.apache.comet.exec.CometEmptyRelationExecSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 593fb10417a..6f01fa686e2 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -218,6 +218,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometSharedMetadataCacheRestartSuite org.apache.comet.exec.CometMergeRowsSuite org.apache.comet.exec.CometTopKSuite org.apache.comet.exec.CometEmptyRelationExecSuite From d14249fd94d0fbf69a1035d9d757e226c516d140 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 22 Oct 2026 00:24:54 -0600 Subject: [PATCH 5/9] address review: scale the shared metadata cache per task slot, document its memory Give a stage's cache the per-task metadata limit for each task slot, and carry that limit in every task's cache manager config so that building a task's RuntimeEnv does not reset it. Say what the cache holds and that retained caches outlive their tasks in the config doc and memory_management.md. Test the unshared case in one RuntimeEnv, check the deleted-file read is native and fails on the missing file, and drop a redundant registry test. --- .../contributor-guide/memory_management.md | 7 ++ native/core/src/execution/jni_api.rs | 44 ++++++++-- native/core/src/parquet/parquet_exec.rs | 44 ++++++---- .../core/src/parquet/shared_metadata_cache.rs | 88 ++++++++++++------- .../scala/org/apache/comet/CometConf.scala | 12 ++- .../apache/comet/exec/CometExecSuite.scala | 13 ++- 6 files changed, 147 insertions(+), 61 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 184bf0de275..146fdd51739 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -508,6 +508,13 @@ diverge for several structural reasons: because several names suggest otherwise: the other Parquet codecs are pure Rust in this build, `snap` for Snappy, `lz4_flex` for LZ4 and `zlib-rs` for gzip, as is `libbz2-rs-sys` despite its name, so those allocations do pass through `GlobalAlloc` and are counted. +- **Native caches that outlive their task.** The native Parquet scan's cache of decoded footers and + page indexes is not reserved from the pool. With + `spark.comet.parquet.metadataCache.sharedWithinStage.enabled`, the default, the tasks of a stage + share one cache per executor, which holds up to 50 MiB for each task the executor runs at once, + and the two most recently used caches are kept after their tasks finish, for the next task of + their stage. So up to twice that limit stays resident with no task running. With the setting off, + each task has its own 50 MiB cache, freed with the task. - **Batches in flight across the FFI boundary.** Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 6a3638bccbb..ad8acf3d9e6 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -47,7 +47,7 @@ use arrow::array::{Array, RecordBatch}; use arrow::datatypes::DataType as ArrowDataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::config::SpillCompression; -use datafusion::execution::cache::cache_manager::CacheManagerConfig; +use datafusion::execution::cache::cache_manager::{CacheManagerConfig, FileMetadataCache}; use datafusion::execution::memory_pool::MemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::logical_expr::ScalarUDF; @@ -995,6 +995,15 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato } } +/// A task's cache manager config over the metadata cache it shares with its stage. The limit is +/// set here too: `CacheManager::try_new` applies the config's `metadata_cache_limit` to the cache +/// it is given, so leaving the default would reset the stage's cache to one task's limit. +fn shared_cache_manager_config(cache: Arc, limit: usize) -> CacheManagerConfig { + CacheManagerConfig::default() + .with_file_metadata_cache(Some(cache)) + .with_metadata_cache_limit(limit) +} + /// Configure DataFusion session context. fn prepare_datafusion_session_context( batch_size: usize, @@ -1013,10 +1022,9 @@ fn prepare_datafusion_session_context( // Let the tasks of one stage on this executor share the Parquet footer and page index cache // instead of each fetching and parsing the footers of the files it reads. - if let Some(cache) = shared_metadata_cache_for_task(spark_config) { - rt_config = rt_config.with_cache_manager( - CacheManagerConfig::default().with_file_metadata_cache(Some(cache)), - ); + let task_slots = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1) / task_cpus.max(1); + if let Some((cache, limit)) = shared_metadata_cache_for_task(spark_config, task_slots) { + rt_config = rt_config.with_cache_manager(shared_cache_manager_config(cache, limit)); } let mut session_config = SessionConfig::new() @@ -4325,4 +4333,30 @@ mod tests { ); release_plan(exec_context, &mut |_| Ok(())).unwrap(); } + + /// Every task builds a `RuntimeEnv` over its stage's metadata cache, and building one applies + /// the config's `metadata_cache_limit` to that cache. The config has to carry the stage's + /// limit, or each new task would shrink the cache back to one task's. + #[test] + fn a_task_keeps_the_limit_of_its_stage_metadata_cache() { + use crate::parquet::shared_metadata_cache::shared_cache_limit; + use datafusion::execution::cache::cache_manager::DEFAULT_METADATA_CACHE_LIMIT; + use datafusion::execution::cache::default_cache::DefaultCache; + + let limit = shared_cache_limit(8); + let cache: Arc = Arc::new(DefaultCache::new(limit)); + RuntimeEnvBuilder::new() + .with_cache_manager(shared_cache_manager_config(Arc::clone(&cache), limit)) + .build() + .unwrap(); + assert_eq!(cache.cache_limit(), limit); + + RuntimeEnvBuilder::new() + .with_cache_manager( + CacheManagerConfig::default().with_file_metadata_cache(Some(Arc::clone(&cache))), + ) + .build() + .unwrap(); + assert_eq!(cache.cache_limit(), DEFAULT_METADATA_CACHE_LIMIT); + } } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 6c60e517bbd..7981d3c9669 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -2166,7 +2166,9 @@ mod tests { #[tokio::test] async fn equal_paths_in_two_stores_do_not_share_footers() { use crate::parquet::shared_metadata_cache::{SharedMetadataCaches, StageKey}; - use datafusion::execution::cache::cache_manager::CacheManagerConfig; + use datafusion::execution::cache::cache_manager::{ + CacheManagerConfig, DEFAULT_METADATA_CACHE_LIMIT, + }; use datafusion::execution::runtime_env::RuntimeEnvBuilder; let location = object_store::path::Path::from("table/part-00000.parquet"); @@ -2185,23 +2187,33 @@ mod tests { for shared in [true, false] { let registry = SharedMetadataCaches::default(); - let stage_cache = registry.cache_for(StageKey { - app_id: "app-1".to_string(), - stage_id: 0, - stage_attempt: 0, - }); + // Without sharing, the two scans can only meet in one task's cache, so they run in + // one plan with one RuntimeEnv. + let task_env = (!shared).then(|| RuntimeEnvBuilder::new().build_arc().unwrap()); let mut scanned = Vec::new(); for (url, file) in urls.iter().zip(&files) { - // One RuntimeEnv per task, as Comet builds them. A shared cache outlives them. - let mut cache_config = CacheManagerConfig::default(); - if shared { - cache_config = - cache_config.with_file_metadata_cache(Some(Arc::clone(&stage_cache))); - } - let runtime_env = RuntimeEnvBuilder::new() - .with_cache_manager(cache_config) - .build_arc() - .unwrap(); + // With sharing, one RuntimeEnv per task, as Comet builds them, over the cache of + // the stage, which outlives them. + let runtime_env = match &task_env { + Some(env) => Arc::clone(env), + None => { + let stage_cache = registry.cache_for( + StageKey { + app_id: "app-1".to_string(), + stage_id: 0, + stage_attempt: 0, + }, + DEFAULT_METADATA_CACHE_LIMIT, + ); + RuntimeEnvBuilder::new() + .with_cache_manager( + CacheManagerConfig::default() + .with_file_metadata_cache(Some(stage_cache)), + ) + .build_arc() + .unwrap() + } + }; let store: Arc = Arc::new(InMemory::new()); store.put(&location, file.clone().into()).await.unwrap(); let object_store_url = ObjectStoreUrl::parse(url).unwrap(); diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs index 481f7d17f07..cbbc7be522a 100644 --- a/native/core/src/parquet/shared_metadata_cache.rs +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -39,6 +39,13 @@ //! time: the next task of the stage starts after the previous one released its plan. Whatever the //! tasks do, at most [`MAX_IDLE_CACHES`] caches outlive the tasks that used them. //! +//! A stage's cache holds DataFusion's per-task limit for every task the executor runs at once, so +//! that sharing does not shrink the room each task had with a cache of its own: tasks that read +//! different files would otherwise evict each other's entries, among them the footers the Delta +//! deletion vector planning loads before the scan opens them. Shared entries are stored once, so +//! the tasks of a stage hold no more than they did, but the retained caches can hold as much after +//! their tasks finish. +//! //! The stage scopes the cache; it does not identify the files in it. DataFusion keys an entry by //! the file's path within its object store, which drops the bucket, so [`scope_to_store`] gives //! each scan a view of the cache that also keys by the scan's object store. @@ -83,19 +90,16 @@ pub(crate) struct SharedMetadataCaches { } impl SharedMetadataCaches { - /// The cache shared by the tasks of `key`, created if no task of that stage has asked for one - /// yet. - pub(crate) fn cache_for(&self, key: StageKey) -> Arc { + /// The cache shared by the tasks of `key`, created with a memory limit of `limit` bytes if no + /// task of that stage has asked for one yet. + pub(crate) fn cache_for(&self, key: StageKey, limit: usize) -> Arc { let mut registry = self.registry.lock().unwrap(); let cache = registry .alive .get(&key) .and_then(Weak::upgrade) .unwrap_or_else(|| { - Arc::new( - DefaultCache::new(DEFAULT_METADATA_CACHE_LIMIT) - .with_name("SharedFileMetadataCache"), - ) + Arc::new(DefaultCache::new(limit).with_name("SharedFileMetadataCache")) }); registry.alive.insert(key.clone(), Arc::downgrade(&cache)); @@ -212,11 +216,22 @@ pub(crate) fn scope_to_store( }) } -/// The cache the task with these Spark configs should use, or `None` when sharing is off or the -/// task's stage is unknown, in which case DataFusion's per-task cache applies. +/// The memory limit of the cache shared by the tasks of a stage on an executor that runs +/// `task_slots` tasks at once: DataFusion's limit for one task's cache, for each of them. +pub(crate) fn shared_cache_limit(task_slots: usize) -> usize { + DEFAULT_METADATA_CACHE_LIMIT.saturating_mul(task_slots.max(1)) +} + +/// The cache the task with these Spark configs should use, with its memory limit, or `None` when +/// sharing is off or the task's stage is unknown, in which case DataFusion's per-task cache +/// applies. `task_slots` is how many tasks the executor runs at once. +/// +/// The limit must also be the `metadata_cache_limit` of the task's `CacheManagerConfig`, since +/// `CacheManager::try_new` applies that limit to the cache it is given. pub(crate) fn shared_metadata_cache_for_task( spark_config: &HashMap, -) -> Option> { + task_slots: usize, +) -> Option<(Arc, usize)> { static CACHES: OnceLock = OnceLock::new(); if !spark_config.get_bool(COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED) { @@ -228,11 +243,13 @@ pub(crate) fn shared_metadata_cache_for_task( stage_id: parse(COMET_STAGE_ID)?, stage_attempt: parse(COMET_STAGE_ATTEMPT)?, }; - Some( + let limit = shared_cache_limit(task_slots); + Some(( CACHES .get_or_init(SharedMetadataCaches::default) - .cache_for(key), - ) + .cache_for(key, limit), + limit, + )) } #[cfg(test)] @@ -250,17 +267,17 @@ mod tests { #[test] fn tasks_of_one_stage_share_a_cache() { let caches = SharedMetadataCaches::default(); - let a = caches.cache_for(key(1, 0)); - let b = caches.cache_for(key(1, 0)); + let a = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); + let b = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); assert!(Arc::ptr_eq(&a, &b)); } #[test] fn stages_and_attempts_do_not_share() { let caches = SharedMetadataCaches::default(); - let first = caches.cache_for(key(1, 0)); - let other_stage = caches.cache_for(key(2, 0)); - let retry = caches.cache_for(key(1, 1)); + let first = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); + let other_stage = caches.cache_for(key(2, 0), DEFAULT_METADATA_CACHE_LIMIT); + let retry = caches.cache_for(key(1, 1), DEFAULT_METADATA_CACHE_LIMIT); assert!(!Arc::ptr_eq(&first, &other_stage)); assert!(!Arc::ptr_eq(&first, &retry)); } @@ -268,10 +285,10 @@ mod tests { #[test] fn an_idle_cache_is_kept_for_the_next_task_of_its_stage() { let caches = SharedMetadataCaches::default(); - let first = caches.cache_for(key(1, 0)); + let first = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); let address = Arc::as_ptr(&first) as *const () as usize; drop(first); - let next = caches.cache_for(key(1, 0)); + let next = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); assert_eq!(Arc::as_ptr(&next) as *const () as usize, address); } @@ -279,20 +296,23 @@ mod tests { fn applications_do_not_share() { // A context started after another one stopped in the same JVM reuses its stage ids. let caches = SharedMetadataCaches::default(); - let old = caches.cache_for(key(1, 0)); + let old = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); drop(old); let mut other = key(1, 0); other.app_id = "app-2".to_string(); - let new = caches.cache_for(other); + let new = caches.cache_for(other, DEFAULT_METADATA_CACHE_LIMIT); assert_eq!(new.len(), 0); - assert!(!Arc::ptr_eq(&new, &caches.cache_for(key(1, 0)))); + assert!(!Arc::ptr_eq( + &new, + &caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT) + )); } #[test] fn idle_retention_is_bounded_after_all_tasks_finish() { let caches = SharedMetadataCaches::default(); let handles: Vec<_> = (0..16) - .map(|stage| caches.cache_for(key(stage, 0))) + .map(|stage| caches.cache_for(key(stage, 0), DEFAULT_METADATA_CACHE_LIMIT)) .collect(); let weak: Vec<_> = handles.iter().map(Arc::downgrade).collect(); drop(handles); @@ -300,25 +320,25 @@ mod tests { } #[test] - fn a_dropped_cache_is_forgotten_by_the_next_request() { + fn a_stage_cache_holds_the_per_task_limit_for_every_task_slot() { let caches = SharedMetadataCaches::default(); - for stage in 0..16 { - drop(caches.cache_for(key(stage, 0))); - } - let _running = caches.cache_for(key(100, 0)); - // The caches of the first sixteen stages are gone but for the ones retained. - assert_eq!(caches.len(), MAX_IDLE_CACHES); + let cache = caches.cache_for(key(1, 0), shared_cache_limit(8)); + assert_eq!(cache.cache_limit(), 8 * DEFAULT_METADATA_CACHE_LIMIT); + assert_eq!(shared_cache_limit(0), DEFAULT_METADATA_CACHE_LIMIT); } #[test] fn idle_caches_are_bounded_but_running_ones_are_not_dropped() { let caches = SharedMetadataCaches::default(); - let running = caches.cache_for(key(0, 0)); + let running = caches.cache_for(key(0, 0), DEFAULT_METADATA_CACHE_LIMIT); for stage in 1..=10 { - drop(caches.cache_for(key(stage, 0))); + drop(caches.cache_for(key(stage, 0), DEFAULT_METADATA_CACHE_LIMIT)); } // The running stage plus the idle ones that fit. assert_eq!(caches.len(), 1 + MAX_IDLE_CACHES); - assert!(Arc::ptr_eq(&running, &caches.cache_for(key(0, 0)))); + assert!(Arc::ptr_eq( + &running, + &caches.cache_for(key(0, 0), DEFAULT_METADATA_CACHE_LIMIT) + )); } } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 9a11e9c674f..7b59a28d9a8 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -196,10 +196,14 @@ object CometConf extends ShimCometConf { .doc( "When enabled, the tasks of one Spark stage that run on the same executor share the " + "native Parquet scan's cache of file footers and page indexes, so a file split " + - "across many tasks has its footer fetched and parsed once per executor instead of " + - "once per task. The cache is never shared between stages or stage attempts, so a " + - "query always reads the footers of its files afresh. When disabled, each task " + - "keeps its own cache.") + "across many tasks can have its footer fetched and parsed once per executor instead " + + "of once per task. Tasks that start together each fetch it, so the saving grows " + + "with the number of splits per file. The cache is never shared between stages or " + + "stage attempts, so a query always reads the footers of its files afresh. A stage's " + + "cache holds up to 50 MiB of decoded footers and page indexes for each task the " + + "executor runs at once, outside Spark's memory accounting, and up to two such " + + "caches are kept after their tasks finish, for the next task of their stage. When " + + "disabled, each task keeps its own 50 MiB cache, which is freed when the task ends.") .booleanConf .createWithDefault(true) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 9073e0edb90..0e966e35ab6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3569,13 +3569,22 @@ class CometExecSuite extends CometTestBase { // The relation keeps its file listing, as a cached table does, so the second query is // given a file that no longer exists. It must find that out by opening the file. val table = spark.read.parquet(path) - assert(table.count() == 1000) + val first = table.selectExpr("count(*)") + assert(first.collect().head.getLong(0) == 1000) + assert( + find(first.queryExecution.executedPlan)(_.isInstanceOf[CometNativeScanExec]).isDefined, + s"Expected a native scan:\n${first.queryExecution.executedPlan}") val deleted = new java.io.File(dir, "table") .listFiles() .filter(_.getName.endsWith(".parquet")) .head assert(deleted.delete()) - intercept[SparkException](table.count()) + val error = intercept[SparkException](table.selectExpr("count(*)").collect()) + // FAILED_READ_FILE.FILE_NOT_EXIST on Spark 4.0+, as HiveMetadataCacheSuite expects, and + // Spark 3.x's file-not-found read error, both caused by the missing file. + assert( + causeChain(error).exists(_.isInstanceOf[java.io.FileNotFoundException]), + s"Expected the missing file to fail the read, got: $error") } } } From 40beedd7483db4f12408694c35c07a2662758e94 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 22 Oct 2026 01:02:08 -0600 Subject: [PATCH 6/9] fix: size the shared metadata cache by the stage's running tasks The executor's core count that reaches native code is an estimate on a standalone executor without spark.executor.cores, so sizing the cache by it could give a two-core executor the limit of the host's processors. Count each task's plan in while it holds its stage's cache and out when it drops it, and size the cache by the tasks running. A cache with no task running shrinks to one task's limit, so a retained cache holds no more than a task's own cache would. The handle ignores the limit a task's CacheManager sets. --- .../contributor-guide/memory_management.md | 8 +- native/core/src/execution/jni_api.rs | 44 +-- native/core/src/parquet/parquet_exec.rs | 17 +- .../core/src/parquet/shared_metadata_cache.rs | 278 +++++++++++++----- .../scala/org/apache/comet/CometConf.scala | 9 +- 5 files changed, 223 insertions(+), 133 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 146fdd51739..51d8ae92da0 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -511,10 +511,10 @@ diverge for several structural reasons: - **Native caches that outlive their task.** The native Parquet scan's cache of decoded footers and page indexes is not reserved from the pool. With `spark.comet.parquet.metadataCache.sharedWithinStage.enabled`, the default, the tasks of a stage - share one cache per executor, which holds up to 50 MiB for each task the executor runs at once, - and the two most recently used caches are kept after their tasks finish, for the next task of - their stage. So up to twice that limit stays resident with no task running. With the setting off, - each task has its own 50 MiB cache, freed with the task. + share one cache per executor, which holds up to 50 MiB for each of its running tasks. The two + most recently used caches are kept after their tasks finish, for the next task of their stage, + and shrink to 50 MiB each, so up to 100 MiB stays resident with no task running. With the + setting off, each task has its own 50 MiB cache, freed with the task. - **Batches in flight across the FFI boundary.** Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index ad8acf3d9e6..6a3638bccbb 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -47,7 +47,7 @@ use arrow::array::{Array, RecordBatch}; use arrow::datatypes::DataType as ArrowDataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::config::SpillCompression; -use datafusion::execution::cache::cache_manager::{CacheManagerConfig, FileMetadataCache}; +use datafusion::execution::cache::cache_manager::CacheManagerConfig; use datafusion::execution::memory_pool::MemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::logical_expr::ScalarUDF; @@ -995,15 +995,6 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato } } -/// A task's cache manager config over the metadata cache it shares with its stage. The limit is -/// set here too: `CacheManager::try_new` applies the config's `metadata_cache_limit` to the cache -/// it is given, so leaving the default would reset the stage's cache to one task's limit. -fn shared_cache_manager_config(cache: Arc, limit: usize) -> CacheManagerConfig { - CacheManagerConfig::default() - .with_file_metadata_cache(Some(cache)) - .with_metadata_cache_limit(limit) -} - /// Configure DataFusion session context. fn prepare_datafusion_session_context( batch_size: usize, @@ -1022,9 +1013,10 @@ fn prepare_datafusion_session_context( // Let the tasks of one stage on this executor share the Parquet footer and page index cache // instead of each fetching and parsing the footers of the files it reads. - let task_slots = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1) / task_cpus.max(1); - if let Some((cache, limit)) = shared_metadata_cache_for_task(spark_config, task_slots) { - rt_config = rt_config.with_cache_manager(shared_cache_manager_config(cache, limit)); + if let Some(cache) = shared_metadata_cache_for_task(spark_config) { + rt_config = rt_config.with_cache_manager( + CacheManagerConfig::default().with_file_metadata_cache(Some(cache)), + ); } let mut session_config = SessionConfig::new() @@ -4333,30 +4325,4 @@ mod tests { ); release_plan(exec_context, &mut |_| Ok(())).unwrap(); } - - /// Every task builds a `RuntimeEnv` over its stage's metadata cache, and building one applies - /// the config's `metadata_cache_limit` to that cache. The config has to carry the stage's - /// limit, or each new task would shrink the cache back to one task's. - #[test] - fn a_task_keeps_the_limit_of_its_stage_metadata_cache() { - use crate::parquet::shared_metadata_cache::shared_cache_limit; - use datafusion::execution::cache::cache_manager::DEFAULT_METADATA_CACHE_LIMIT; - use datafusion::execution::cache::default_cache::DefaultCache; - - let limit = shared_cache_limit(8); - let cache: Arc = Arc::new(DefaultCache::new(limit)); - RuntimeEnvBuilder::new() - .with_cache_manager(shared_cache_manager_config(Arc::clone(&cache), limit)) - .build() - .unwrap(); - assert_eq!(cache.cache_limit(), limit); - - RuntimeEnvBuilder::new() - .with_cache_manager( - CacheManagerConfig::default().with_file_metadata_cache(Some(Arc::clone(&cache))), - ) - .build() - .unwrap(); - assert_eq!(cache.cache_limit(), DEFAULT_METADATA_CACHE_LIMIT); - } } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 7981d3c9669..55d10bc764c 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -2166,9 +2166,7 @@ mod tests { #[tokio::test] async fn equal_paths_in_two_stores_do_not_share_footers() { use crate::parquet::shared_metadata_cache::{SharedMetadataCaches, StageKey}; - use datafusion::execution::cache::cache_manager::{ - CacheManagerConfig, DEFAULT_METADATA_CACHE_LIMIT, - }; + use datafusion::execution::cache::cache_manager::{CacheManagerConfig, FileMetadataCache}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; let location = object_store::path::Path::from("table/part-00000.parquet"); @@ -2197,14 +2195,11 @@ mod tests { let runtime_env = match &task_env { Some(env) => Arc::clone(env), None => { - let stage_cache = registry.cache_for( - StageKey { - app_id: "app-1".to_string(), - stage_id: 0, - stage_attempt: 0, - }, - DEFAULT_METADATA_CACHE_LIMIT, - ); + let stage_cache: Arc = registry.cache_for(StageKey { + app_id: "app-1".to_string(), + stage_id: 0, + stage_attempt: 0, + }); RuntimeEnvBuilder::new() .with_cache_manager( CacheManagerConfig::default() diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs index cbbc7be522a..4ef7cb9c545 100644 --- a/native/core/src/parquet/shared_metadata_cache.rs +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -39,12 +39,14 @@ //! time: the next task of the stage starts after the previous one released its plan. Whatever the //! tasks do, at most [`MAX_IDLE_CACHES`] caches outlive the tasks that used them. //! -//! A stage's cache holds DataFusion's per-task limit for every task the executor runs at once, so -//! that sharing does not shrink the room each task had with a cache of its own: tasks that read -//! different files would otherwise evict each other's entries, among them the footers the Delta -//! deletion vector planning loads before the scan opens them. Shared entries are stored once, so -//! the tasks of a stage hold no more than they did, but the retained caches can hold as much after -//! their tasks finish. +//! A stage's cache holds DataFusion's per-task limit for each of its tasks that is running, and one +//! task's limit once none is, so that sharing does not shrink the room each task had with a cache +//! of its own: tasks that read different files would otherwise evict each other's entries, among +//! them the footers the Delta deletion vector planning loads before the scan opens them. Each task +//! takes its share when it asks for the cache and gives it back when its `RuntimeEnv` drops it, so +//! the limit follows the tasks that actually run, whatever the executor's configured cores say. +//! Shared entries are stored once, so the tasks of a stage hold no more than they did, and a +//! retained cache holds no more than one task's. //! //! The stage scopes the cache; it does not identify the files in it. DataFusion keys an entry by //! the file's path within its object store, which drops the bucket, so [`scope_to_store`] gives @@ -76,12 +78,127 @@ pub(crate) struct StageKey { pub stage_attempt: i32, } +/// The cache of one stage, and how many of its tasks are running. +struct StageCache { + cache: Arc>, + running: Mutex, +} + +impl StageCache { + fn new() -> Self { + Self { + cache: Arc::new( + DefaultCache::new(DEFAULT_METADATA_CACHE_LIMIT) + .with_name("SharedFileMetadataCache"), + ), + running: Mutex::new(0), + } + } + + /// Counts a task in or out, and sizes the cache for the tasks left running. Lowering the limit + /// evicts the least recently used entries down to it. + fn update_running(&self, update: impl FnOnce(usize) -> usize) { + let mut running = self.running.lock().unwrap(); + *running = update(*running); + self.cache + .update_cache_limit(DEFAULT_METADATA_CACHE_LIMIT.saturating_mul((*running).max(1))); + } +} + +/// One task's handle on its stage's cache, which counts the task as running for as long as it +/// lives. The task's `RuntimeEnv` and the scans' views of the cache hold it, so it drops when the +/// task's plan does. Strictly it counts native plans, of which a task usually has one; a task that +/// runs several at once is counted once for each, which errs toward a larger cache. +/// +/// It is the cache the task's `CacheManager` is given. `CacheManager::try_new` sets the limit of +/// that cache to the one in its config, so this ignores limit updates: the stage's running tasks +/// decide it. +pub(crate) struct TaskCache { + stage: Arc, +} + +impl TaskCache { + fn new(stage: Arc) -> Self { + stage.update_running(|running| running + 1); + Self { stage } + } +} + +impl Drop for TaskCache { + fn drop(&mut self) { + self.stage + .update_running(|running| running.saturating_sub(1)); + } +} + +impl std::fmt::Debug for TaskCache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TaskCache") + .field("running", &*self.stage.running.lock().unwrap()) + .finish() + } +} + +impl Cache for TaskCache { + fn get(&self, key: &Path) -> Option { + self.stage.cache.get(key) + } + + fn put(&self, key: &Path, value: CachedFileMetadataEntry) -> Option { + self.stage.cache.put(key, value) + } + + fn remove(&self, key: &Path) -> Option { + self.stage.cache.remove(key) + } + + fn contains_key(&self, key: &Path) -> bool { + self.stage.cache.contains_key(key) + } + + fn len(&self) -> usize { + self.stage.cache.len() + } + + fn clear(&self) { + self.stage.cache.clear() + } + + fn name(&self) -> String { + self.stage.cache.name() + } + + fn cache_limit(&self) -> usize { + self.stage.cache.cache_limit() + } + + fn update_cache_limit(&self, _limit: usize) {} + + fn cache_ttl(&self) -> Option { + self.stage.cache.cache_ttl() + } + + fn update_cache_ttl(&self, ttl: Option) { + self.stage.cache.update_cache_ttl(ttl) + } + + fn drop_table_entries(&self, table_ref: &TableReference) -> datafusion::error::Result<()> { + self.stage.cache.drop_table_entries(table_ref) + } + + fn list_entries( + &self, + ) -> datafusion::common::HashMap> { + self.stage.cache.list_entries() + } +} + #[derive(Default)] struct Registry { - /// Every cache that is still alive, running or retained. - alive: HashMap>, + /// Every stage cache that is still alive, running or retained. + alive: HashMap>, /// The caches kept alive for the next task of their stage, least recently requested first. - retained: VecDeque<(StageKey, Arc)>, + retained: VecDeque<(StageKey, Arc)>, } #[derive(Default)] @@ -90,26 +207,24 @@ pub(crate) struct SharedMetadataCaches { } impl SharedMetadataCaches { - /// The cache shared by the tasks of `key`, created with a memory limit of `limit` bytes if no - /// task of that stage has asked for one yet. - pub(crate) fn cache_for(&self, key: StageKey, limit: usize) -> Arc { + /// A task's handle on the cache shared by the tasks of `key`, created if no task of that stage + /// has asked for one yet. + pub(crate) fn cache_for(&self, key: StageKey) -> Arc { let mut registry = self.registry.lock().unwrap(); - let cache = registry + let stage = registry .alive .get(&key) .and_then(Weak::upgrade) - .unwrap_or_else(|| { - Arc::new(DefaultCache::new(limit).with_name("SharedFileMetadataCache")) - }); - registry.alive.insert(key.clone(), Arc::downgrade(&cache)); + .unwrap_or_else(|| Arc::new(StageCache::new())); + registry.alive.insert(key.clone(), Arc::downgrade(&stage)); registry.retained.retain(|(retained, _)| *retained != key); - registry.retained.push_back((key, Arc::clone(&cache))); + registry.retained.push_back((key, Arc::clone(&stage))); while registry.retained.len() > MAX_IDLE_CACHES { registry.retained.pop_front(); } registry.alive.retain(|_, weak| weak.strong_count() > 0); - cache + Arc::new(TaskCache::new(stage)) } #[cfg(test)] @@ -216,22 +331,11 @@ pub(crate) fn scope_to_store( }) } -/// The memory limit of the cache shared by the tasks of a stage on an executor that runs -/// `task_slots` tasks at once: DataFusion's limit for one task's cache, for each of them. -pub(crate) fn shared_cache_limit(task_slots: usize) -> usize { - DEFAULT_METADATA_CACHE_LIMIT.saturating_mul(task_slots.max(1)) -} - -/// The cache the task with these Spark configs should use, with its memory limit, or `None` when -/// sharing is off or the task's stage is unknown, in which case DataFusion's per-task cache -/// applies. `task_slots` is how many tasks the executor runs at once. -/// -/// The limit must also be the `metadata_cache_limit` of the task's `CacheManagerConfig`, since -/// `CacheManager::try_new` applies that limit to the cache it is given. +/// The cache the task with these Spark configs should use, or `None` when sharing is off or the +/// task's stage is unknown, in which case DataFusion's per-task cache applies. pub(crate) fn shared_metadata_cache_for_task( spark_config: &HashMap, - task_slots: usize, -) -> Option<(Arc, usize)> { +) -> Option> { static CACHES: OnceLock = OnceLock::new(); if !spark_config.get_bool(COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED) { @@ -243,13 +347,10 @@ pub(crate) fn shared_metadata_cache_for_task( stage_id: parse(COMET_STAGE_ID)?, stage_attempt: parse(COMET_STAGE_ATTEMPT)?, }; - let limit = shared_cache_limit(task_slots); - Some(( - CACHES - .get_or_init(SharedMetadataCaches::default) - .cache_for(key, limit), - limit, - )) + let cache: Arc = CACHES + .get_or_init(SharedMetadataCaches::default) + .cache_for(key); + Some(cache) } #[cfg(test)] @@ -264,81 +365,108 @@ mod tests { } } + fn same_stage(a: &TaskCache, b: &TaskCache) -> bool { + Arc::ptr_eq(&a.stage, &b.stage) + } + #[test] fn tasks_of_one_stage_share_a_cache() { let caches = SharedMetadataCaches::default(); - let a = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); - let b = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); - assert!(Arc::ptr_eq(&a, &b)); + let a = caches.cache_for(key(1, 0)); + let b = caches.cache_for(key(1, 0)); + assert!(same_stage(&a, &b)); } #[test] fn stages_and_attempts_do_not_share() { let caches = SharedMetadataCaches::default(); - let first = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); - let other_stage = caches.cache_for(key(2, 0), DEFAULT_METADATA_CACHE_LIMIT); - let retry = caches.cache_for(key(1, 1), DEFAULT_METADATA_CACHE_LIMIT); - assert!(!Arc::ptr_eq(&first, &other_stage)); - assert!(!Arc::ptr_eq(&first, &retry)); + let first = caches.cache_for(key(1, 0)); + let other_stage = caches.cache_for(key(2, 0)); + let retry = caches.cache_for(key(1, 1)); + assert!(!same_stage(&first, &other_stage)); + assert!(!same_stage(&first, &retry)); } #[test] fn an_idle_cache_is_kept_for_the_next_task_of_its_stage() { let caches = SharedMetadataCaches::default(); - let first = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); - let address = Arc::as_ptr(&first) as *const () as usize; + let first = caches.cache_for(key(1, 0)); + let stage = Arc::downgrade(&first.stage); drop(first); - let next = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); - assert_eq!(Arc::as_ptr(&next) as *const () as usize, address); + let next = caches.cache_for(key(1, 0)); + assert!(Arc::ptr_eq(&next.stage, &stage.upgrade().unwrap())); } #[test] fn applications_do_not_share() { // A context started after another one stopped in the same JVM reuses its stage ids. let caches = SharedMetadataCaches::default(); - let old = caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT); + let old = caches.cache_for(key(1, 0)); drop(old); let mut other = key(1, 0); other.app_id = "app-2".to_string(); - let new = caches.cache_for(other, DEFAULT_METADATA_CACHE_LIMIT); + let new = caches.cache_for(other); assert_eq!(new.len(), 0); - assert!(!Arc::ptr_eq( - &new, - &caches.cache_for(key(1, 0), DEFAULT_METADATA_CACHE_LIMIT) - )); + assert!(!same_stage(&new, &caches.cache_for(key(1, 0)))); } #[test] fn idle_retention_is_bounded_after_all_tasks_finish() { let caches = SharedMetadataCaches::default(); let handles: Vec<_> = (0..16) - .map(|stage| caches.cache_for(key(stage, 0), DEFAULT_METADATA_CACHE_LIMIT)) + .map(|stage| caches.cache_for(key(stage, 0))) .collect(); - let weak: Vec<_> = handles.iter().map(Arc::downgrade).collect(); + let stages: Vec<_> = handles.iter().map(|h| Arc::downgrade(&h.stage)).collect(); drop(handles); - assert!(weak.iter().filter(|w| w.upgrade().is_some()).count() <= MAX_IDLE_CACHES); - } - - #[test] - fn a_stage_cache_holds_the_per_task_limit_for_every_task_slot() { - let caches = SharedMetadataCaches::default(); - let cache = caches.cache_for(key(1, 0), shared_cache_limit(8)); - assert_eq!(cache.cache_limit(), 8 * DEFAULT_METADATA_CACHE_LIMIT); - assert_eq!(shared_cache_limit(0), DEFAULT_METADATA_CACHE_LIMIT); + assert!(stages.iter().filter(|w| w.upgrade().is_some()).count() <= MAX_IDLE_CACHES); } #[test] fn idle_caches_are_bounded_but_running_ones_are_not_dropped() { let caches = SharedMetadataCaches::default(); - let running = caches.cache_for(key(0, 0), DEFAULT_METADATA_CACHE_LIMIT); + let running = caches.cache_for(key(0, 0)); for stage in 1..=10 { - drop(caches.cache_for(key(stage, 0), DEFAULT_METADATA_CACHE_LIMIT)); + drop(caches.cache_for(key(stage, 0))); } // The running stage plus the idle ones that fit. assert_eq!(caches.len(), 1 + MAX_IDLE_CACHES); - assert!(Arc::ptr_eq( - &running, - &caches.cache_for(key(0, 0), DEFAULT_METADATA_CACHE_LIMIT) - )); + assert!(same_stage(&running, &caches.cache_for(key(0, 0)))); + } + + #[test] + fn a_stage_cache_holds_one_task_limit_per_running_task() { + let caches = SharedMetadataCaches::default(); + let tasks: Vec<_> = (0..3).map(|_| caches.cache_for(key(1, 0))).collect(); + assert_eq!(tasks[0].cache_limit(), 3 * DEFAULT_METADATA_CACHE_LIMIT); + // The limit a task's CacheManager sets does not apply to the shared cache. + tasks[0].update_cache_limit(DEFAULT_METADATA_CACHE_LIMIT / 2); + assert_eq!(tasks[0].cache_limit(), 3 * DEFAULT_METADATA_CACHE_LIMIT); + + let stage = Arc::clone(&tasks[0].stage); + drop(tasks); + // Idle, retained for the next task, with one task's limit. + assert_eq!(stage.cache.cache_limit(), DEFAULT_METADATA_CACHE_LIMIT); + let next = caches.cache_for(key(1, 0)); + assert_eq!(next.cache_limit(), DEFAULT_METADATA_CACHE_LIMIT); + } + + #[test] + fn a_task_runtime_env_does_not_reset_the_stage_limit() { + use datafusion::execution::cache::cache_manager::CacheManagerConfig; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + + let caches = SharedMetadataCaches::default(); + let other = caches.cache_for(key(1, 0)); + let task: Arc = caches.cache_for(key(1, 0)); + // CacheManager::try_new applies the config's default limit to the cache it is given. + RuntimeEnvBuilder::new() + .with_cache_manager( + CacheManagerConfig::default().with_file_metadata_cache(Some(Arc::clone(&task))), + ) + .build() + .unwrap(); + assert_eq!(task.cache_limit(), 2 * DEFAULT_METADATA_CACHE_LIMIT); + drop(other); + assert_eq!(task.cache_limit(), DEFAULT_METADATA_CACHE_LIMIT); } } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 7b59a28d9a8..7e98eb1ea24 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -200,10 +200,11 @@ object CometConf extends ShimCometConf { "of once per task. Tasks that start together each fetch it, so the saving grows " + "with the number of splits per file. The cache is never shared between stages or " + "stage attempts, so a query always reads the footers of its files afresh. A stage's " + - "cache holds up to 50 MiB of decoded footers and page indexes for each task the " + - "executor runs at once, outside Spark's memory accounting, and up to two such " + - "caches are kept after their tasks finish, for the next task of their stage. When " + - "disabled, each task keeps its own 50 MiB cache, which is freed when the task ends.") + "cache holds up to 50 MiB of decoded footers and page indexes for each of its " + + "running tasks, outside Spark's memory accounting. Up to two caches are kept after " + + "their tasks finish, for the next task of their stage, and each of those holds up " + + "to 50 MiB. When disabled, each task keeps its own 50 MiB cache, which is freed " + + "when the task ends.") .booleanConf .createWithDefault(true) From d80bb2195aa30b0f8615597d6d6696c1c44280d1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 23 Oct 2026 00:16:26 -0600 Subject: [PATCH 7/9] fix: free retained shared metadata caches at plugin shutdown Clear the registry in release_runtime, as clear_file_io_cache does for the Iceberg FileIO cache, and document when a retained cache is freed. --- .../contributor-guide/memory_management.md | 6 ++- native/core/src/execution/jni_api.rs | 4 +- .../core/src/parquet/shared_metadata_cache.rs | 43 ++++++++++++++++--- .../scala/org/apache/comet/CometConf.scala | 3 +- 4 files changed, 45 insertions(+), 11 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 51d8ae92da0..fc7a87e5a00 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -513,8 +513,10 @@ diverge for several structural reasons: `spark.comet.parquet.metadataCache.sharedWithinStage.enabled`, the default, the tasks of a stage share one cache per executor, which holds up to 50 MiB for each of its running tasks. The two most recently used caches are kept after their tasks finish, for the next task of their stage, - and shrink to 50 MiB each, so up to 100 MiB stays resident with no task running. With the - setting off, each task has its own 50 MiB cache, freed with the task. + and shrink to 50 MiB each, so up to 100 MiB stays resident with no task running. A kept cache is + freed once tasks of two other stages run on the executor, or when the plugin shuts down. An + executor is not told when a stage ends, so one that goes idle after a query holds its last two + caches until then. With the setting off, each task has its own 50 MiB cache, freed with the task. - **Batches in flight across the FFI boundary.** Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 6a3638bccbb..ab0c5392005 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -450,7 +450,8 @@ pub fn get_runtime() -> Handle { .clone() } -/// Tears down the global Tokio runtime, if it has been initialized. +/// Tears down the global Tokio runtime, if it has been initialized, and frees the process-wide +/// caches that outlive tasks: the Iceberg `FileIO`s and the retained Parquet metadata caches. /// /// The runtime is moved out of the global slot, so the next [`init_runtime`] or /// [`get_runtime`] call builds a new one. Shutting it down blocks the calling (JNI) thread @@ -463,6 +464,7 @@ pub fn get_runtime() -> Handle { /// would deadlock/panic. pub fn release_runtime() { crate::execution::operators::clear_file_io_cache(); + crate::parquet::shared_metadata_cache::clear_shared_metadata_caches(); let runtime = TOKIO_RUNTIME.lock().take(); if let Some(runtime) = runtime { runtime.shutdown_timeout(Duration::from_secs(3)); diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs index 4ef7cb9c545..8d2e0c2a3cf 100644 --- a/native/core/src/parquet/shared_metadata_cache.rs +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -37,7 +37,11 @@ //! tasks own it through their `RuntimeEnv`. It additionally holds strong references to the last //! [`MAX_IDLE_CACHES`] caches it handed out. Those matter when an executor runs one task at a //! time: the next task of the stage starts after the previous one released its plan. Whatever the -//! tasks do, at most [`MAX_IDLE_CACHES`] caches outlive the tasks that used them. +//! tasks do, at most [`MAX_IDLE_CACHES`] caches outlive the tasks that used them. A retained cache +//! is freed when tasks of [`MAX_IDLE_CACHES`] other stages have asked for theirs since, or when +//! the plugin shuts down and [`clear_shared_metadata_caches`] runs. Nothing frees it sooner: an +//! executor is not told that a stage has finished, so one that goes idle after a query keeps the +//! caches of its last stages until it shuts down. //! //! A stage's cache holds DataFusion's per-task limit for each of its tasks that is running, and one //! task's limit once none is, so that sharing does not shrink the room each task had with a cache @@ -65,7 +69,7 @@ use datafusion::execution::cache::default_cache::DefaultCache; use datafusion::execution::cache::{Cache, CacheEntryInfo}; use object_store::path::Path; use std::collections::{HashMap, VecDeque}; -use std::sync::{Arc, Mutex, OnceLock, Weak}; +use std::sync::{Arc, LazyLock, Mutex, Weak}; use std::time::Duration; /// How many caches of stages with no running task the registry holds on to. @@ -227,12 +231,26 @@ impl SharedMetadataCaches { Arc::new(TaskCache::new(stage)) } + /// Forgets every stage, freeing the retained caches. Running tasks keep the caches they hold. + fn clear(&self) { + let dropped = std::mem::take(&mut *self.registry.lock().unwrap()); + drop(dropped); + } + #[cfg(test)] fn len(&self) -> usize { self.registry.lock().unwrap().alive.len() } } +/// The registry of the executor's stages. +static CACHES: LazyLock = LazyLock::new(SharedMetadataCaches::default); + +/// Frees the caches retained for stages with no running task, as the plugin shuts down. +pub fn clear_shared_metadata_caches() { + CACHES.clear(); +} + /// A view of a [`FileMetadataCache`] that keys every entry by the object store as well as the /// path within it. /// @@ -336,8 +354,6 @@ pub(crate) fn scope_to_store( pub(crate) fn shared_metadata_cache_for_task( spark_config: &HashMap, ) -> Option> { - static CACHES: OnceLock = OnceLock::new(); - if !spark_config.get_bool(COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED) { return None; } @@ -347,9 +363,7 @@ pub(crate) fn shared_metadata_cache_for_task( stage_id: parse(COMET_STAGE_ID)?, stage_attempt: parse(COMET_STAGE_ATTEMPT)?, }; - let cache: Arc = CACHES - .get_or_init(SharedMetadataCaches::default) - .cache_for(key); + let cache: Arc = CACHES.cache_for(key); Some(cache) } @@ -433,6 +447,21 @@ mod tests { assert!(same_stage(&running, &caches.cache_for(key(0, 0)))); } + #[test] + fn clearing_frees_the_retained_caches_but_not_a_running_one() { + let caches = SharedMetadataCaches::default(); + let running = caches.cache_for(key(0, 0)); + let idle = Arc::downgrade(&caches.cache_for(key(1, 0)).stage); + assert!(idle.upgrade().is_some()); + + caches.clear(); + assert!(idle.upgrade().is_none()); + assert_eq!(caches.len(), 0); + assert_eq!(running.cache_limit(), DEFAULT_METADATA_CACHE_LIMIT); + drop(running); + assert_eq!(caches.len(), 0); + } + #[test] fn a_stage_cache_holds_one_task_limit_per_running_task() { let caches = SharedMetadataCaches::default(); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 7e98eb1ea24..6916b335923 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -203,7 +203,8 @@ object CometConf extends ShimCometConf { "cache holds up to 50 MiB of decoded footers and page indexes for each of its " + "running tasks, outside Spark's memory accounting. Up to two caches are kept after " + "their tasks finish, for the next task of their stage, and each of those holds up " + - "to 50 MiB. When disabled, each task keeps its own 50 MiB cache, which is freed " + + "to 50 MiB until tasks of two other stages run on the executor or the executor " + + "shuts down. When disabled, each task keeps its own 50 MiB cache, which is freed " + "when the task ends.") .booleanConf .createWithDefault(true) From 8fe1aa2db431358b0c1973fa445595129605d0a7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 23 Oct 2026 00:33:09 -0600 Subject: [PATCH 8/9] feat: free a retained shared metadata cache after 10 seconds idle An executor is not told when a stage ends, so a retained cache stayed until two other stages asked for one. A task on the native runtime now frees each retained cache once no task of its stage has run for 10 seconds, and stops when none is retained. --- .../contributor-guide/memory_management.md | 8 +- .../core/src/parquet/shared_metadata_cache.rs | 155 ++++++++++++++++-- .../scala/org/apache/comet/CometConf.scala | 6 +- 3 files changed, 151 insertions(+), 18 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index fc7a87e5a00..3d6558a9700 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -513,10 +513,10 @@ diverge for several structural reasons: `spark.comet.parquet.metadataCache.sharedWithinStage.enabled`, the default, the tasks of a stage share one cache per executor, which holds up to 50 MiB for each of its running tasks. The two most recently used caches are kept after their tasks finish, for the next task of their stage, - and shrink to 50 MiB each, so up to 100 MiB stays resident with no task running. A kept cache is - freed once tasks of two other stages run on the executor, or when the plugin shuts down. An - executor is not told when a stage ends, so one that goes idle after a query holds its last two - caches until then. With the setting off, each task has its own 50 MiB cache, freed with the task. + and shrink to 50 MiB each, so up to 100 MiB stays resident with no task running. An executor is + not told when a stage ends, so a kept cache is freed once no task of its stage has run for 10 + seconds, once tasks of two other stages run on the executor, or when the plugin shuts down. With + the setting off, each task has its own 50 MiB cache, freed with the task. - **Batches in flight across the FFI boundary.** Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs index 8d2e0c2a3cf..839cd6c74bd 100644 --- a/native/core/src/parquet/shared_metadata_cache.rs +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -37,11 +37,11 @@ //! tasks own it through their `RuntimeEnv`. It additionally holds strong references to the last //! [`MAX_IDLE_CACHES`] caches it handed out. Those matter when an executor runs one task at a //! time: the next task of the stage starts after the previous one released its plan. Whatever the -//! tasks do, at most [`MAX_IDLE_CACHES`] caches outlive the tasks that used them. A retained cache -//! is freed when tasks of [`MAX_IDLE_CACHES`] other stages have asked for theirs since, or when -//! the plugin shuts down and [`clear_shared_metadata_caches`] runs. Nothing frees it sooner: an -//! executor is not told that a stage has finished, so one that goes idle after a query keeps the -//! caches of its last stages until it shuts down. +//! tasks do, at most [`MAX_IDLE_CACHES`] caches outlive the tasks that used them. An executor is +//! not told that a stage has finished, so a retained cache is freed once no task of its stage has +//! run for [`IDLE_CACHE_TIMEOUT`], once tasks of [`MAX_IDLE_CACHES`] other stages have asked for +//! theirs, or when the plugin shuts down and [`clear_shared_metadata_caches`] runs, whichever comes +//! first. An executor that goes idle after a query frees them within the timeout. //! //! A stage's cache holds DataFusion's per-task limit for each of its tasks that is running, and one //! task's limit once none is, so that sharing does not shrink the room each task had with a cache @@ -56,6 +56,7 @@ //! the file's path within its object store, which drops the bucket, so [`scope_to_store`] gives //! each scan a view of the cache that also keys by the scan's object store. +use crate::execution::jni_api::get_runtime; use crate::execution::spark_config::{ SparkConfig, COMET_APP_ID, COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED, COMET_STAGE_ATTEMPT, COMET_STAGE_ID, @@ -69,12 +70,17 @@ use datafusion::execution::cache::default_cache::DefaultCache; use datafusion::execution::cache::{Cache, CacheEntryInfo}; use object_store::path::Path; use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, LazyLock, Mutex, Weak}; -use std::time::Duration; +use std::time::{Duration, Instant}; /// How many caches of stages with no running task the registry holds on to. const MAX_IDLE_CACHES: usize = 2; +/// How long the registry holds on to the cache of a stage with no running task. The next task of a +/// stage that runs one task at a time on an executor starts well within it. +const IDLE_CACHE_TIMEOUT: Duration = Duration::from_secs(10); + #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(crate) struct StageKey { pub app_id: String, @@ -82,10 +88,17 @@ pub(crate) struct StageKey { pub stage_attempt: i32, } +/// How many tasks of a stage are running, and since when none has been. +#[derive(Default)] +struct Running { + tasks: usize, + idle_since: Option, +} + /// The cache of one stage, and how many of its tasks are running. struct StageCache { cache: Arc>, - running: Mutex, + running: Mutex, } impl StageCache { @@ -95,7 +108,7 @@ impl StageCache { DefaultCache::new(DEFAULT_METADATA_CACHE_LIMIT) .with_name("SharedFileMetadataCache"), ), - running: Mutex::new(0), + running: Mutex::new(Running::default()), } } @@ -103,9 +116,15 @@ impl StageCache { /// evicts the least recently used entries down to it. fn update_running(&self, update: impl FnOnce(usize) -> usize) { let mut running = self.running.lock().unwrap(); - *running = update(*running); + running.tasks = update(running.tasks); + running.idle_since = (running.tasks == 0).then(Instant::now); self.cache - .update_cache_limit(DEFAULT_METADATA_CACHE_LIMIT.saturating_mul((*running).max(1))); + .update_cache_limit(DEFAULT_METADATA_CACHE_LIMIT.saturating_mul(running.tasks.max(1))); + } + + /// When the cache has gone unused for `timeout`, or `None` while a task of the stage runs. + fn expires_at(&self, timeout: Duration) -> Option { + Some(self.running.lock().unwrap().idle_since? + timeout) } } @@ -138,7 +157,7 @@ impl Drop for TaskCache { impl std::fmt::Debug for TaskCache { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("TaskCache") - .field("running", &*self.stage.running.lock().unwrap()) + .field("running", &self.stage.running.lock().unwrap().tasks) .finish() } } @@ -208,6 +227,8 @@ struct Registry { #[derive(Default)] pub(crate) struct SharedMetadataCaches { registry: Mutex, + /// Whether a task that frees the idle caches is running. + expiring: AtomicBool, } impl SharedMetadataCaches { @@ -231,10 +252,60 @@ impl SharedMetadataCaches { Arc::new(TaskCache::new(stage)) } + /// Frees the retained caches that have gone unused for `timeout` at `now`, and returns when + /// the next of the remaining ones would, or `None` when none is retained. + fn expire_idle(&self, now: Instant, timeout: Duration) -> Option { + let mut registry = self.registry.lock().unwrap(); + let mut expired = Vec::new(); + let mut next: Option = None; + registry + .retained + .retain(|(_, stage)| match stage.expires_at(timeout) { + Some(at) if at <= now => { + expired.push(Arc::clone(stage)); + false + } + at => { + // A stage with a running task is checked again a whole timeout later. + let at = at.unwrap_or(now + timeout); + next = Some(next.map_or(at, |next| next.min(at))); + true + } + }); + drop(expired); + registry.alive.retain(|_, weak| weak.strong_count() > 0); + next + } + + /// Starts a task on `runtime` that frees retained caches as they time out, unless one is + /// running. It stops once no cache is retained, and the next request for a cache starts another. + fn expire_idle_on(&'static self, runtime: &tokio::runtime::Handle, timeout: Duration) { + if self.expiring.swap(true, Ordering::AcqRel) { + return; + } + runtime.spawn(async move { + loop { + while let Some(next) = self.expire_idle(Instant::now(), timeout) { + tokio::time::sleep_until(next.into()).await; + } + self.expiring.store(false, Ordering::Release); + // A cache retained after the last check found this task still running, so carry + // on unless no cache is retained or another task has started since. + if self.registry.lock().unwrap().retained.is_empty() + || self.expiring.swap(true, Ordering::AcqRel) + { + break; + } + } + }); + } + /// Forgets every stage, freeing the retained caches. Running tasks keep the caches they hold. fn clear(&self) { let dropped = std::mem::take(&mut *self.registry.lock().unwrap()); drop(dropped); + // The runtime that ran the expiring task is shut down along with the caches. + self.expiring.store(false, Ordering::Release); } #[cfg(test)] @@ -364,6 +435,7 @@ pub(crate) fn shared_metadata_cache_for_task( stage_attempt: parse(COMET_STAGE_ATTEMPT)?, }; let cache: Arc = CACHES.cache_for(key); + CACHES.expire_idle_on(&get_runtime(), IDLE_CACHE_TIMEOUT); Some(cache) } @@ -447,6 +519,67 @@ mod tests { assert!(same_stage(&running, &caches.cache_for(key(0, 0)))); } + #[test] + fn an_idle_cache_is_freed_after_the_timeout() { + let caches = SharedMetadataCaches::default(); + let running = caches.cache_for(key(0, 0)); + let idle = Arc::downgrade(&caches.cache_for(key(1, 0)).stage); + let idle_since = Instant::now(); + let timeout = Duration::from_secs(10); + + // Not yet timed out; the running stage is checked a timeout from now. + let next = caches.expire_idle(idle_since, timeout).unwrap(); + assert!(idle.upgrade().is_some()); + assert!(next <= idle_since + timeout + Duration::from_secs(1)); + + let later = idle_since + timeout + Duration::from_secs(1); + assert!(caches.expire_idle(later, timeout).is_some()); + assert!(idle.upgrade().is_none()); + assert_eq!(caches.len(), 1); + // A stage with a running task is never timed out. + assert!(same_stage(&running, &caches.cache_for(key(0, 0)))); + + drop(running); + let much_later = later + 2 * timeout; + assert!(caches.expire_idle(much_later, timeout).is_none()); + assert_eq!(caches.len(), 0); + } + + #[test] + fn a_task_of_the_stage_resets_its_timeout() { + let caches = SharedMetadataCaches::default(); + drop(caches.cache_for(key(1, 0))); + let timeout = Duration::from_secs(10); + let start = Instant::now(); + std::thread::sleep(Duration::from_millis(20)); + // The next task of the stage runs and finishes, so the cache went idle after `start`. + drop(caches.cache_for(key(1, 0))); + assert!(caches + .expire_idle(start + timeout, timeout) + .is_some_and(|next| next > start + timeout)); + assert_eq!(caches.len(), 1); + } + + #[test] + fn the_expiring_task_frees_idle_caches_and_stops() { + static CACHES: LazyLock = + LazyLock::new(SharedMetadataCaches::default); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + let idle = Arc::downgrade(&CACHES.cache_for(key(1, 0)).stage); + let timeout = Duration::from_millis(50); + CACHES.expire_idle_on(runtime.handle(), timeout); + runtime.block_on(async { + while CACHES.expiring.load(Ordering::Acquire) { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }); + assert!(idle.upgrade().is_none()); + assert_eq!(CACHES.len(), 0); + } + #[test] fn clearing_frees_the_retained_caches_but_not_a_running_one() { let caches = SharedMetadataCaches::default(); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6916b335923..3aad4330df5 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -203,9 +203,9 @@ object CometConf extends ShimCometConf { "cache holds up to 50 MiB of decoded footers and page indexes for each of its " + "running tasks, outside Spark's memory accounting. Up to two caches are kept after " + "their tasks finish, for the next task of their stage, and each of those holds up " + - "to 50 MiB until tasks of two other stages run on the executor or the executor " + - "shuts down. When disabled, each task keeps its own 50 MiB cache, which is freed " + - "when the task ends.") + "to 50 MiB until no task of its stage has run for 10 seconds, tasks of two other " + + "stages run on the executor, or the executor shuts down. When disabled, each task " + + "keeps its own 50 MiB cache, which is freed when the task ends.") .booleanConf .createWithDefault(true) From 2f43ec3ad78ca9f788866ade23b84c7f5c4294c0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 23 Oct 2026 11:53:37 +0100 Subject: [PATCH 9/9] test: a runtime started after a plugin shutdown frees idle shared metadata caches --- .../core/src/parquet/shared_metadata_cache.rs | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/native/core/src/parquet/shared_metadata_cache.rs b/native/core/src/parquet/shared_metadata_cache.rs index 839cd6c74bd..b078edf865e 100644 --- a/native/core/src/parquet/shared_metadata_cache.rs +++ b/native/core/src/parquet/shared_metadata_cache.rs @@ -580,6 +580,45 @@ mod tests { assert_eq!(CACHES.len(), 0); } + #[test] + fn a_runtime_started_after_a_shutdown_frees_idle_caches() { + static CACHES: LazyLock = + LazyLock::new(SharedMetadataCaches::default); + let timeout = Duration::from_millis(50); + let runtime = || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + }; + + // The plugin shuts down before the first context's timer has run, which drops the runtime + // the timer was spawned on. + let first = runtime(); + drop(CACHES.cache_for(key(1, 0))); + CACHES.expire_idle_on(first.handle(), timeout); + CACHES.clear(); + drop(first); + + // A context started later in the same JVM gets a new runtime. Its tasks leave a cache + // behind, and the timer has to start again on the new runtime to free it. + let second = runtime(); + let idle = Arc::downgrade(&CACHES.cache_for(key(1, 0)).stage); + CACHES.expire_idle_on(second.handle(), timeout); + second + .block_on(async { + tokio::time::timeout(Duration::from_secs(10), async { + while CACHES.expiring.load(Ordering::Acquire) { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + }) + .expect("the timer should stop once it has freed the idle cache"); + assert!(idle.upgrade().is_none()); + assert_eq!(CACHES.len(), 0); + } + #[test] fn clearing_frees_the_retained_caches_but_not_a_running_one() { let caches = SharedMetadataCaches::default();