andygrove commented on code in PR #7082:
URL: https://github.com/apache/datafusion-comet/pull/7082#discussion_r4308772382
##########
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. DataFusion keys the cache by `ObjectMeta.location`, and the
location of a scanned file has no bucket, so equal-size files at the same path
in two stores passed the size check and shared a footer.
Each scan now looks up footers through a view of the cache that also keys by
the scan's registered object-store URL (`scope_to_store`). That URL already
carries the backend and a hash of the store config, so I didn't add a separate
identity. It's applied at `parquet_exec.rs:230` and in the Delta
deletion-vector footer read, which gets the scan's URL through
`attach_access_plans` so both find the same entries. The view is used whether
sharing is on or off.
The regression is `equal_paths_in_two_stores_do_not_share_footers`. It puts
equal-size files holding `[1,2,3]` and `[101,102,103]` at the same path in two
stores and scans each with `a > 100`, with the cache shared and not shared. It
fails without the scoping (`[[], []]`). It runs two sequential scans rather
than a UNION. Is that enough for what you had in mind?
##########
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:
You're right. The bound was only enforced on the next `cache_for` call, so
dropping 16 handles left 16 caches pinned.
I took the weak-reference model you suggested. The registry now keeps a
`Weak` for every cache, which the running tasks own through their `RuntimeEnv`,
plus a queue of strong references to the last two it handed out, refreshed on
acquisition. Expired weak entries are pruned on each request, and the
`strong_count` arithmetic is gone. At most `MAX_IDLE_CACHES` caches outlive
their tasks, and sequential tasks of one stage still get hits.
I added your test as `idle_retention_is_bounded_after_all_tasks_finish`,
plus one that checks dropped caches are forgotten on the next request.
##########
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:
Agreed, stage ids restart at zero in every `SparkContext` and the registry
outlives a context stopped in the same JVM.
`serializeCometSQLConfs` now sends `spark.app.id` with the stage fields, and
it's part of the key. If there's no app id, sharing is off for that task rather
than guessing.
I ran the restart scenario. `CometSharedMetadataCacheRestartSuite` stops and
recreates the local context, replaces an equal-size file at the same path, and
counts rows with `id >= 1000`. With the app id taken out of the key it fails
with `0 did not equal 1000`, because the stale footer prunes every row. With
it, the suite passes.
--
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]
