Skip to content

Create expressions by SessionContext API #684

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.

Sign up for GitHub

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

Merged
merged 8 commits into from
Jul 22, 2024
Merged

Create expressions by SessionContext API #684

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
62cecbd
upgrade to 40.0.0 and fix the api
goldmedal Jul 18, 2024
a53aff2
fix the test cases
goldmedal Jul 18, 2024
4463720
update lock
goldmedal Jul 18, 2024
0cb55b0
adjust Dataset location
goldmedal Jul 19, 2024
7f828a2
use seesion_state to create the required expr
goldmedal Jul 19, 2024
175e9c8
Merge branch 'main' into feature/create-epxr-by-ctx
goldmedal Jul 19, 2024
5c381f8
fix test
goldmedal Jul 19, 2024
98d9996
fix format and clippy
goldmedal Jul 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions wren-modeling-py/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions wren-modeling-rs/core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,4 @@ petgraph-evcxr = "*"
serde = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros"] }
parking_lot = "0.12.3"
54 changes: 44 additions & 10 deletions wren-modeling-rs/core/src/logical_plan/analyze/plan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@ use crate::mdl::utils::{
create_remote_expr_for_model, create_wren_calculated_field_expr,
create_wren_expr_for_model, is_dag,
};
use crate::mdl::{AnalyzedWrenMDL, ColumnReference, Dataset};
use crate::mdl::Dataset;
use crate::mdl::{AnalyzedWrenMDL, ColumnReference, SessionStateRef};

#[derive(Debug)]
pub(crate) enum WrenPlan {
Expand Down Expand Up @@ -66,8 +67,9 @@ impl ModelPlanNode {
required_fields: Vec<Expr>,
original_table_scan: Option<LogicalPlan>,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state: SessionStateRef,
) -> Result<Self> {
ModelPlanNodeBuilder::new(analyzed_wren_mdl).build(
ModelPlanNodeBuilder::new(analyzed_wren_mdl, session_state).build(
model,
required_fields,
original_table_scan,
Expand All @@ -90,17 +92,22 @@ struct ModelPlanNodeBuilder {
required_calculation: Vec<WrenPlan>,
fields: VecDeque<(Option<TableReference>, Arc<Field>)>,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state: SessionStateRef,
}

impl ModelPlanNodeBuilder {
fn new(analyzed_wren_mdl: Arc<AnalyzedWrenMDL>) -> Self {
fn new(
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state: SessionStateRef,
) -> Self {
Self {
required_exprs_buffer: BTreeSet::new(),
directed_graph: Graph::new(),
model_required_fields: HashMap::new(),
required_calculation: vec![],
fields: VecDeque::new(),
analyzed_wren_mdl,
session_state,
}
}

Expand Down Expand Up @@ -139,6 +146,7 @@ impl ModelPlanNodeBuilder {
let expr = create_wren_calculated_field_expr(
column_rf,
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
)?;
let expr_plan = expr.alias(column.name());
expr_plan
Expand Down Expand Up @@ -177,6 +185,7 @@ impl ModelPlanNodeBuilder {
if self.is_contain_calculation_source(&qualified_column) {
collect_partial_model_plan(
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
&qualified_column,
&mut self.model_required_fields,
)?;
Expand All @@ -186,6 +195,7 @@ impl ModelPlanNodeBuilder {
let _ = collect_model_required_fields(
&qualified_column,
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
&mut self.model_required_fields,
);
}
Expand All @@ -194,6 +204,7 @@ impl ModelPlanNodeBuilder {
&column,
Arc::clone(&model),
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
)?;
self.model_required_fields
.entry(model_ref.clone())
Expand Down Expand Up @@ -254,6 +265,7 @@ impl ModelPlanNodeBuilder {
source,
source_required_fields,
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
)?
} else {
let Some(first_calculation) = calculate_iter.next() else {
Expand All @@ -271,6 +283,7 @@ impl ModelPlanNodeBuilder {
self.directed_graph.clone(),
&self.model_required_fields.clone(),
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
)?;

for calculation_plan in calculate_iter {
Expand Down Expand Up @@ -360,6 +373,7 @@ impl ModelPlanNodeBuilder {
if self.is_contain_calculation_source(qualified_column) {
collect_partial_model_plan(
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
qualified_column,
&mut partial_model_required_fields,
)?;
Expand All @@ -368,6 +382,7 @@ impl ModelPlanNodeBuilder {
collect_model_required_fields(
qualified_column,
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
&mut partial_model_required_fields,
)?;

Expand All @@ -384,6 +399,7 @@ impl ModelPlanNodeBuilder {
source,
source_required_fields,
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
)?;

let partial_chain = RelationChain::with_chain(
Expand All @@ -393,6 +409,7 @@ impl ModelPlanNodeBuilder {
column_graph.clone(),
&partial_model_required_fields,
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
)?;
let Some(column_rf) = self
.analyzed_wren_mdl
Expand All @@ -405,7 +422,7 @@ impl ModelPlanNodeBuilder {
column_rf,
col_expr,
partial_chain,
Arc::clone(&self.analyzed_wren_mdl),
Arc::clone(&self.session_state),
)?)))
}
}
Expand All @@ -421,6 +438,7 @@ fn is_required_column(expr: &Expr, name: &str) -> bool {

fn collect_partial_model_plan(
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
qualified_column: &Column,
required_fields: &mut HashMap<TableReference, BTreeSet<OrdExpr>>,
) -> Result<()> {
Expand All @@ -446,7 +464,7 @@ fn collect_partial_model_plan(
let expr = create_wren_expr_for_model(
&c.name,
dataset.try_as_model().unwrap(),
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
)?;
required_fields
.entry(relation_ref.clone())
Expand All @@ -463,6 +481,7 @@ fn collect_partial_model_plan(
fn collect_model_required_fields(
qualified_column: &Column,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
model_required_fields: &mut HashMap<TableReference, BTreeSet<OrdExpr>>,
) -> Result<()> {
let Some(set) = analyzed_wren_mdl
Expand All @@ -488,7 +507,7 @@ fn collect_model_required_fields(
let Ok(expr) = create_wren_expr_for_model(
expression,
Arc::clone(&m),
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
) else {
// skip the semantic expression (e.g. calculated field or relationship column)
debug!(
Expand All @@ -512,6 +531,7 @@ fn collect_model_required_fields(
&column,
Arc::clone(&m),
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
)?;
debug!("Required field: {}", &expr_plan);
model_required_fields
Expand All @@ -530,11 +550,22 @@ fn get_remote_column_exp(
column: &mdl::manifest::Column,
model: Arc<Model>,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
) -> Result<Expr> {
let expr = if let Some(expression) = &column.expression {
create_remote_expr_for_model(expression, model, analyzed_wren_mdl)?
create_remote_expr_for_model(
expression,
model,
analyzed_wren_mdl,
session_state_ref,
)?
} else {
create_remote_expr_for_model(&column.name, model, analyzed_wren_mdl)?
create_remote_expr_for_model(
&column.name,
model,
analyzed_wren_mdl,
session_state_ref,
)?
};
Ok(expr.alias(column.name.clone()))
}
Expand Down Expand Up @@ -656,6 +687,7 @@ impl ModelSourceNode {
model: Arc<Model>,
required_exprs: Vec<Expr>,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
original_table_scan: Option<LogicalPlan>,
) -> Result<Self> {
let mut required_exprs_buffer = BTreeSet::new();
Expand Down Expand Up @@ -689,6 +721,7 @@ impl ModelSourceNode {
&column,
Arc::clone(&model),
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
)?));
}
} else {
Expand All @@ -711,6 +744,7 @@ impl ModelSourceNode {
&column,
Arc::clone(&model),
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
)?;
required_exprs_buffer.insert(OrdExpr::new(expr_plan.clone()));
}
Expand Down Expand Up @@ -793,7 +827,7 @@ impl CalculationPlanNode {
calculation: ColumnReference,
calculation_expr: Expr,
relation_chain: RelationChain,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
) -> Result<Self> {
let Some(model) = calculation.dataset.try_as_model() else {
return plan_err!("Only support model as source dataset");
Expand Down Expand Up @@ -822,7 +856,7 @@ impl CalculationPlanNode {
let dimensions = vec![create_wren_expr_for_model(
&pk_column.name,
Arc::clone(&model),
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
)?
.alias(pk_column.name())];
let schema_ref = DFSchemaRef::new(
Expand Down
8 changes: 7 additions & 1 deletion wren-modeling-rs/core/src/logical_plan/analyze/relation_chain.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@ use crate::logical_plan::utils::create_schema;
use crate::mdl;
use crate::mdl::lineage::DatasetLink;
use crate::mdl::manifest::JoinType;
use crate::mdl::{AnalyzedWrenMDL, Dataset};
use crate::mdl::Dataset;
use crate::mdl::{AnalyzedWrenMDL, SessionStateRef};
use datafusion::catalog::TableReference;
use datafusion::common::{internal_err, not_impl_err, plan_err, DFSchema, DFSchemaRef};
use datafusion::logical_expr::{
Expand All @@ -33,6 +34,7 @@ impl RelationChain {
dataset: &Dataset,
required_fields: Vec<Expr>,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
) -> datafusion::common::Result<Self> {
match dataset {
Dataset::Model(source_model) => {
Expand All @@ -41,6 +43,7 @@ impl RelationChain {
Arc::clone(source_model),
required_fields,
analyzed_wren_mdl,
session_state_ref,
None,
)?),
})))
Expand All @@ -58,6 +61,7 @@ impl RelationChain {
directed_graph: Graph<Dataset, DatasetLink>,
model_required_fields: &HashMap<TableReference, BTreeSet<OrdExpr>>,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
) -> datafusion::common::Result<Self> {
let mut relation_chain = source;

Expand Down Expand Up @@ -88,6 +92,7 @@ impl RelationChain {
fields.iter().cloned().map(|c| c.expr).collect(),
None,
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
)?;

let df_schema =
Expand All @@ -101,6 +106,7 @@ impl RelationChain {
Arc::clone(target_model),
fields.iter().cloned().map(|c| c.expr).collect(),
Arc::clone(&analyzed_wren_mdl),
Arc::clone(&session_state_ref),
None,
)?),
})
Expand Down
Loading
Loading