Function datafusion

benchmarks/src/cancellation.rs:162–205

(store: Arc<dyn ObjectStore>)

}

async fn datafusion(store: Arc<dyn ObjectStore>) -> Result<()> {
    let query = "SELECT distinct \"A\", \"B\", \"C\", \"D\", \"E\" FROM \"test_table\"";

    let config = SessionConfig::new()
        .with_target_partitions(4)
        .set_bool("datafusion.execution.parquet.pushdown_filters", true);
    let ctx = SessionContext::new_with_config(config);
    let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
    ctx.register_object_store(object_store_url.as_ref(), Arc::clone(&store));

    let file_format = ParquetFormat::default().with_enable_pruning(true);
    let listing_options = ListingOptions::new(Arc::new(file_format))
        .with_file_extension(ParquetFormat::default().get_ext());

    let table_path = ListingTableUrl::parse("test:///data/")?;

    ctx.register_listing_table(
        "test_table",
        &table_path,
        listing_options.clone(),
        None,
        None,
    )
    .await?;

    println!("Creating logical plan...");
    let logical_plan = ctx.state().create_logical_plan(query).await?;

    println!("Creating physical plan...");
    let physical_plan = Arc::new(CoalescePartitionsExec::new(
        ctx.state().create_physical_plan(&logical_plan).await?,
    ));

    println!("Executing physical plan...");
    let partition = 0;
    let task_context = Arc::new(TaskContext::from(&ctx));
    let stream = physical_plan.execute(partition, task_context).unwrap();

    println!("Getting results...");
    let results: Vec<_> = stream.try_collect().await?;
    println!("Got {} record batches", results.len());

    Ok(())
}

async fn find_or_generate_files(
    data_dir: impl AsRef<Path>,