diff --git a/crates/polars-pipe/src/pipeline/convert.rs b/crates/polars-pipe/src/pipeline/convert.rs index b27226729ee2..4b33ac87f7fc 100644 --- a/crates/polars-pipe/src/pipeline/convert.rs +++ b/crates/polars-pipe/src/pipeline/convert.rs @@ -379,7 +379,7 @@ where .map(|name| { let name: Arc = Arc::from(name.as_str()); let node = expr_arena.add(AExpr::Column(name.clone())); - ExprIR::new(node, Some(name.clone()), OutputName::Alias(name)) + ExprIR::new(node, OutputName::Alias(name)) }) .collect::>(); let aggs = vec![]; @@ -395,7 +395,7 @@ where group_by_out_schema.with_column(name.clone(), dtype.clone()); let name: Arc = Arc::from(key.as_str()); let node = expr_arena.add(AExpr::Column(name.clone())); - ExprIR::new(node, Some(name.clone()), OutputName::Alias(name)) + ExprIR::new(node, OutputName::Alias(name)) }) .collect::>(); @@ -422,11 +422,7 @@ where unreachable!() }, }; - Some(ExprIR::new( - node, - Some(name.clone()), - OutputName::Alias(name), - )) + Some(ExprIR::new(node, OutputName::Alias(name))) } }) .collect(); diff --git a/crates/polars-plan/src/logical_plan/builder_alp.rs b/crates/polars-plan/src/logical_plan/builder_alp.rs index 38867c922484..74283ad343d1 100644 --- a/crates/polars-plan/src/logical_plan/builder_alp.rs +++ b/crates/polars-plan/src/logical_plan/builder_alp.rs @@ -190,7 +190,6 @@ impl<'a> ALogicalPlanBuilder<'a> { expr_irs.push(ExprIR::new( node, - None, OutputName::ColumnLhs(Arc::from(field.name.as_ref())), )); new_schema.with_column(field.name().clone(), field.data_type().clone()); diff --git a/crates/polars-plan/src/logical_plan/conversion.rs b/crates/polars-plan/src/logical_plan/conversion.rs index 3ac951bfdabb..38e12936bf6f 100644 --- a/crates/polars-plan/src/logical_plan/conversion.rs +++ b/crates/polars-plan/src/logical_plan/conversion.rs @@ -7,13 +7,24 @@ use crate::prelude::*; pub fn to_expr_ir(expr: Expr, arena: &mut Arena) -> ExprIR { let mut state = ConversionState::new(); let node = to_aexpr_impl(expr, arena, &mut state); - ExprIR::new(node, state.left_most_input_name, state.output_name) + ExprIR::new(node, state.output_name) } fn to_expr_irs(input: Vec, arena: &mut Arena) -> Vec { input.convert_owned(|e| to_expr_ir(e, arena)) } +pub fn to_expr_ir_ignore_alias(expr: Expr, arena: &mut Arena) -> ExprIR { + let mut state = ConversionState::new(); + state.ignore_alias = true; + let node = to_aexpr_impl(expr, arena, &mut state); + ExprIR::new(node, state.output_name) +} + +fn to_expr_irs_ignore_alias(input: Vec, arena: &mut Arena) -> Vec { + input.convert_owned(|e| to_expr_ir_ignore_alias(e, arena)) +} + /// converts expression to AExpr and adds it to the arena, which uses an arena (Vec) for allocation pub fn to_aexpr(expr: Expr, arena: &mut Arena) -> Node { to_aexpr_impl( @@ -28,9 +39,11 @@ pub fn to_aexpr(expr: Expr, arena: &mut Arena) -> Node { #[derive(Default)] struct ConversionState { - left_most_input_name: Option>, output_name: OutputName, + /// Remove alias from the expressions and set as [`OutputName`]. prune_alias: bool, + /// If an `alias` is encountered prune and ignore it. + ignore_alias: bool, } impl ConversionState { @@ -55,7 +68,7 @@ fn to_aexpr_impl(expr: Expr, arena: &mut Arena, state: &mut ConversionSta Expr::Explode(expr) => AExpr::Explode(to_aexpr_impl(*expr, arena, state)), Expr::Alias(e, name) => { if state.prune_alias { - if state.output_name.is_none() { + if state.output_name.is_none() && !state.ignore_alias { state.output_name = OutputName::Alias(name); } to_aexpr_impl(*e, arena, state); @@ -78,9 +91,6 @@ fn to_aexpr_impl(expr: Expr, arena: &mut Arena, state: &mut ConversionSta if state.output_name.is_none() { state.output_name = OutputName::ColumnLhs(name.clone()) } - if state.left_most_input_name.is_none() { - state.left_most_input_name = Some(name.clone()) - } AExpr::Column(name) }, Expr::BinaryExpr { left, op, right } => { @@ -394,8 +404,8 @@ pub fn to_alp( let input_left = to_alp(*input_left, expr_arena, lp_arena)?; let input_right = to_alp(*input_right, expr_arena, lp_arena)?; - let left_on = to_expr_irs(left_on, expr_arena); - let right_on = to_expr_irs(right_on, expr_arena); + let left_on = to_expr_irs_ignore_alias(left_on, expr_arena); + let right_on = to_expr_irs_ignore_alias(right_on, expr_arena); ALogicalPlan::Join { input_left, diff --git a/crates/polars-plan/src/logical_plan/expr_ir.rs b/crates/polars-plan/src/logical_plan/expr_ir.rs index 924dfa9bd989..633a359dfd3b 100644 --- a/crates/polars-plan/src/logical_plan/expr_ir.rs +++ b/crates/polars-plan/src/logical_plan/expr_ir.rs @@ -29,9 +29,6 @@ impl OutputName { #[derive(Clone, Debug)] pub struct ExprIR { - /// Name that this expression refers to via `col()` - /// This is `None` for literals. - left_most_input_name: Option>, /// Output name of this expression. output_name: OutputName, /// Reduced expression. @@ -40,23 +37,14 @@ pub struct ExprIR { } impl ExprIR { - pub fn new( - node: Node, - left_most_input_name: Option>, - output_name: OutputName, - ) -> Self { + pub fn new(node: Node, output_name: OutputName) -> Self { debug_assert!(!output_name.is_none()); - ExprIR { - left_most_input_name, - output_name, - node, - } + ExprIR { output_name, node } } pub fn from_node(node: Node, arena: &Arena) -> Self { let mut out = Self { node, - left_most_input_name: None, output_name: OutputName::None, }; out.node = node; @@ -108,14 +96,6 @@ impl ExprIR { self.output_name_arc().as_ref() } - pub(crate) fn left_most_input_name_arc(&self) -> &Name { - self.left_most_input_name.as_ref().unwrap() - } - - pub(crate) fn left_most_input_name(&self) -> &str { - self.left_most_input_name_arc().as_ref() - } - pub(crate) fn to_expr(&self, expr_arena: &Arena) -> Expr { let out = node_to_expr(self.node, expr_arena); @@ -169,7 +149,7 @@ impl From<&ExprIR> for Node { pub(crate) fn name_to_expr_ir(name: &str, expr_arena: &mut Arena) -> ExprIR { let name: Name = Arc::from(name); let node = expr_arena.add(AExpr::Column(name.clone())); - ExprIR::new(node, Some(name.clone()), OutputName::ColumnLhs(name)) + ExprIR::new(node, OutputName::ColumnLhs(name)) } pub(crate) fn names_to_expr_irs, S: AsRef>( diff --git a/crates/polars-plan/src/logical_plan/optimizer/cse_expr.rs b/crates/polars-plan/src/logical_plan/optimizer/cse_expr.rs index 2dabc1de03ed..9c29eb869126 100644 --- a/crates/polars-plan/src/logical_plan/optimizer/cse_expr.rs +++ b/crates/polars-plan/src/logical_plan/optimizer/cse_expr.rs @@ -688,7 +688,7 @@ impl<'a> CommonSubExprOptimizer<'a> { for id in &self.replaced_identifiers { let (node, _count) = self.se_count.get(id).unwrap(); let name = id.materialize(); - let out_e = ExprIR::new(*node, None, OutputName::Alias(Arc::from(name))); + let out_e = ExprIR::new(*node, OutputName::Alias(Arc::from(name))); new_expr.push(out_e) } let expr = ProjectionExprs::new_with_cse(new_expr, self.replaced_identifiers.len()); diff --git a/crates/polars-plan/src/logical_plan/optimizer/projection_pushdown/joins.rs b/crates/polars-plan/src/logical_plan/optimizer/projection_pushdown/joins.rs index 5ac8b9fcee05..e22fa8b4e499 100644 --- a/crates/polars-plan/src/logical_plan/optimizer/projection_pushdown/joins.rs +++ b/crates/polars-plan/src/logical_plan/optimizer/projection_pushdown/joins.rs @@ -259,7 +259,7 @@ pub(super) fn process_join( // In case of outer joins we also add the columns. // But before we do that we must check if the column wasn't already added by the lhs. let add_local = if add_local { - !already_added_local_to_local_projected.contains(e.left_most_input_name()) + !already_added_local_to_local_projected.contains(e.output_name()) } else { false }; @@ -439,13 +439,9 @@ fn resolve_join_suffixes( let downstream_name = &name.as_ref()[..name.len() - suffix.len()]; let col = AExpr::Column(Arc::from(downstream_name)); let node = expr_arena.add(col); - ExprIR::new( - node, - Some(Arc::from(downstream_name)), - OutputName::Alias(name), - ) + ExprIR::new(node, OutputName::Alias(name)) } else { - ExprIR::new(proj.0, Some(name.clone()), OutputName::ColumnLhs(name)) + ExprIR::new(proj.0, OutputName::ColumnLhs(name)) } }) .collect::>();