Add a memory bound FileStatisticsCache for the Listing Table

benchmarks/src/bin/external_aggr.rs

@@ -326,7 +326,9 @@ impl ExternalAggrConfig {
         let config = ListingTableConfig::new(table_path).with_listing_options(options);
         let config = config.infer_schema(&state).await?;
 
-        Ok(Arc::new(ListingTable::try_new(config)?))
+        Ok(Arc::new(ListingTable::try_new(config)?.with_cache(
+            ctx.runtime_env().cache_manager.get_file_statistic_cache(),
+        )))
     }
 
     fn iterations(&self) -> usize {

benchmarks/src/imdb/run.rs

@@ -470,7 +470,9 @@ impl RunOpt {
             _ => unreachable!(),
         };
 
-        Ok(Arc::new(ListingTable::try_new(config)?))
+        Ok(Arc::new(ListingTable::try_new(config)?.with_cache(
+            ctx.runtime_env().cache_manager.get_file_statistic_cache(),
+        )))
     }
 
     fn iterations(&self) -> usize {

benchmarks/src/sort_pushdown.rs

@@ -273,7 +273,9 @@ impl RunOpt {
             .with_listing_options(options)
             .with_schema(schema);
 
-        Ok(Arc::new(ListingTable::try_new(config)?))
+        Ok(Arc::new(ListingTable::try_new(config)?.with_cache(
+            ctx.runtime_env().cache_manager.get_file_statistic_cache(),
+        )))
     }
 
     fn iterations(&self) -> usize {

benchmarks/src/sort_tpch.rs

@@ -351,7 +351,9 @@ impl RunOpt {
             .with_listing_options(options)
             .with_schema(schema);
 
-        Ok(Arc::new(ListingTable::try_new(config)?))
+        Ok(Arc::new(ListingTable::try_new(config)?.with_cache(
+            ctx.runtime_env().cache_manager.get_file_statistic_cache(),
+        )))
     }
 
     fn iterations(&self) -> usize {

benchmarks/src/tpcds/run.rs

@@ -347,7 +347,9 @@ impl RunOpt {
             .with_listing_options(options)
             .with_schema(schema);
 
-        Ok(Arc::new(ListingTable::try_new(config)?))
+        Ok(Arc::new(ListingTable::try_new(config)?.with_cache(
+            ctx.runtime_env().cache_manager.get_file_statistic_cache(),
+        )))
     }
 
     fn iterations(&self) -> usize {

benchmarks/src/tpch/run.rs

@@ -342,7 +342,9 @@ impl RunOpt {
             .with_listing_options(options)
             .with_schema(schema);
 
-        Ok(Arc::new(ListingTable::try_new(config)?))
+        Ok(Arc::new(ListingTable::try_new(config)?.with_cache(
+            ctx.runtime_env().cache_manager.get_file_statistic_cache(),
+        )))
     }
 
     fn iterations(&self) -> usize {

datafusion-cli/src/functions.rs

@@ -633,6 +633,7 @@ impl TableFunctionImpl for StatisticsCacheFunc {
 
         let schema = Arc::new(Schema::new(vec![
             Field::new("path", DataType::Utf8, false),
+            Field::new("table", DataType::Utf8, false),
             Field::new(
                 "file_modified",
                 DataType::Timestamp(TimeUnit::Millisecond, None),
@@ -649,6 +650,7 @@ impl TableFunctionImpl for StatisticsCacheFunc {
 
         // construct record batch from metadata
         let mut path_arr = vec![];
+        let mut table_arr = vec![];
         let mut file_modified_arr = vec![];
         let mut file_size_bytes_arr = vec![];
         let mut e_tag_arr = vec![];
@@ -661,7 +663,9 @@ impl TableFunctionImpl for StatisticsCacheFunc {
         if let Some(file_statistics_cache) = self.cache_manager.get_file_statistic_cache()
         {
             for (path, entry) in file_statistics_cache.list_entries() {
-                path_arr.push(path.to_string());
+                path_arr.push(path.path.to_string());
+                table_arr
+                    .push(path.table.map_or_else(|| "".to_string(), |t| t.to_string()));
                 file_modified_arr
                     .push(Some(entry.object_meta.last_modified.timestamp_millis()));
                 file_size_bytes_arr.push(entry.object_meta.size);
@@ -678,6 +682,7 @@ impl TableFunctionImpl for StatisticsCacheFunc {
             schema.clone(),
             vec![
                 Arc::new(StringArray::from(path_arr)),
+                Arc::new(StringArray::from(table_arr)),
                 Arc::new(TimestampMillisecondArray::from(file_modified_arr)),
                 Arc::new(UInt64Array::from(file_size_bytes_arr)),
                 Arc::new(StringArray::from(e_tag_arr)),

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,68 +682,17 @@ 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 sql = "SELECT split_part(path, '/', -1) as filename, table, 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?;
-        assert_snapshot!(batches_to_string(&rbs),@r"
-        +-----------------------------------+-----------------+--------------+-------------+------------------+
-        | filename                          | file_size_bytes | num_rows     | num_columns | table_size_bytes |
-        +-----------------------------------+-----------------+--------------+-------------+------------------+
-        | alltypes_plain.parquet            | 1851            | Exact(8)     | 11          | Absent           |
-        | alltypes_tiny_pages.parquet       | 454233          | Exact(7300)  | 13          | Absent           |
-        | lz4_raw_compressed_larger.parquet | 380836          | Exact(10000) | 1           | Absent           |
-        +-----------------------------------+-----------------+--------------+-------------+------------------+
+          +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+
+          | filename                          | table                     | file_size_bytes | num_rows     | num_columns | table_size_bytes |
+          +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+
+          | alltypes_plain.parquet            | alltypes_plain            | 1851            | Exact(8)     | 11          | Absent           |
+          | alltypes_tiny_pages.parquet       | alltypes_tiny_pages       | 454233          | Exact(7300)  | 13          | Absent           |
+          | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836          | Exact(10000) | 1           | Absent           |
+          +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+
         ");
 
         Ok(())

datafusion/catalog-listing/src/helpers.rs

@@ -21,9 +21,11 @@ use std::mem;
 use std::sync::Arc;
 
 use datafusion_catalog::Session;
-use datafusion_common::{HashMap, Result, ScalarValue, assert_or_internal_err};
-use datafusion_datasource::ListingTableUrl;
+use datafusion_common::{
+    HashMap, Result, ScalarValue, TableReference, assert_or_internal_err,
+};
 use datafusion_datasource::PartitionedFile;
+use datafusion_datasource::{FileExtensions, ListingTableUrl};
 use datafusion_expr::{BinaryExpr, Operator, lit, utils};
 
 use arrow::{
@@ -382,6 +384,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(Some(pf))
 }
@@ -416,8 +419,15 @@ pub async fn pruned_partition_list<'a>(
             table_path
         );
 
-        // if no partition col => simply list all the files
-        Ok(objects.map_ok(|object_meta| object_meta.into()).boxed())
+        // if no partition col => list all the files
+        Ok(objects
+            .try_filter_map(|object_meta| {
+                futures::future::ready(object_meta_to_partitioned_file(
+                    object_meta,
+                    table_path.get_table_ref(),
+                ))
+            })
+            .boxed())
     } else {
         let df_schema = DFSchema::from_unqualified_fields(
             partition_cols
@@ -442,6 +452,22 @@ pub async fn pruned_partition_list<'a>(
     }
 }
 
+fn object_meta_to_partitioned_file(
+    object_meta: ObjectMeta,
+    table_ref: &Option<TableReference>,
+) -> Result<Option<PartitionedFile>> {
+    Ok(Some(PartitionedFile {
+        object_meta,
+        partition_values: vec![],
+        range: None,
+        statistics: None,
+        ordering: None,
+        extensions: FileExtensions::new(),
+        metadata_size_hint: None,
+        table_reference: table_ref.clone(),
+    }))
+}
+
 /// Extract the partition values for the given `file_path` (in the given `table_path`)
 /// associated to the partitions defined by `table_partition_cols`
 pub fn parse_partitions_for_path<'a, I>(

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))
     }