Repository: roapi/roapi (language: python)

Issue:
Allow engine (datafusion) configuration in roapi yaml
Referring to the [configuration section](https://arrow.apache.org/datafusion/user-guide/configs.html) of datafusion docs. I was thinking if we should allow the end users to set options on the roapi server level set options like change of dialect etc.

This would look something like the below in the YAML config.
```
addr:
    ...

tables:
    ...

configuration:
    datafusion:
        sql_parser.dialect: "Hive"
        explain.physical_plan_only: true
        ...

```


Discussion / hints:
Great idea :+1: 
can i take a stab at this ?
go for it! let me know if you need help reviewing it.

On Tue, Aug 1, 2023, 3:00 PM Joe ***@***.***> wrote:

> can i take a stab at this ?
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/roapi/roapi/issues/274#issuecomment-1659932209>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ACU6QOBEC2WHY3NHCQVPCXDXTDECHANCNFSM6AAAAAAXQ4TTVY>
> .
> You are receiving this because you authored the thread.Message ID:
> ***@***.***>
>


Fix patch:
diff --git a/roapi/src/config.rs b/roapi/src/config.rs
index 18b5f735e..cb9cdf767 100644
--- a/roapi/src/config.rs
+++ b/roapi/src/config.rs
@@ -1,3 +1,6 @@
+use anyhow::Ok;
+use columnq::datafusion::config::ConfigOptions;
+use columnq::SessionConfig;
 use serde_derive::Deserialize;
 
 use anyhow::{bail, Context, Result};
@@ -6,6 +9,7 @@ use columnq::encoding;
 use columnq::table::parse_table_uri_arg;
 use columnq::table::KeyValueSource;
 use columnq::table::TableSource;
+use std::collections::HashMap;
 use std::fs;
 use std::time::Duration;
 
@@ -26,6 +30,8 @@ pub struct Config {
     pub kvstores: Vec<KeyValueSource>,
     #[serde(default)]
     pub response_format: encoding::ContentType,
+    #[serde(default)]
+    pub datafusion: Option<HashMap<String, String>>,
 }
 
 fn table_arg() -> clap::Arg<'static> {
@@ -164,3 +170,18 @@ pub fn get_configuration() -> Result<Config, anyhow::Error> {
 
     Ok(config)
 }
+
+impl Config {
+    pub fn get_datafusion_config(&self) -> Result<SessionConfig> {
+        match &self.datafusion {
+            Some(df_cfg) => {
+                let mut opt = ConfigOptions::default();
+                for (k, v) in df_cfg {
+                    opt.set(format!("datafusion.{}", k).as_str(), v)?;
+                }
+                Ok(opt.into())
+            }
+            None => Ok(SessionConfig::default()),
+        }
+    }
+}
diff --git a/roapi/src/context.rs b/roapi/src/context.rs
index eac35a0eb..6fa66d158 100644
--- a/roapi/src/context.rs
+++ b/roapi/src/context.rs
@@ -23,7 +23,10 @@ pub struct RawRoapiContext {
 
 impl RawRoapiContext {
     pub async fn new(config: &Config) -> anyhow::Result<Self> {
-        let mut cq = ColumnQ::new();
+        let mut cq = match config.get_datafusion_config() {
+            Ok(df_cfg) => ColumnQ::new_with_config(df_cfg),
+            _ => ColumnQ::new(),
+        };
 
         if config.tables.is_empty() && config.kvstores.is_empty() {
             anyhow::bail!("No table nor kvstore found in config");


Test patch:
diff --git a/roapi/tests/config_test.rs b/roapi/tests/config_test.rs
new file mode 100644
index 000000000..25f5b7732
--- /dev/null
+++ b/roapi/tests/config_test.rs
@@ -0,0 +1,22 @@
+use anyhow::Result;
+use roapi::config::Config;
+use std::fs;
+
+mod helpers;
+
+#[test]
+fn test_load_yaml_datafusion_config() -> Result<()> {
+    let config_path = helpers::test_data_path("./test_datafusion_config.yml");
+    let config_content = fs::read_to_string(config_path)?;
+
+    let cfg: Config = serde_yaml::from_str(&config_content)?;
+    let df_cfg = cfg.get_datafusion_config()?;
+
+    assert_eq!(df_cfg.options().sql_parser.dialect, "Hive");
+    assert_eq!(df_cfg.options().explain.physical_plan_only, true);
+    assert_eq!(df_cfg.options().optimizer.max_passes, 10);
+    assert_eq!(df_cfg.options().execution.batch_size, 100);
+    assert_eq!(df_cfg.options().catalog.format, Some("parquet".to_string()));
+
+    Ok(())
+}
diff --git a/roapi/tests/helpers.rs b/roapi/tests/helpers.rs
index fde7db7c6..f11067191 100644
--- a/roapi/tests/helpers.rs
+++ b/roapi/tests/helpers.rs
@@ -13,6 +13,7 @@ pub fn test_data_path(relative_path: &str) -> String {
     d.to_string_lossy().to_string()
 }
 
+#[allow(dead_code)]
 pub async fn test_api_app_with_tables(tables: Vec<TableSource>) -> (Application, String) {
     test_api_app(tables, vec![]).await
 }
@@ -66,6 +67,7 @@ pub async fn http_post(url: &str, payload: impl Into<reqwest::Body>) -> reqwest:
         .expect("Unable to execute POST request")
 }
 
+#[allow(dead_code)]
 pub fn get_spacex_table() -> TableSource {
     let json_source_path = test_data_path("spacex_launches.json");
     TableSource::new("spacex_launches".to_string(), json_source_path)
diff --git a/test_data/test_datafusion_config.yml b/test_data/test_datafusion_config.yml
new file mode 100644
index 000000000..d9102094b
--- /dev/null
+++ b/test_data/test_datafusion_config.yml
@@ -0,0 +1,37 @@
+addr:
+  # binding address for TCP port that speaks HTTP protocol
+  http: 0.0.0.0:8000
+  # binding address for TCP port that speaks Postgres wire protocol
+  postgres: 0.0.0.0:5432
+tables:
+  - name: "s3_blogs"
+    uri: "s3://test-data/blogs.parquet"
+    option:
+      format: "parquet"
+      use_memory_table: true
+  - name: "s3_blogs_space_encode"
+    uri: "s3://test-data/blogs%20space.parquet"
+    option:
+      format: "parquet"
+      use_memory_table: true
+  - name: "s3_blogs_dir"
+    uri: "s3://test-data/blogs/"
+    option:
+      format: "parquet"
+      use_memory_table: true
+  - name: "gcs_blogs"
+    uri: "gs://test-data/blogs.parquet"
+    option:
+      format: "parquet"
+      use_memory_table: true
+  - name: "azure_blogs"
+    uri: "az://test-data/blogs.parquet"
+    option:
+      format: "parquet"
+      use_memory_table: true
+datafusion:
+  sql_parser.dialect: "Hive"
+  explain.physical_plan_only: true
+  optimizer.max_passes: 10
+  execution.batch_size: 100
+  catalog.format: "parquet"
\ No newline at end of file