Skip to content

Commit

Permalink
Introduce the calculation for TO_MANY relationship (#626)
Browse files Browse the repository at this point in the history
* upgrade to datafusio v39.0.0

* fix the 39.0.0 change

* introduce to many calculation

* add an example

* stop panic and use DataFusionError instead

* change the sqllogictests flow

* stop panic and unwrap directly

* fix count(*) aggregation for model and registered physical table

* fmt and clippy

* add one-one-many test case

* fmt and clippy

* fix test

* update cargo lock

* fmt and fix test
  • Loading branch information
goldmedal authored Jun 24, 2024
1 parent 5fe82bc commit 8431723
Show file tree
Hide file tree
Showing 25 changed files with 1,440 additions and 928 deletions.
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

0 comments on commit 8431723

Please sign in to comment.