feat: Add the ability to review the contents of the Metadata Cache (#17126)

* feat: Add the ability to review the contents of the Metadata Cache

* Remove e_tag from test_metadata_cache

* Add entry in the user doc about this function

* Change type to UInt64

* Fix prettier

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: nuno-faria

datafusion-cli/src/functions.rs

@@ -22,15 +22,16 @@ use std::fs::File;
 use std::str::FromStr;
 use std::sync::Arc;
 
-use arrow::array::{Int64Array, StringArray};
-use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow::array::{Int64Array, StringArray, TimestampMillisecondArray, UInt64Array};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
 use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::pretty_format_batches;
 use datafusion::catalog::{Session, TableFunctionImpl};
 use datafusion::common::{plan_err, Column};
 use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::datasource::TableProvider;
 use datafusion::error::Result;
+use datafusion::execution::cache::cache_manager::CacheManager;
 use datafusion::logical_expr::Expr;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::scalar::ScalarValue;
@@ -460,3 +461,121 @@ impl TableFunctionImpl for ParquetMetadataFunc {
         Ok(Arc::new(parquet_metadata))
     }
 }
+
+/// METADATA_CACHE table function
+#[derive(Debug)]
+struct MetadataCacheTable {
+    schema: SchemaRef,
+    batch: RecordBatch,
+}
+
+#[async_trait]
+impl TableProvider for MetadataCacheTable {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn schema(&self) -> arrow::datatypes::SchemaRef {
+        self.schema.clone()
+    }
+
+    fn table_type(&self) -> datafusion::logical_expr::TableType {
+        datafusion::logical_expr::TableType::Base
+    }
+
+    async fn scan(
+        &self,
+        _state: &dyn Session,
+        projection: Option<&Vec<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(MemorySourceConfig::try_new_exec(
+            &[vec![self.batch.clone()]],
+            TableProvider::schema(self),
+            projection.cloned(),
+        )?)
+    }
+}
+
+#[derive(Debug)]
+pub struct MetadataCacheFunc {
+    cache_manager: Arc<CacheManager>,
+}
+
+impl MetadataCacheFunc {
+    pub fn new(cache_manager: Arc<CacheManager>) -> Self {
+        Self { cache_manager }
+    }
+}
+
+impl TableFunctionImpl for MetadataCacheFunc {
+    fn call(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+        if !exprs.is_empty() {
+            return plan_err!("metadata_cache should have no arguments");
+        }
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("path", DataType::Utf8, false),
+            Field::new(
+                "file_modified",
+                DataType::Timestamp(TimeUnit::Millisecond, None),
+                false,
+            ),
+            Field::new("file_size_bytes", DataType::UInt64, false),
+            Field::new("e_tag", DataType::Utf8, true),
+            Field::new("version", DataType::Utf8, true),
+            Field::new("metadata_size_bytes", DataType::UInt64, false),
+            Field::new("hits", DataType::UInt64, false),
+            Field::new("extra", DataType::Utf8, true),
+        ]));
+
+        // construct record batch from metadata
+        let mut path_arr = vec![];
+        let mut file_modified_arr = vec![];
+        let mut file_size_bytes_arr = vec![];
+        let mut e_tag_arr = vec![];
+        let mut version_arr = vec![];
+        let mut metadata_size_bytes = vec![];
+        let mut hits_arr = vec![];
+        let mut extra_arr = vec![];
+
+        let cached_entries = self.cache_manager.get_file_metadata_cache().list_entries();
+
+        for (path, entry) in cached_entries {
+            path_arr.push(path.to_string());
+            file_modified_arr
+                .push(Some(entry.object_meta.last_modified.timestamp_millis()));
+            file_size_bytes_arr.push(entry.object_meta.size);
+            e_tag_arr.push(entry.object_meta.e_tag);
+            version_arr.push(entry.object_meta.version);
+            metadata_size_bytes.push(entry.size_bytes as u64);
+            hits_arr.push(entry.hits as u64);
+
+            let mut extra = entry
+                .extra
+                .iter()
+                .map(|(k, v)| format!("{k}={v}"))
+                .collect::<Vec<_>>();
+            extra.sort();
+            extra_arr.push(extra.join(" "));
+        }
+
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(StringArray::from(path_arr)),
+                Arc::new(TimestampMillisecondArray::from(file_modified_arr)),
+                Arc::new(UInt64Array::from(file_size_bytes_arr)),
+                Arc::new(StringArray::from(e_tag_arr)),
+                Arc::new(StringArray::from(version_arr)),
+                Arc::new(UInt64Array::from(metadata_size_bytes)),
+                Arc::new(UInt64Array::from(hits_arr)),
+                Arc::new(StringArray::from(extra_arr)),
+            ],
+        )?;
+
+        let metadata_cache = MetadataCacheTable { schema, batch };
+        Ok(Arc::new(metadata_cache))
+    }
+}

datafusion-cli/src/main.rs

@@ -30,7 +30,7 @@ use datafusion::execution::memory_pool::{
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::prelude::SessionContext;
 use datafusion_cli::catalog::DynamicObjectStoreCatalog;
-use datafusion_cli::functions::ParquetMetadataFunc;
+use datafusion_cli::functions::{MetadataCacheFunc, ParquetMetadataFunc};
 use datafusion_cli::{
     exec,
     pool_type::PoolType,
@@ -219,6 +219,14 @@ async fn main_inner() -> Result<()> {
     // register `parquet_metadata` table function to get metadata from parquet files
     ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {}));
 
+    // register `metadata_cache` table function to get the contents of the file metadata cache
+    ctx.register_udtf(
+        "metadata_cache",
+        Arc::new(MetadataCacheFunc::new(
+            ctx.task_ctx().runtime_env().cache_manager.clone(),
+        )),
+    );
+
     let mut print_options = PrintOptions {
         format: args.format,
         quiet: args.quiet,
@@ -397,7 +405,7 @@ pub fn extract_disk_limit(size: &str) -> Result<usize, String> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use datafusion::common::test_util::batches_to_string;
+    use datafusion::{common::test_util::batches_to_string, prelude::ParquetReadOptions};
     use insta::assert_snapshot;
 
     fn assert_conversion(input: &str, expected: Result<usize, String>) {
@@ -512,4 +520,97 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_metadata_cache() -> Result<(), DataFusionError> {
+        let ctx = SessionContext::new();
+        ctx.register_udtf(
+            "metadata_cache",
+            Arc::new(MetadataCacheFunc::new(
+                ctx.task_ctx().runtime_env().cache_manager.clone(),
+            )),
+        );
+
+        ctx.register_parquet(
+            "alltypes_plain",
+            "../parquet-testing/data/alltypes_plain.parquet",
+            ParquetReadOptions::new(),
+        )
+        .await?;
+
+        ctx.register_parquet(
+            "alltypes_tiny_pages",
+            "../parquet-testing/data/alltypes_tiny_pages.parquet",
+            ParquetReadOptions::new(),
+        )
+        .await?;
+
+        ctx.register_parquet(
+            "lz4_raw_compressed_larger",
+            "../parquet-testing/data/lz4_raw_compressed_larger.parquet",
+            ParquetReadOptions::new(),
+        )
+        .await?;
+
+        ctx.sql("select * from alltypes_plain")
+            .await?
+            .collect()
+            .await?;
+        ctx.sql("select * from alltypes_tiny_pages")
+            .await?
+            .collect()
+            .await?;
+        ctx.sql("select * from lz4_raw_compressed_larger")
+            .await?
+            .collect()
+            .await?;
+
+        // initial state
+        let sql = "SELECT split_part(path, '/', -1) as filename, file_size_bytes, metadata_size_bytes, hits, extra from metadata_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 | metadata_size_bytes | hits | extra            |
+        +-----------------------------------+-----------------+---------------------+------+------------------+
+        | alltypes_plain.parquet            | 1851            | 10181               | 2    | page_index=false |
+        | alltypes_tiny_pages.parquet       | 454233          | 881634              | 2    | page_index=true  |
+        | lz4_raw_compressed_larger.parquet | 380836          | 2939                | 2    | page_index=false |
+        +-----------------------------------+-----------------+---------------------+------+------------------+
+        "#);
+
+        // increase the number of hits
+        ctx.sql("select * from alltypes_plain")
+            .await?
+            .collect()
+            .await?;
+        ctx.sql("select * from alltypes_plain")
+            .await?
+            .collect()
+            .await?;
+        ctx.sql("select * from alltypes_plain")
+            .await?
+            .collect()
+            .await?;
+        ctx.sql("select * from lz4_raw_compressed_larger")
+            .await?
+            .collect()
+            .await?;
+        let sql = "select split_part(path, '/', -1) as filename, file_size_bytes, metadata_size_bytes, hits, extra from metadata_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 | metadata_size_bytes | hits | extra            |
+        +-----------------------------------+-----------------+---------------------+------+------------------+
+        | alltypes_plain.parquet            | 1851            | 10181               | 5    | page_index=false |
+        | alltypes_tiny_pages.parquet       | 454233          | 881634              | 2    | page_index=true  |
+        | lz4_raw_compressed_larger.parquet | 380836          | 2939                | 3    | page_index=false |
+        +-----------------------------------+-----------------+---------------------+------+------------------+
+        "#);
+
+        Ok(())
+    }
 }

datafusion/datasource-parquet/src/reader.rs

@@ -30,6 +30,7 @@ use parquet::arrow::arrow_reader::ArrowReaderOptions;
 use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader};
 use parquet::file::metadata::ParquetMetaData;
 use std::any::Any;
+use std::collections::HashMap;
 use std::fmt::Debug;
 use std::ops::Range;
 use std::sync::Arc;
@@ -296,4 +297,10 @@ impl FileMetadata for CachedParquetMetaData {
     fn memory_size(&self) -> usize {
         self.0.memory_size()
     }
+
+    fn extra_info(&self) -> HashMap<String, String> {
+        let page_index =
+            self.0.column_index().is_some() && self.0.offset_index().is_some();
+        HashMap::from([("page_index".to_owned(), page_index.to_string())])
+    }
 }

datafusion/execution/src/cache/cache_manager.rs

@@ -21,6 +21,7 @@ use datafusion_common::{Result, Statistics};
 use object_store::path::Path;
 use object_store::ObjectMeta;
 use std::any::Any;
+use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
@@ -42,17 +43,36 @@ pub trait FileMetadata: Any + Send + Sync {
 
     /// Returns the size of the metadata in bytes.
     fn memory_size(&self) -> usize;
+
+    /// Returns extra information about this entry (used by [`FileMetadataCache::list_entries`]).
+    fn extra_info(&self) -> HashMap<String, String>;
 }
 
 /// Cache to store file-embedded metadata.
 pub trait FileMetadataCache:
     CacheAccessor<ObjectMeta, Arc<dyn FileMetadata>, Extra = ObjectMeta>
 {
-    // Returns the cache's memory limit in bytes.
+    /// Returns the cache's memory limit in bytes.
     fn cache_limit(&self) -> usize;
 
-    // Updates the cache with a new memory limit in bytes.
+    /// Updates the cache with a new memory limit in bytes.
     fn update_cache_limit(&self, limit: usize);
+
+    /// Retrieves the information about the entries currently cached.
+    fn list_entries(&self) -> HashMap<Path, FileMetadataCacheEntry>;
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+/// Represents information about a cached metadata entry.
+/// This is used to expose the metadata cache contents to outside modules.
+pub struct FileMetadataCacheEntry {
+    pub object_meta: ObjectMeta,
+    /// Size of the cached metadata, in bytes.
+    pub size_bytes: usize,
+    /// Number of times this entry was retrieved.
+    pub hits: usize,
+    /// Additional object-specific information.
+    pub extra: HashMap<String, String>,
 }
 
 impl Debug for dyn CacheAccessor<Path, Arc<Statistics>, Extra = ObjectMeta> {