rich7420 commented on code in PR #7082: URL: https://github.com/apache/datafusion-comet/pull/7082#discussion_r4305678604
########## native/core/src/parquet/shared_metadata_cache.rs: ########## @@ -0,0 +1,187 @@ +// 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 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. +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<FileMetadataCache>, +} + +#[derive(Default)] +pub(crate) struct SharedMetadataCaches { + /// Ordered from least to most recently requested. + entries: Mutex<Vec<Entry>>, +} + +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<FileMetadataCache> { + 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(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 } Review Comment: Could we make the idle bound hold when tasks release their caches, without waiting for another `cache_for` call? Holding 16 stage handles and then dropping all of them leaves 16 idle entries pinned indefinitely. Each can retain 50 MiB outside Spark's memory accounting. A smaller ownership model would be a weak-reference registry plus a bounded queue of strong references, refreshed on acquisition, with expired weak entries pruned. That keeps sequential-task reuse and removes the `strong_count` arithmetic. This regression can go in the existing test module; it fails with the current registry: ```rust #[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); } ``` ########## native/core/src/execution/jni_api.rs: ########## @@ -1009,6 +1011,14 @@ 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 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)), Review Comment: Could we scope this cache by the resolved object-store identity too? DataFusion keys it by `ObjectMeta.location`, but `get_partitioned_files` strips the bucket/authority and sets the modification time to the epoch. Tasks reading equal-size files at the same path in different buckets can therefore reuse each other's footer. I reproduced this with two stores containing `table/part-00000.parquet`, both 284 bytes: A contains `[1,2,3]`, B `[101,102,103]`. After warming A, B receives A's exact footer Arc and min/max `1..3`; without sharing B reports `101..103`. A filter can then prune valid rows. Please resolve the cache per scan using the existing store identity, including backend/config distinctions, at both `parquet_exec.rs:230` and the footer lookup in `delta_dv.rs:683`. Passing modification time alone does not establish identity. Could we add a native UNION regression over same-path, same-size files in two stores, comparing sharing on/off with a predicate that selects only B? ########## native/core/src/parquet/shared_metadata_cache.rs: ########## @@ -0,0 +1,187 @@ +// 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 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. +const MAX_IDLE_CACHES: usize = 2; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct StageKey { + pub stage_id: i32, + pub stage_attempt: i32, +} Review Comment: Could the namespace also include an application/context identity? Spark's `DAGScheduler.nextStageId` starts at zero for each SparkContext, while this native static registry survives a stop/start in the same local JVM. A new context can match an old retained stage/attempt and reuse stale metadata for the same path and size. Please send an application/context identifier alongside the stage fields in `serializeCometSQLConfs` and include it in this key. A regression should stop/recreate the local SparkContext, reuse the stage ID and an equal-size replaced file, and verify the new contents. This finding is source-traced; I have not run that full restart scenario. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
