Include key into memory estimation

Author: mkleen

datafusion/common/Cargo.toml

@@ -75,7 +75,7 @@ indexmap = { workspace = true }
 itertools = { workspace = true }
 libc = "0.2.180"
 log = { workspace = true }
-object_store = { workspace = true, optional = true }
+object_store = { workspace = true, optional = true, default-features = true }
 parquet = { workspace = true, optional = true, default-features = true }
 paste = { workspace = true }
 recursive = { workspace = true, optional = true }

datafusion/common/src/heap_size.rs

@@ -29,6 +29,7 @@ use half::f16;
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::sync::Arc;
+use object_store::path::Path;
 
 /// This is a temporary solution until <https://github.com/apache/datafusion/pull/19599> and
 /// <https://github.com/apache/arrow-rs/pull/9138> are resolved.
@@ -179,6 +180,12 @@ impl DFHeapSize for DataType {
     }
 }
 
+impl DFHeapSize for Path {
+    fn heap_size(&self) -> usize {
+      self.as_ref().heap_size()
+    }
+}
+
 impl<T: DFHeapSize> DFHeapSize for Vec<T> {
     fn heap_size(&self) -> usize {
         let item_size = size_of::<T>();

datafusion/execution/src/cache/cache_unit.rs

@@ -97,16 +97,17 @@ impl DefaultFileStatisticsCacheState {
         key: &Path,
         value: CachedFileMetadata,
     ) -> Option<CachedFileMetadata> {
+        let key_size = key.heap_size();
         let entry_size = value.heap_size();
 
-        if entry_size > self.memory_limit {
+        if entry_size + key_size > self.memory_limit {
             // Remove stale entry if exists
             self.remove(key);
             return None;
         }
 
         let old_value = self.lru_queue.put(key.clone(), value);
-        self.memory_used += entry_size;
+        self.memory_used += entry_size + key_size;
 
         if let Some(old_entry) = &old_value {
             self.memory_used -= old_entry.heap_size();
@@ -119,6 +120,7 @@ impl DefaultFileStatisticsCacheState {
 
     fn remove(&mut self, k: &Path) -> Option<CachedFileMetadata> {
         if let Some(old_entry) = self.lru_queue.remove(k) {
+            self.memory_used -= k.heap_size();
             self.memory_used -= old_entry.heap_size();
             Some(old_entry)
         } else {
@@ -142,6 +144,7 @@ impl DefaultFileStatisticsCacheState {
     fn evict_entries(&mut self) {
         while self.memory_used > self.memory_limit {
             if let Some(removed) = self.lru_queue.pop() {
+                self.memory_used -= removed.0.heap_size();
                 self.memory_used -= removed.1.heap_size();
             } else {
                 // cache is empty while memory_used > memory_limit, cannot happen
@@ -547,7 +550,7 @@ mod tests {
         let (meta_2, value_2) = create_cached_file_metadata_with_stats("test2.parquet");
         let (meta_3, value_3) = create_cached_file_metadata_with_stats("test3.parquet");
 
-        let limit_for_2_entries = value_1.heap_size() + value_2.heap_size();
+        let limit_for_2_entries = &meta_1.location.heap_size() + value_1.heap_size() +  &meta_2.location.heap_size() + value_2.heap_size();
 
         // create a cache with a limit which fits exactly 2 entries
         let cache = DefaultFileStatisticsCache::new(limit_for_2_entries);
@@ -579,7 +582,7 @@ mod tests {
 
         cache.remove(&meta_2.location);
         assert_eq!(cache.len(), 1);
-        assert_eq!(cache.memory_used(), value_3.heap_size());
+        assert_eq!(cache.memory_used(),  &meta_3.location.heap_size()  + value_3.heap_size());
 
         cache.clear();
         assert_eq!(cache.len(), 0);