Function test_metadata_cache

datafusion-cli/src/main.rs:564–654

()

    #[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            | 8882                | 2    | page_index=false |
        | alltypes_tiny_pages.parquet       | 454233          | 269074              | 2    | page_index=true  |
        | lz4_raw_compressed_larger.parquet | 380836          | 1339                | 2    | page_index=false |
        +-----------------------------------+-----------------+---------------------+------+------------------+
        ");