use arrow::util::pretty::print_batches;
use arrow_schema::DataType;
use datafusion::error::Result;
use datafusion::physical_plan::display::DisplayableExecutionPlan;
use datafusion::physical_plan::{collect, ExecutionPlan};
use datafusion::{
    config::{ConfigField, TableParquetOptions},
    datasource::{file_format::parquet::ParquetFormat, listing::ListingOptions},
    prelude::*,
};
use std::sync::Arc;
use tokio;

#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
    let ctx = SessionContext::new();
    ctx.state_ref()
        .write()
        .config_mut()
        .options_mut()
        .sql_parser
        .enable_ident_normalization = false;
    ctx.state_ref()
        .write()
        .config_mut()
        .options_mut()
        .execution
        .target_partitions = 1;
    let mut opts = TableParquetOptions::default();
    opts.set("pushdown_filters", "true").unwrap();
    let format = ParquetFormat::new().with_options(opts);
    let options = ListingOptions::new(Arc::new(format))
        .with_table_partition_cols(vec![("A".to_owned(), DataType::Int32)]);
    ctx.register_listing_table(
        "hits",
        "/path",
        options,
        None,
        None,
    )
    .await
    .unwrap();

    let pp = plan(&ctx, "select count(*) from hits where Age > 10")
        .await
        .unwrap();
    let rb = collect(pp, ctx.task_ctx()).await.unwrap();
    print_batches(&rb).unwrap();

    Ok(())
}

async fn plan(ctx: &SessionContext, sql: &str) -> Result<Arc<dyn ExecutionPlan>> {
    let now = std::time::Instant::now();
    let lp = Arc::new(ctx.state().create_logical_plan(sql).await?);
    println!("LogicalPlan:\n{:?}", lp);
    let state = ctx.state();
    let lp = state.optimize(&lp).unwrap();
    println!("Optimized LogicalPlan:\n{:?}", lp);
    let pp = state.create_physical_plan(&lp).await.unwrap();
    println!(
        "Plan duration: {}us. Optimized PhysicalPlan:\n{}",
        now.elapsed().as_micros(),
        DisplayableExecutionPlan::new(&*pp).indent(true)
    );
    Ok(pp)
}

let stats = if self.pushdown_filters() && self.predicate.is_some() {
    self.projected_statistics.clone().into_inexact()
} else {
    self.projected_statistics.clone()
};