make lfc table scoped

jizezhang · jizezhang · commit 1789f023e02e · 2026-01-02T16:19:37.000-08:00
diff --git a/datafusion-cli/src/functions.rs b/datafusion-cli/src/functions.rs
@@ -771,6 +771,7 @@ impl TableFunctionImpl for ListFilesCacheFunc {
             Field::new("metadata", DataType::Struct(nested_fields.clone()), true);
 
         let schema = Arc::new(Schema::new(vec![
+            Field::new("table", DataType::Utf8, false),
             Field::new("path", DataType::Utf8, false),
             Field::new("metadata_size_bytes", DataType::UInt64, false),
             // expires field in ListFilesEntry has type Instant when set, from which we cannot get "the number of seconds", hence using Duration instead of Timestamp as data type.
@@ -785,7 +786,8 @@ impl TableFunctionImpl for ListFilesCacheFunc {
                 true,
             ),
         ]));
-
+        
+        let mut table_arr = vec![];
         let mut path_arr = vec![];
         let mut metadata_size_bytes_arr = vec![];
         let mut expires_arr = vec![];
@@ -802,7 +804,8 @@ impl TableFunctionImpl for ListFilesCacheFunc {
             let mut current_offset: i32 = 0;
 
             for (path, entry) in list_files_cache.list_entries() {
-                path_arr.push(path.to_string());
+                table_arr.push(path.0.map_or("NULL".to_string(), |t| t.to_string()));
+                path_arr.push(path.1.to_string());
                 metadata_size_bytes_arr.push(entry.size_bytes as u64);
                 // calculates time left before entry expires
                 expires_arr.push(
@@ -841,6 +844,7 @@ impl TableFunctionImpl for ListFilesCacheFunc {
         let batch = RecordBatch::try_new(
             schema.clone(),
             vec![
+                Arc::new(StringArray::from(table_arr)),
                 Arc::new(StringArray::from(path_arr)),
                 Arc::new(UInt64Array::from(metadata_size_bytes_arr)),
                 Arc::new(DurationMillisecondArray::from(expires_arr)),
diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs
@@ -34,6 +34,7 @@ use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
 use datafusion_datasource::{
     ListingTableUrl, PartitionedFile, TableSchema, compute_all_files_statistics,
 };
+use datafusion_execution::cache::TableScopedPath;
 use datafusion_execution::cache::cache_manager::FileStatisticsCache;
 use datafusion_execution::cache::cache_unit::DefaultFileStatisticsCache;
 use datafusion_expr::dml::InsertOp;
@@ -565,7 +566,11 @@ impl TableProvider for ListingTable {
 
         // Invalidate cache entries for this table if they exist
         if let Some(lfc) = state.runtime_env().cache_manager.get_list_files_cache() {
-            let _ = lfc.remove(table_path.prefix());
+            let key = TableScopedPath(
+                table_path.get_table_ref().clone(),
+                table_path.prefix().clone(),
+            );
+            let _ = lfc.remove(&key);
         }
 
         // Sink related option, apart from format
diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs
@@ -63,7 +63,8 @@ impl TableProviderFactory for ListingTableFactory {
             ))?
             .create(session_state, &cmd.options)?;
 
-        let mut table_path = ListingTableUrl::parse(&cmd.location)?;
+        let mut table_path =
+            ListingTableUrl::parse(&cmd.location)?.with_table_ref(cmd.name.clone());
         let file_extension = match table_path.is_collection() {
             // Setting the extension to be empty instead of allowing the default extension seems
             // odd, but was done to ensure existing behavior isn't modified. It seems like this
diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs
@@ -74,6 +74,7 @@ use datafusion_common::{
     tree_node::{TreeNodeRecursion, TreeNodeVisitor},
 };
 pub use datafusion_execution::TaskContext;
+use datafusion_execution::cache::TableScopedPath;
 use datafusion_execution::cache::cache_manager::{
     DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, DEFAULT_LIST_FILES_CACHE_TTL,
     DEFAULT_METADATA_CACHE_LIMIT,
@@ -101,6 +102,7 @@ use datafusion_session::SessionStore;
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
+use log::warn;
 use object_store::ObjectStore;
 use parking_lot::RwLock;
 use url::Url;
@@ -1315,7 +1317,7 @@ impl SessionContext {
         let table = table_ref.table().to_owned();
         let maybe_schema = {
             let state = self.state.read();
-            let resolved = state.resolve_table_ref(table_ref);
+            let resolved = state.resolve_table_ref(table_ref.clone());
             state
                 .catalog_list()
                 .catalog(&resolved.catalog)
@@ -1327,12 +1329,35 @@ impl SessionContext {
             && table_provider.table_type() == table_type
         {
             schema.deregister_table(&table)?;
+            if table_type == TableType::Base {
+                self.drop_list_files_cache_entries(&table_ref, &table_provider);
+            }
             return Ok(true);
         }
 
         Ok(false)
     }
 
+    fn drop_list_files_cache_entries(
+        &self,
+        table_ref: &TableReference,
+        table_provider: &Arc<dyn TableProvider>,
+    ) {
+        if let Some(lfc) = self.runtime_env().cache_manager.get_list_files_cache()
+            && let Some(listing_table) =
+                table_provider.as_any().downcast_ref::<ListingTable>()
+        {
+            for table_path in listing_table.table_paths() {
+                let key =
+                    TableScopedPath(Some(table_ref.clone()), table_path.prefix().clone());
+                println!("removing entry for key {key:?}");
+                if lfc.remove(&key).is_none() {
+                    warn!("list files cache for key {key:?} not found");
+                }
+            }
+        }
+    }
+
     async fn create_function(&self, stmt: CreateFunction) -> Result<DataFrame> {
         let function = {
             let state = self.state.read().clone();
diff --git a/datafusion/core/tests/catalog_listing/pruned_partition_list.rs b/datafusion/core/tests/catalog_listing/pruned_partition_list.rs
@@ -18,6 +18,7 @@
 use std::sync::Arc;
 
 use arrow_schema::DataType;
+use datafusion_sql::TableReference;
 use futures::{FutureExt, StreamExt as _, TryStreamExt as _};
 use object_store::{ObjectStore as _, memory::InMemory, path::Path};
 
diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs
@@ -17,7 +17,8 @@
 
 use std::sync::Arc;
 
-use datafusion_common::{DataFusionError, Result};
+use datafusion_common::{DataFusionError, Result, TableReference};
+use datafusion_execution::cache::TableScopedPath;
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_session::Session;
 
@@ -41,6 +42,8 @@ pub struct ListingTableUrl {
     prefix: Path,
     /// An optional glob expression used to filter files
     glob: Option<Pattern>,
+
+    table_ref: Option<TableReference>,
 }
 
 impl ListingTableUrl {
@@ -145,7 +148,12 @@ impl ListingTableUrl {
     /// to create a [`ListingTableUrl`].
     pub fn try_new(url: Url, glob: Option<Pattern>) -> Result<Self> {
         let prefix = Path::from_url_path(url.path())?;
-        Ok(Self { url, prefix, glob })
+        Ok(Self {
+            url,
+            prefix,
+            glob,
+            table_ref: None,
+        })
     }
 
     /// Returns the URL scheme
@@ -255,7 +263,14 @@ impl ListingTableUrl {
         };
 
         let list: BoxStream<'a, Result<ObjectMeta>> = if self.is_collection() {
-            list_with_cache(ctx, store, &self.prefix, prefix.as_ref()).await?
+            list_with_cache(
+                ctx,
+                store,
+                self.table_ref.as_ref(),
+                &self.prefix,
+                prefix.as_ref(),
+            )
+            .await?
         } else {
             match store.head(&full_prefix).await {
                 Ok(meta) => futures::stream::once(async { Ok(meta) })
@@ -264,7 +279,14 @@ impl ListingTableUrl {
                 // If the head command fails, it is likely that object doesn't exist.
                 // Retry as though it were a prefix (aka a collection)
                 Err(object_store::Error::NotFound { .. }) => {
-                    list_with_cache(ctx, store, &self.prefix, prefix.as_ref()).await?
+                    list_with_cache(
+                        ctx,
+                        store,
+                        self.table_ref.as_ref(),
+                        &self.prefix,
+                        prefix.as_ref(),
+                    )
+                    .await?
                 }
                 Err(e) => return Err(e.into()),
             }
@@ -323,6 +345,15 @@ impl ListingTableUrl {
             Pattern::new(glob).map_err(|e| DataFusionError::External(Box::new(e)))?;
         Self::try_new(self.url, Some(glob))
     }
+
+    pub fn with_table_ref(mut self, table_ref: TableReference) -> Self {
+        self.table_ref = Some(table_ref);
+        self
+    }
+
+    pub fn get_table_ref(&self) -> &Option<TableReference> {
+        &self.table_ref
+    }
 }
 
 /// Lists files with cache support, using prefix-aware lookups.
@@ -345,6 +376,7 @@ impl ListingTableUrl {
 async fn list_with_cache<'b>(
     ctx: &'b dyn Session,
     store: &'b dyn ObjectStore,
+    table_ref: Option<&TableReference>,
     table_base_path: &Path,
     prefix: Option<&Path>,
 ) -> Result<BoxStream<'b, Result<ObjectMeta>>> {
@@ -367,9 +399,16 @@ async fn list_with_cache<'b>(
             // Convert prefix to Option<Path> for cache lookup
             let prefix_filter = prefix.cloned();
 
+            println!("list files cache is set");
+
+            let table_scoped_base_path =
+                TableScopedPath(table_ref.cloned(), table_base_path.clone());
+
+            println!("table scoped path {table_scoped_base_path:?}");
+
             // Try cache lookup with optional prefix filter
             let vec = if let Some(res) =
-                cache.get_with_extra(table_base_path, &prefix_filter)
+                cache.get_with_extra(&table_scoped_base_path, &prefix_filter)
             {
                 debug!("Hit list files cache");
                 res.as_ref().clone()
@@ -380,7 +419,7 @@ async fn list_with_cache<'b>(
                     .list(Some(table_base_path))
                     .try_collect::<Vec<ObjectMeta>>()
                     .await?;
-                cache.put(table_base_path, Arc::new(vec.clone()));
+                cache.put(&table_scoped_base_path, Arc::new(vec.clone()));
 
                 // If a prefix filter was requested, apply it to the results
                 if prefix.is_some() {
diff --git a/datafusion/execution/src/cache/cache_manager.rs b/datafusion/execution/src/cache/cache_manager.rs
@@ -17,6 +17,7 @@
 
 use crate::cache::cache_unit::DefaultFilesMetadataCache;
 use crate::cache::list_files_cache::ListFilesEntry;
+use crate::cache::list_files_cache::TableScopedPath;
 use crate::cache::{CacheAccessor, DefaultListFilesCache};
 use datafusion_common::stats::Precision;
 use datafusion_common::{Result, Statistics};
@@ -81,7 +82,7 @@ pub struct FileStatisticsCacheEntry {
 ///
 /// See [`crate::runtime_env::RuntimeEnv`] for more details.
 pub trait ListFilesCache:
-    CacheAccessor<Path, Arc<Vec<ObjectMeta>>, Extra = Option<Path>>
+    CacheAccessor<TableScopedPath, Arc<Vec<ObjectMeta>>, Extra = Option<Path>>
 {
     /// Returns the cache's memory limit in bytes.
     fn cache_limit(&self) -> usize;
@@ -96,7 +97,7 @@ pub trait ListFilesCache:
     fn update_cache_ttl(&self, ttl: Option<Duration>);
 
     /// Retrieves the information about the entries currently cached.
-    fn list_entries(&self) -> HashMap<Path, ListFilesEntry>;
+    fn list_entries(&self) -> HashMap<TableScopedPath, ListFilesEntry>;
 }
 
 /// Generic file-embedded metadata used with [`FileMetadataCache`].
diff --git a/datafusion/execution/src/cache/list_files_cache.rs b/datafusion/execution/src/cache/list_files_cache.rs
diff --git a/datafusion/execution/src/cache/mod.rs b/datafusion/execution/src/cache/mod.rs
