[V1.0.6] Refactored and added new logic for builder

Author: DatacollectorVN

engine/.DS_Store

@@ -1,3 +1,6 @@
+[workspace]
+members = ["py_lib"]
+
 [package]
 name = "engine"
 version = "0.1.0"
@@ -12,4 +15,8 @@ tempfile = "3.3.0"
 async-trait = "0.1"
 object_store = { version = "0.11.2", features=["aws"] }
 regex = "1.11.1"
-url = "2.3.1"
+url = "2.3.1"
+
+[lib]
+name = "engine"
+crate-type = ["rlib"]

engine/Cargo.lock

@@ -0,0 +1,17 @@
+[package]
+name = "py_lib"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+pyo3 = { version = "0.18", features = ["extension-module"] }
+pyo3-asyncio = { version = "0.18", features = ["tokio-runtime"] }
+arrow = "54.0.0"
+arrow-array = "54.0.0"
+futures = "0.3"
+
+engine = {path = "../"}
+
+[lib]
+name = "sa_rust"
+crate-type = ["cdylib"]

engine/Cargo.toml

@@ -0,0 +1,33 @@
+use pyo3::prelude::PyResult;
+use pyo3::exceptions::PyRuntimeError;
+use engine::builder::pipelines;
+use pyo3::prelude::*;
+use pyo3_asyncio::tokio::future_into_py;
+
+
+async fn py_sa_to_arrow_ipc_pipeline(stm: &str) -> PyResult<Vec<u8>> {
+    Ok(
+        pipelines::sa_to_arrow_ipc_pipeline(stm).await.map_err(|e| {
+                PyRuntimeError::new_err(format!("Failed to convert PyError: {}", e))
+            }
+        )?
+    )
+}
+
+
+#[pyfunction]
+fn execute_sql<'py>(
+    py: Python<'py>,
+    query: String,
+) -> PyResult<&'py PyAny> {
+    future_into_py(py, async move { py_sa_to_arrow_ipc_pipeline(query.as_str()).await })
+}
+
+
+/// Python module definition
+#[warn(unused_variables)]
+#[pymodule]
+fn sa_rust(py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(execute_sql, m)?)?;
+    Ok(())
+}

engine/py_lib/Cargo.toml

@@ -0,0 +1,3 @@
+fn main() {
+    println!("Hello, world!");
+}

engine/py_lib/src/lib.rs

@@ -0,0 +1,22 @@
+use engine::builder::pipelines::sa_to_arrow_ipc_pipeline;
+use datafusion::common::Result;
+
+#[tokio::main]
+async fn main() -> Result<()>{
+    let stm: &str = r#"SELECT
+            st."id" AS "Student Id",
+            st."name" AS "Student Name",
+            st."age" AS "Student Age",
+            s."subject" AS "Subject",
+            s."score" AS "Score"
+        FROM
+            "file:///Users/nathanngo/Projects/my/SQLAnyWhere/engine/.data/bin/ex-local-storage-application/students.csv" AS st
+        JOIN
+            "file:///Users/nathanngo/Projects/my/SQLAnyWhere/engine/.data/bin/ex-local-storage-application/scores.csv" AS s
+            ON
+                st.id = s.student_id
+    "#;
+    let buffer: Vec<u8> = sa_to_arrow_ipc_pipeline(stm).await?;
+    println!("{:?}", buffer);
+    Ok(())
+}

engine/py_lib/src/main.rs

@@ -0,0 +1 @@
+pub mod pipelines;

engine/src/bin/ex_sa_to_arrow_ipc_pipeline.rs

@@ -0,0 +1,80 @@
+use crate::datafusion::SaDataFusion;
+use crate::object_storage::{SaS3, SaLocalStorage};
+use crate::helper;
+use std::sync::Arc;
+use datafusion::common::Result;
+use datafusion::datasource::file_format::{
+    FileFormat,
+    csv::CsvFormat,
+    parquet::ParquetFormat
+};
+use datafusion::prelude::DataFrame;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::ipc::writer::StreamWriter;
+use std::env;
+
+
+async fn sa_query(sa_datafusion: SaDataFusion, stm: &str) -> Result<DataFrame> {
+    let uris: Vec<&str> = helper::sql_parser(stm);
+    for uri in uris {
+        println!("[sa_query]: Detected URI: {}", uri);
+        let file_format: Arc<dyn FileFormat> = if uri.ends_with(".csv") {
+            Arc::new(CsvFormat::default())
+        } else if uri.ends_with(".parquet") {
+            Arc::new(ParquetFormat::default())
+        } else {
+            panic!("[sa_query]: Unsupported file format")
+        };
+
+        if uri.starts_with("file://") {
+            let local_storage: SaLocalStorage = SaLocalStorage::new_with_file_uri(uri)
+                .init_table_provider(
+                    &sa_datafusion,
+                    file_format,
+                    Some(false)
+            ).await?;
+            sa_datafusion.register_sa_storage(Arc::new(local_storage)).await?;
+
+        }
+        else if uri.starts_with("s3://") {
+            let s3_region: String  = env::var("AWS_S3_REGION").unwrap_or("us-east-1".to_string());
+            let s3_storage: SaS3 = SaS3::new_from_s3_uri(uri)
+                .init_table_provider(
+                    s3_region.as_str(),
+                    &sa_datafusion,
+                    file_format,
+                    Some(false)
+                ).await?;
+                sa_datafusion.register_sa_storage(Arc::new(s3_storage)).await?;
+        }
+        else {
+            panic!("[sa_query]: Unsupported file protocal")
+        }
+    }
+    let df: DataFrame = sa_datafusion.execute_sql(&stm.to_string()).await?;
+    Ok(df)
+}
+
+
+pub async fn sa_to_dataframe_pipeline(stm: &str) -> Result<DataFrame> {
+    let sa_datafusion: SaDataFusion = SaDataFusion::new();
+    Ok(sa_query(sa_datafusion, stm).await?)
+}
+
+
+pub async fn sa_to_arrow_ipc_pipeline(stm: &str) -> Result<Vec<u8>> {
+    let sa_datafusion: SaDataFusion = SaDataFusion::new();
+    let df: DataFrame = sa_query(sa_datafusion, stm).await?;
+    let record_batches: Vec<RecordBatch> = df.collect().await?;
+    // Serialize RecordBatches to Arrow IPC
+    let mut buffer: Vec<u8> = Vec::new();
+    {
+        let mut writer: StreamWriter<&mut Vec<u8>> = StreamWriter::try_new(&mut buffer, &record_batches[0].schema())?;
+        for batch in record_batches {
+            writer.write(&batch)?
+        }
+        writer.finish()?;
+    }
+
+    Ok(buffer)
+}

engine/src/builder/mod.rs

@@ -9,6 +9,7 @@ use datafusion::execution::SessionState;
 use url::Url;
 
 
+#[derive(Clone)]
 pub struct SaDataFusion {
     pub ctx: SessionContext,
 }
@@ -27,7 +28,6 @@ impl SaDataFusion {
         self.ctx.sql(stm).await
     }
 
-
     pub fn register_object_store(&self, url: &Url, object_store: Arc<dyn ObjectStore>) {
         self.ctx.runtime_env().register_object_store(url, object_store);
     }