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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

// object_store/reader.rs

#[cfg(feature = "avro")]
use datafusion::datasource::file_format::options::AvroReadOptions;
use datafusion::{
    datasource::file_format::options::{CsvReadOptions, JsonReadOptions, ParquetReadOptions},
    execution::context::SessionContext,
};

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

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

    let uri = location::uri::to_parts(&data_source.location)?;

    let location = if location::uri::scheme(&uri).unwrap_or(location::uri::SupportedScheme::File)
        == location::uri::SupportedScheme::File
    {
        location::file::create_data_file_path(&data_source.location)?
    } else {
        data_source.location.clone()
    };

    match data_source.format {
        DataSourceFormat::Csv => register_csv(ctx, data_source, &location).await?,
        DataSourceFormat::NdJson => register_ndjson(ctx, data_source, &location).await?,
        DataSourceFormat::Parquet => register_parquet(ctx, data_source, &location).await?,
        #[cfg(feature = "avro")]
        DataSourceFormat::Avro => register_avro(ctx, data_source, &location).await?,
        _ => {
            return Err(ResponseError::unsupported_type(format!(
                "read from object store is not supported {:?}",
                &data_source.format
            )));
        }
    }

    Ok(())
}

async fn register_csv(
    ctx: &SessionContext,
    data_source: &DataSource,
    location: &str,
) -> Result<(), ResponseError> {
    log::debug!("object_store::register_csv(): {location}");

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

    let csv_options = CsvReadOptions::default()
        .has_header(options.has_header.unwrap_or(true))
        .delimiter(options.delimiter.unwrap_or(',') as u8);

    let arrow_schema;
    let csv_options = if let Some(schema) = &data_source.schema {
        arrow_schema = schema.to_arrow_schema()?;
        csv_options.schema(&arrow_schema)
    } else {
        csv_options.schema_infer_max_records(options.infer_schema_rows.unwrap_or(100))
    };

    ctx.register_csv(&data_source.name, location, csv_options)
        .await?;

    Ok(())
}

async fn register_ndjson(
    ctx: &SessionContext,
    data_source: &DataSource,
    location: &str,
) -> Result<(), ResponseError> {
    log::debug!("object_store::register_ndjson(): {location}");

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

    let mut ndjson_options = JsonReadOptions::default();

    let arrow_schema;
    let ndjson_options = if let Some(schema) = &data_source.schema {
        arrow_schema = schema.to_arrow_schema()?;
        ndjson_options.schema(&arrow_schema)
    } else {
        ndjson_options.schema_infer_max_records = options.infer_schema_rows.unwrap_or(100);
        ndjson_options
    };

    ctx.register_json(&data_source.name, location, ndjson_options)
        .await?;

    Ok(())
}

async fn register_parquet(
    ctx: &SessionContext,
    data_source: &DataSource,
    location: &str,
) -> Result<(), ResponseError> {
    log::debug!("object_store::register_parquet(): {location}");
    ctx.register_parquet(&data_source.name, location, ParquetReadOptions::default())
        .await?;
    Ok(())
}

#[cfg(feature = "avro")]
async fn register_avro(
    ctx: &SessionContext,
    data_source: &DataSource,
    location: &str,
) -> Result<(), ResponseError> {
    log::debug!("object_store::register_avro(): {location}");

    let mut avro_options = AvroReadOptions::default();

    let arrow_schema;
    if data_source.schema.is_some() {
        arrow_schema = data_source.schema.as_ref().unwrap().to_arrow_schema()?;
        avro_options = avro_options.schema(&arrow_schema);
    }

    ctx.register_avro(&data_source.name, location, avro_options)
        .await?;

    Ok(())
}