Add a memory bound FileStatisticsCache for the Listing Table

datafusion-cli/src/main.rs

@@ -443,10 +443,7 @@ mod tests {
     use super::*;
     use datafusion::{
         common::test_util::batches_to_string,
-        execution::cache::{
-            DefaultListFilesCache, cache_manager::CacheManagerConfig,
-            cache_unit::DefaultFileStatisticsCache,
-        },
+        execution::cache::{DefaultListFilesCache, cache_manager::CacheManagerConfig},
         prelude::{ParquetReadOptions, col, lit, split_part},
     };
     use insta::assert_snapshot;
@@ -656,8 +653,6 @@ mod tests {
         Ok(())
     }
 
-    /// Shows that the statistics cache is not enabled by default yet
-    /// See https://github.com/apache/datafusion/issues/19217
     #[tokio::test]
     async fn test_statistics_cache_default() -> Result<(), DataFusionError> {
         let ctx = SessionContext::new();
@@ -687,57 +682,6 @@ mod tests {
             .await?;
         }
 
-        // When the cache manager creates a StatisticsCache by default,
-        // the contents will show up here
-        let sql = "SELECT split_part(path, '/', -1) as filename, file_size_bytes, num_rows, num_columns, table_size_bytes from statistics_cache() order by filename";
-        let df = ctx.sql(sql).await?;
-        let rbs = df.collect().await?;
-        assert_snapshot!(batches_to_string(&rbs),@r"
-        ++
-        ++
-        ");
-
-        Ok(())
-    }
-
-    // Can be removed when https://github.com/apache/datafusion/issues/19217 is resolved
-    #[tokio::test]
-    async fn test_statistics_cache_override() -> Result<(), DataFusionError> {
-        // Install a specific StatisticsCache implementation
-        let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default());
-        let cache_config = CacheManagerConfig::default()
-            .with_files_statistics_cache(Some(file_statistics_cache.clone()));
-        let runtime = RuntimeEnvBuilder::new()
-            .with_cache_manager(cache_config)
-            .build()?;
-        let config = SessionConfig::new().with_collect_statistics(true);
-        let ctx = SessionContext::new_with_config_rt(config, Arc::new(runtime));
-
-        ctx.register_udtf(
-            "statistics_cache",
-            Arc::new(StatisticsCacheFunc::new(
-                ctx.task_ctx().runtime_env().cache_manager.clone(),
-            )),
-        );
-
-        for filename in [
-            "alltypes_plain",
-            "alltypes_tiny_pages",
-            "lz4_raw_compressed_larger",
-        ] {
-            ctx.sql(
-                format!(
-                    "create external table {filename}
-                    stored as parquet
-                    location '../parquet-testing/data/{filename}.parquet'",
-                )
-                .as_str(),
-            )
-            .await?
-            .collect()
-            .await?;
-        }
-
         let sql = "SELECT split_part(path, '/', -1) as filename, file_size_bytes, num_rows, num_columns, table_size_bytes from statistics_cache() order by filename";
         let df = ctx.sql(sql).await?;
         let rbs = df.collect().await?;

datafusion/catalog-listing/src/helpers.rs

@@ -359,7 +359,7 @@ fn try_into_partitioned_file(
 
     let mut pf: PartitionedFile = object_meta.into();
     pf.partition_values = partition_values;
-
+    pf.table_reference.clone_from(table_path.get_table_ref());
     Ok(pf)
 }
 

datafusion/catalog-listing/src/table.rs

@@ -36,7 +36,6 @@ use datafusion_datasource::{
 };
 use datafusion_execution::cache::TableScopedPath;
 use datafusion_execution::cache::cache_manager::FileStatisticsCache;
-use datafusion_execution::cache::cache_unit::DefaultFileStatisticsCache;
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType};
@@ -187,7 +186,7 @@ pub struct ListingTable {
     /// The SQL definition for this table, if any
     definition: Option<String>,
     /// Cache for collected file statistics
-    collected_statistics: Arc<dyn FileStatisticsCache>,
+    collected_statistics: Option<Arc<dyn FileStatisticsCache>>,
     /// Constraints applied to this table
     constraints: Constraints,
     /// Column default expressions for columns that are not physically present in the data files
@@ -231,7 +230,7 @@ impl ListingTable {
             schema_source,
             options,
             definition: None,
-            collected_statistics: Arc::new(DefaultFileStatisticsCache::default()),
+            collected_statistics: None,
             constraints: Constraints::default(),
             column_defaults: HashMap::new(),
             expr_adapter_factory: config.expr_adapter_factory,
@@ -260,10 +259,8 @@ impl ListingTable {
     /// Setting a statistics cache on the `SessionContext` can avoid refetching statistics
     /// multiple times in the same session.
     ///
-    /// If `None`, creates a new [`DefaultFileStatisticsCache`] scoped to this query.
     pub fn with_cache(mut self, cache: Option<Arc<dyn FileStatisticsCache>>) -> Self {
-        self.collected_statistics =
-            cache.unwrap_or_else(|| Arc::new(DefaultFileStatisticsCache::default()));
+        self.collected_statistics = cache;
         self
     }
 
@@ -802,11 +799,15 @@ impl ListingTable {
     ) -> datafusion_common::Result<(Arc<Statistics>, Option<LexOrdering>)> {
         use datafusion_execution::cache::cache_manager::CachedFileMetadata;
 
-        let path = &part_file.object_meta.location;
+        let path = TableScopedPath {
+            table: part_file.table_reference.clone(),
+            path: part_file.object_meta.location.clone(),
+        };
         let meta = &part_file.object_meta;
 
         // Check cache first - if we have valid cached statistics and ordering
-        if let Some(cached) = self.collected_statistics.get(path)
+        if let Some(cache) = &self.collected_statistics
+            && let Some(cached) = cache.get(&path)
             && cached.is_valid_for(meta)
         {
             // Return cached statistics and ordering
@@ -823,14 +824,16 @@ impl ListingTable {
         let statistics = Arc::new(file_meta.statistics);
 
         // Store in cache
-        self.collected_statistics.put(
-            path,
-            CachedFileMetadata::new(
-                meta.clone(),
-                Arc::clone(&statistics),
-                file_meta.ordering.clone(),
-            ),
-        );
+        if let Some(cache) = &self.collected_statistics {
+            cache.put(
+                &path,
+                CachedFileMetadata::new(
+                    meta.clone(),
+                    Arc::clone(&statistics),
+                    file_meta.ordering.clone(),
+                ),
+            );
+        }
 
         Ok((statistics, file_meta.ordering))
     }