Function test_list_files_cache

datafusion-cli/src/main.rs:702–794

()

    #[tokio::test]
    async fn test_list_files_cache() -> Result<(), DataFusionError> {
        let list_files_cache = Arc::new(DefaultListFilesCache::new(
            1024,
            Some(Duration::from_secs(1)),
        ));

        let rt = RuntimeEnvBuilder::new()
            .with_cache_manager(
                CacheManagerConfig::default()
                    .with_list_files_cache(Some(list_files_cache)),
            )
            .build_arc()
            .unwrap();

        let ctx = SessionContext::new_with_config_rt(SessionConfig::default(), rt);

        ctx.register_object_store(
            &Url::parse("mem://test_table").unwrap(),
            Arc::new(InMemory::new()),
        );

        ctx.register_udtf(
            "list_files_cache",
            Arc::new(ListFilesCacheFunc::new(
                ctx.task_ctx().runtime_env().cache_manager.clone(),
            )),
        );

        ctx.sql(
            "CREATE EXTERNAL TABLE src_table
            STORED AS PARQUET
            LOCATION '../parquet-testing/data/alltypes_plain.parquet'",
        )
        .await?
        .collect()
        .await?;

        ctx.sql("COPY (SELECT * FROM src_table) TO 'mem://test_table/0.parquet' STORED AS PARQUET").await?.collect().await?;

        ctx.sql("COPY (SELECT * FROM src_table) TO 'mem://test_table/1.parquet' STORED AS PARQUET").await?.collect().await?;

        ctx.sql(
            "CREATE EXTERNAL TABLE test_table 
            STORED AS PARQUET
            LOCATION 'mem://test_table/'
        ",
        )
        .await?
        .collect()
        .await?;

        let sql = "SELECT metadata_size_bytes, expires_in, metadata_list FROM list_files_cache()";
        let df = ctx
            .sql(sql)
            .await?
            .unnest_columns(&["metadata_list"])?
            .with_column_renamed("metadata_list", "metadata")?
            .unnest_columns(&["metadata"])?;