Function write_out

datafusion-examples/examples/dataframe/dataframe.rs:203–270

Use the DataFrame API to: 1. Write out a DataFrame to a table 2. Write out a DataFrame to a parquet file 3. Write out a DataFrame to a csv file 4. Write out a DataFrame to a json file

(ctx: &SessionContext)

/// 3. Write out a DataFrame to a csv file
/// 4. Write out a DataFrame to a json file
async fn write_out(ctx: &SessionContext) -> Result<()> {
    let array = StringViewArray::from(vec!["a", "b", "c"]);
    let schema = Arc::new(Schema::new(vec![Field::new(
        "tablecol1",
        DataType::Utf8View,
        false,
    )]));
    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)])?;
    let mem_table = MemTable::try_new(schema.clone(), vec![vec![batch]])?;
    ctx.register_table("initial_data", Arc::new(mem_table))?;
    let df = ctx.table("initial_data").await?;

    // Create a single temp root with subdirectories
    let tmp_root = TempDir::new()?;
    let examples_root = tmp_root.path().join("datafusion-examples");
    create_dir_all(&examples_root).await?;
    let table_dir = examples_root.join("test_table");
    let parquet_dir = examples_root.join("test_parquet");
    let csv_dir = examples_root.join("test_csv");
    let json_dir = examples_root.join("test_json");
    create_dir_all(&table_dir).await?;
    create_dir_all(&parquet_dir).await?;
    create_dir_all(&csv_dir).await?;
    create_dir_all(&json_dir).await?;

    let create_sql = format!(
        "CREATE EXTERNAL TABLE test(tablecol1 varchar)
         STORED AS parquet
         LOCATION '{}'",
        table_dir.display()
    );
    ctx.sql(&create_sql).await?.collect().await?;

    // This is equivalent to INSERT INTO test VALUES ('a'), ('b'), ('c').
    // The behavior of write_table depends on the TableProvider's implementation
    // of the insert_into method.
    df.clone()
        .write_table("test", DataFrameWriteOptions::new())
        .await?;

    df.clone()
        .write_parquet(
            parquet_dir.to_str().unwrap(),
            DataFrameWriteOptions::new(),
            None,
        )
        .await?;

    df.clone()
        .write_csv(
            csv_dir.to_str().unwrap(),
            // DataFrameWriteOptions contains options which control how data is written
            // such as compression codec
            DataFrameWriteOptions::new(),
            Some(CsvOptions::default().with_compression(CompressionTypeVariant::GZIP)),
        )
        .await?;