Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce the calculation for TO_MANY relationship #626

Merged
merged 14 commits into from
Jun 24, 2024
495 changes: 139 additions & 356 deletions wren-modeling-py/Cargo.lock

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion wren-modeling-py/src/errors.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::string::FromUtf8Error;
use base64::DecodeError;
use pyo3::exceptions::PyException;
use pyo3::PyErr;
use std::string::FromUtf8Error;
use thiserror::Error;

#[derive(Error, Debug)]
Expand Down
26 changes: 16 additions & 10 deletions wren-modeling-py/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,24 @@ use base64::prelude::*;
use pyo3::prelude::*;

use wren_core::mdl;
use wren_core::mdl::AnalyzedWrenMDL;
use wren_core::mdl::manifest::Manifest;
use wren_core::mdl::AnalyzedWrenMDL;

use crate::errors::CoreError;

mod errors;

#[pyfunction]
fn transform_sql(mdl_base64: &str, sql: &str) -> Result<String, CoreError> {
let mdl_json_bytes = BASE64_STANDARD.decode(mdl_base64).map_err(CoreError::from)?;
let mdl_json_bytes = BASE64_STANDARD
.decode(mdl_base64)
.map_err(CoreError::from)?;
let mdl_json = String::from_utf8(mdl_json_bytes).map_err(CoreError::from)?;
let manifest = serde_json::from_str::<Manifest>(&mdl_json)?;
let analyzed_mdl = AnalyzedWrenMDL::analyze(manifest);

let Ok(analyzed_mdl) = AnalyzedWrenMDL::analyze(manifest) else {
return Err(CoreError::new("Failed to analyze manifest"));
};
match mdl::transform_sql(Arc::new(analyzed_mdl), sql) {
Ok(transformed_sql) => Ok(transformed_sql),
Err(e) => Err(CoreError::new(&e.to_string())),
Expand All @@ -32,8 +37,8 @@ fn wren_core_wrapper(m: &Bound<'_, PyModule>) -> PyResult<()> {

#[cfg(test)]
mod tests {
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use serde_json::Value;

use crate::transform_sql;
Expand All @@ -58,11 +63,12 @@ mod tests {
}"#;
let v: Value = serde_json::from_str(data).unwrap();
let mdl_base64: String = BASE64_STANDARD.encode(v.to_string().as_bytes());
let transformed_sql = transform_sql(
&mdl_base64,
"SELECT * FROM my_catalog.my_schema.customer",
)
.unwrap();
assert_eq!(transformed_sql, r##"SELECT "customer"."custkey", "customer"."name" FROM (SELECT "customer"."custkey", "customer"."name" FROM (SELECT "main"."customer"."c_custkey" AS "custkey", "main"."customer"."c_name" AS "name" FROM "main"."customer") AS "customer") AS "customer""##);
let transformed_sql =
transform_sql(&mdl_base64, "SELECT * FROM my_catalog.my_schema.customer")
.unwrap();
assert_eq!(
transformed_sql,
r#"SELECT customer.custkey, customer."name" FROM (SELECT customer.custkey, customer."name" FROM (SELECT main.customer.c_custkey AS custkey, main.customer.c_name AS "name" FROM main.customer) AS customer) AS customer"#
);
}
}
2 changes: 1 addition & 1 deletion wren-modeling-py/tests/test_modeling_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,5 +27,5 @@ def test_transform_sql():
rewritten_sql = wren_core.transform_sql(manifest_str, sql)
assert (
rewritten_sql
== 'SELECT "customer"."custkey", "customer"."name" FROM (SELECT "customer"."custkey", "customer"."name" FROM (SELECT "main"."customer"."c_custkey" AS "custkey", "main"."customer"."c_name" AS "name" FROM "main"."customer") AS "customer") AS "customer"'
== 'SELECT customer.custkey, customer."name" FROM (SELECT customer.custkey, customer."name" FROM (SELECT main.customer.c_custkey AS custkey, main.customer.c_name AS "name" FROM main.customer) AS customer) AS customer'
)
3 changes: 1 addition & 2 deletions wren-modeling-rs/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,8 @@ rust-version = "1.78"
version = "0.1.0"

[workspace.dependencies]
arrow-schema = { version = "51.0.0", default-features = false }
async-trait = "0.1.80"
datafusion = { version = "38.0.0" }
datafusion = { version = "39.0.0" }
env_logger = "0.11.3"
log = { version = "0.4.14" }
petgraph = "0.6.5"
Expand Down
1 change: 0 additions & 1 deletion wren-modeling-rs/core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ name = "wren_core"
path = "src/lib.rs"

[dependencies]
arrow-schema = { workspace = true }
async-trait = { workspace = true }
datafusion = { workspace = true }
env_logger = { workspace = true }
Expand Down
Loading