datafusion-server 0.21.0

Web server library for session-based queries using Arrow and other large datasets as data sources.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

// object_store/writer.rs

use datafusion::{
    common::config::{CsvOptions, JsonOptions, ParquetOptions, TableParquetOptions},
    dataframe::DataFrameWriteOptions,
    execution::context::SessionContext,
};

use crate::request::body::{DataSource, DataSourceFormat, DataSourceOption};
use crate::response::http_error::ResponseError;

pub async fn write(ctx: &SessionContext, data_source: &DataSource) -> Result<(), ResponseError> {
    log::debug!("object_store::writer(): {data_source:?}");

    let options = match &data_source.options {
        Some(options) => options.clone(),
        None => DataSourceOption::new(),
    };

    let write_options = DataFrameWriteOptions::new()
        // TODO: raise error when options.overwrite is true and already exists target file
        // .with_overwrite(options.overwrite.unwrap_or(false))
        .with_single_file_output(true);

    let df = ctx
        .sql(&format!("SELECT * FROM {}", data_source.name))
        .await?;

    match data_source.format {
        DataSourceFormat::Csv => {
            let csv_options = CsvOptions::default()
                .with_has_header(options.has_header.unwrap_or(true))
                .with_delimiter(options.delimiter.unwrap_or(',') as u8);

            df.write_csv(&data_source.location, write_options, Some(csv_options))
                .await?;
        }
        DataSourceFormat::NdJson => {
            let ndjson_options = JsonOptions::default();
            df.write_json(&data_source.location, write_options, Some(ndjson_options))
                .await?;
        }
        DataSourceFormat::Parquet => {
            let parquet_options = ParquetOptions {
                compression: Some("snappy".to_string()),
                created_by: format!("datafusion-server v{}", env!("CARGO_PKG_VERSION")),
                ..Default::default()
            };

            df.write_parquet(
                &data_source.location,
                write_options,
                Some(TableParquetOptions {
                    global: parquet_options,
                    ..Default::default()
                }),
            )
            .await?;
        }
        _ => {
            return Err(ResponseError::unsupported_type(format!(
                "write to object store is not supported {:?}",
                &data_source.format
            )));
        }
    }

    Ok(())
}