andygrove commented on code in PR #7082: URL: https://github.com/apache/datafusion-comet/pull/7082#discussion_r4358762156
########## native/core/src/parquet/shared_metadata_cache.rs: ########## @@ -0,0 +1,634 @@ +// 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. 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 +//! 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 +//! 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, +}; +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::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, LazyLock, Mutex, Weak}; +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; Review Comment: I measured it. Spark 4.1 `local-cluster`, 16 files of about 6 MB read as about 12 splits each (192 file opens, 178 tasks), a debug build on a Mac, three runs per shape. I switched retention off with a temporary override of the retained-cache count, so the only difference between the arms is whether the next task of a stage finds the cache. | Executors x cores | Hits, retention on | Hits, retention off | | --- | --- | --- | | 2 x 1 | 160 | 0 | | 4 x 1 | 128 | 0 | | 8 x 1 | 64 to 95 | 0 | | 2 x 4 | 122 to 138 | 107 to 140 | | 4 x 4 | 80 to 104 | 78 to 85 | So your reading holds. With several slots the stage's other running tasks keep the cache alive and retention adds little. With single-core executors it is the whole benefit: no hits without it, and with it the hits fall as the executor count approaches the number of splits per file, from 83% of opens at 2 executors to about a third at 8. I'd keep it. I'd also keep the shutdown hook. The timer runs on the native runtime, so it stops when `release_runtime` shuts that runtime down. A cache retained in the 10 seconds before a shutdown would then stay in a JVM that lives on, as in local mode, and the hook is one line next to `clear_file_io_cache`, which is what mbutrovich asked for. Does that work for you? ########## native/core/src/execution/jni_api.rs: ########## @@ -461,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(); Review Comment: You're right that the second shutdown never reaches `release_runtime`, because `released` stays set. But then the native runtime is not shut down either, so the expiry task from 8fe1aa2db keeps running on it. The first shutdown clears the registry and resets the flag that says a timer is running. The second context's first task then starts a timer on the new runtime, and it frees the retained caches 10 seconds after that context's last task, with no shutdown call needed. Your probe has no runtime, so it can't show that. I added `a_runtime_started_after_a_shutdown_frees_idle_caches` (2f43ec3ad7). It clears the registry and drops the runtime the first timer was spawned on, as a shutdown does. It then starts a second runtime and checks that the timer starts again and frees the idle cache. It fails when `clear` doesn't reset that flag. It covers the registry and the timer, not a full `SparkContext` restart, which I didn't 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]
