andygrove commented on code in PR #7082: URL: https://github.com/apache/datafusion-comet/pull/7082#discussion_r4345444712
########## 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: No, it wasn't meant to, and you're right that it shrinks every task's headroom when the tasks read different files. A stage's cache now holds `DEFAULT_METADATA_CACHE_LIMIT` for each of its running tasks. Each task takes a `TaskCache` handle when it builds its `RuntimeEnv`, which counts it in, and drops it with its plan, which counts it out. So each running task keeps the 50 MiB it had with a cache of its own, and the footers the deletion vector planning loads are no more likely to be evicted before the scan opens them than before. A file the tasks share is stored once, so the tasks of a stage hold no more than they did. When the last task finishes, the cache shrinks back to 50 MiB, evicting the least recently used entries. So each of the two caches the registry keeps holds no more than a task's own cache would, at most 100 MiB with no task running. (40beedd74) An earlier version scaled the limit by `spark.executor.cores / spark.task.cpus`, but sunchao pointed out that this is only an estimate on a standalone executor, so the limit now follows the tasks that actually run. ########## native/core/src/parquet/parquet_exec.rs: ########## @@ -2111,4 +2118,131 @@ mod tests { ); } } + + /// A parquet file of one `a` column holding `values`. + fn parquet_of(values: Vec<i32>) -> 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<Bytes> = 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] { Review Comment: Done. Without sharing, both scans now run in one `RuntimeEnv` with its own cache, as two scans in one plan would. The registry and the stage cache are only built for the shared iteration. I checked that each half needs the scoping. With `scope_to_store` returning the cache unscoped, the shared iteration fails, and so does the unshared one when it runs first. ########## 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 " + Review Comment: Updated. The doc now says that tasks which start together each fetch the footer, so the saving grows with the number of splits per file. It also says what the cache holds: up to 50 MiB of decoded footers and page indexes for each of the stage's running tasks, outside Spark's memory accounting. Up to two caches are kept after their tasks finish, each holding up to 50 MiB. With the setting off, each task's 50 MiB cache is freed when the task ends. ########## 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) + .with_name("SharedFileMetadataCache"), + ) + }); + 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.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<FileMetadataCache>, +} + +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<Path, CachedFileMetadataEntry> for StoreScopedCache { + fn get(&self, key: &Path) -> Option<CachedFileMetadataEntry> { + self.inner.get(&self.scoped(key)) + } + + fn put(&self, key: &Path, value: CachedFileMetadataEntry) -> Option<CachedFileMetadataEntry> { + self.inner.put(&self.scoped(key), value) + } + + fn remove(&self, key: &Path) -> Option<CachedFileMetadataEntry> { + 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<Duration> { + self.inner.cache_ttl() + } + + fn update_cache_ttl(&self, ttl: Option<Duration>) { + 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<Path, CacheEntryInfo<CachedFileMetadataEntry>> { + 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<FileMetadataCache>, + object_store_url: &ObjectStoreUrl, +) -> Arc<FileMetadataCache> { + 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( + spark_config: &HashMap<String, String>, +) -> Option<Arc<FileMetadataCache>> { + static CACHES: OnceLock<SharedMetadataCaches> = OnceLock::new(); + + if !spark_config.get_bool(COMET_PARQUET_SHARED_METADATA_CACHE_ENABLED) { + return None; + } + let parse = |name: &str| spark_config.get(name)?.parse::<i32>().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)?, + }; + Some( + CACHES + .get_or_init(SharedMetadataCaches::default) + .cache_for(key), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(stage_id: i32, stage_attempt: i32) -> StageKey { + StageKey { + app_id: "app-1".to_string(), + 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)); + 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)); + 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)); + } + + #[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 address = Arc::as_ptr(&first) as *const () as usize; + drop(first); + let next = caches.cache_for(key(1, 0)); + 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() { Review Comment: Dropped, thanks. `idle_caches_are_bounded_but_running_ones_are_not_dropped` covers it. ########## 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) Review Comment: Fair question. I measured it on one Linux box with Spark 4.1 in `local-cluster` mode, with a release build, at the default 128 MB `maxPartitionBytes`. The query was `sum(l_quantity), sum(l_extendedprice)` over TPC-H lineitem, three runs per setting. | Data | Executors x cores | Tasks | Footer reads, sharing on | Footer reads, sharing off | |---|---|---|---|---| | SF100, 16 files of 1.47 GB | 2 x 8 | 176 | 56 to 67 (62 to 68% hits) | 176 | | SF100, 16 files of 1.47 GB | 4 x 4 | 176 | 79 to 106 (40 to 55% hits) | 176 | | SF10, one file of 2.4 GB | 2 x 8 and 4 x 4 | 19 | 16 (3 hits) | 19 | So you're right about tasks that start together. A file with few splits gets almost nothing, since its tasks run in one wave and each misses. With files of a size common in data lakes, each split into about 11 tasks, and 176 tasks over 16 slots, the later waves find the footer and footer reads drop by 40 to 68%. More executors means fewer hits, as you predicted. Wall time on this box doesn't move, about 0.8 s either way, because the data is on local NVMe and a footer read costs very little there. The point of the change is the object store case, where each footer and page index read is one or more GET requests with tens of milliseconds of latency. Removing 40 to 68% of those requests from the start of the tasks is the saving, and local disk can't show it. So I'd like to keep it on by default. With the limit now sized by the stage's running tasks, as in the thread on `shared_metadata_cache.rs`, turning it on no longer takes headroom away from any scan. ########## 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: Good catch, thanks. The handle a task's `CacheManager` gets, `TaskCache`, now ignores `update_cache_limit`, so the `metadata_cache_limit` that `CacheManager::try_new` applies has no effect on the shared cache. The registry alone sets the limit, from the stage's running tasks. `a_task_runtime_env_does_not_reset_the_stage_limit` builds a `RuntimeEnv` with the default config over one task's handle while another task is running, and checks that the limit is still two tasks' worth. It then drops the other task and checks that the limit falls to one task's. (40beedd74) -- 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]
