andygrove commented on code in PR #7082: URL: https://github.com/apache/datafusion-comet/pull/7082#discussion_r4481473694
########## native/core/src/parquet/shared_metadata_cache.rs: ########## @@ -0,0 +1,324 @@ +// 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. 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 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_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::{ + CachedFileMetadataEntry, FileMetadataCache, DEFAULT_METADATA_CACHE_LIMIT, +}; +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::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, PartialEq, Eq, Hash)] +pub(crate) struct StageKey { + pub app_id: String, + pub stage_id: i32, + pub stage_attempt: i32, +} + +#[derive(Default)] +struct Registry { + /// Every cache that is still alive, running or retained. + alive: HashMap<StageKey, Weak<FileMetadataCache>>, + /// The caches kept alive for the next task of their stage, least recently requested first. + retained: VecDeque<(StageKey, Arc<FileMetadataCache>)>, +} + +#[derive(Default)] +pub(crate) struct SharedMetadataCaches { + registry: Mutex<Registry>, +} + +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<FileMetadataCache> { + 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) Review Comment: You're right, the `peak` limit covers the gap between tasks but not the next task's own loads. A running task's preloaded footers are older than the finished sibling's entries, so they go first. I reproduced your numbers: when the stage's tasks load more than the limit holds, the preloaded footers are the ones refetched. I'd treat it as an accepted cost. It is a refetch of a footer a per-task cache would have kept, it needs the tasks of a stage to load more than the limit holds, and it never changes results. Passing the preloaded footers straight to the file open would remove it, but it means a second path through `delta_dv.rs` for a case that only costs a read. The module doc now says so. (commit below) I also reran the fork's `rust-test` job that ran out of disk. ########## native/core/src/parquet/shared_metadata_cache.rs: ########## @@ -0,0 +1,324 @@ +// 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. 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 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_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::{ + CachedFileMetadataEntry, FileMetadataCache, DEFAULT_METADATA_CACHE_LIMIT, +}; +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::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, PartialEq, Eq, Hash)] +pub(crate) struct StageKey { + pub app_id: String, + pub stage_id: i32, + pub stage_attempt: i32, +} + +#[derive(Default)] +struct Registry { + /// Every cache that is still alive, running or retained. + alive: HashMap<StageKey, Weak<FileMetadataCache>>, + /// The caches kept alive for the next task of their stage, least recently requested first. + retained: VecDeque<(StageKey, Arc<FileMetadataCache>)>, +} + +#[derive(Default)] +pub(crate) struct SharedMetadataCaches { + registry: Mutex<Registry>, +} + +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<FileMetadataCache> { + 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) Review Comment: You're right, the `peak` limit covers the gap between tasks but not the next task's own loads. The cache evicts the least recently used entry whoever loaded it, and a running task's preloaded footers are older than the finished sibling's entries, so they go first. By that reading, when the stage's tasks load more than the limit holds, the preloaded footers are the ones refetched. I haven't re-run your probe. I'd treat it as an accepted cost. It is a refetch of a footer a per-task cache would have kept, it needs the tasks of a stage to load more than the limit holds, and it never changes results. Passing the preloaded footers straight to the file open would remove it, but it means a second path through `delta_dv.rs` for a case that only costs a read. The module doc now says so (9947811dc4). On the `rust-test` failure: that was `No space left on device`. I reran it, and the push starts a new fork run. -- 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]
