viirya commented on code in PR #7066:
URL: https://github.com/apache/datafusion-comet/pull/7066#discussion_r4300504607
##########
native/core/src/execution/jni_api.rs:
##########
@@ -990,6 +990,44 @@ fn configure_skip_partial_aggregation(config: &mut
SessionConfig, plan: &Operato
}
}
+/// The key of a shared [`DiskManager`]: the local directories and the temp
directory size limit.
+type DiskManagerKey = (Vec<String>, u64);
+
+static DISK_MANAGERS: OnceLock<Mutex<HashMap<DiskManagerKey,
Arc<DiskManager>>>> = OnceLock::new();
+
+/// Returns the [`DiskManager`] shared by every native plan in this process
that uses the same
+/// local directories and size limit.
+///
+/// A `DiskManager` in `Directories` mode creates a `datafusion-*` directory
in every local
+/// directory when it is built, and removes it recursively when it is dropped.
Building one per
+/// plan costs a `stat` and a `mkdir` per local disk when the plan is created,
and a recursive
+/// delete when it is released, whether or not the plan ever spills (#7060).
Sharing one makes
+/// that a once-per-executor cost. Spill files are still deleted when their
last reference is
+/// dropped, so a long-lived `DiskManager` does not accumulate them.
+///
+/// The size limit then covers all plans sharing the manager, not each plan
separately.
+fn shared_disk_manager(
+ local_dirs: Vec<String>,
+ max_temp_directory_size: u64,
+) -> CometResult<Arc<DiskManager>> {
+ let mut managers = DISK_MANAGERS
+ .get_or_init(|| Mutex::new(HashMap::new()))
+ .lock();
+ let key = (local_dirs, max_temp_directory_size);
+ if let Some(disk_manager) = managers.get(&key) {
+ return Ok(Arc::clone(disk_manager));
+ }
Review Comment:
I think keying on the limit is the safer of the two here.
`set_max_temp_directory_size` stores into a shared `AtomicU64`, and
`FileSpillWriter::write` reads `max_temp_directory_size()` again on every
write. If the cache were keyed on the dirs only and each plan set its own limit
on the cached manager, the last plan created would set the limit for every plan
already spilling in that executor. A session running with a 500 GB limit could
hit a 10 GB cap in the middle of a spill because another session just started a
plan.
I'm also not sure this is limited to local-mode tests. The value comes from
`SQLConf.get`, which on an executor reads the session confs that travel with
the task. A Thrift server or Spark Connect app runs many sessions in one
application, so different sessions can set different values and send tasks to
the same executor. That is why I suggested in the comment on line 1016 that the
docs say the limit is shared by plans that use the same value, rather than
changing the keying.
--
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]