From 696b33847a79bcf2a0df250d57323014393b9601 Mon Sep 17 00:00:00 2001 From: ritchie Date: Fri, 28 Feb 2025 08:31:40 +0100 Subject: [PATCH] fix: Respect rewriting flag in Node rewriter --- crates/polars-plan/src/plans/optimizer/cse/cse_lp.rs | 2 +- crates/polars-plan/src/plans/visitor/lp.rs | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/polars-plan/src/plans/optimizer/cse/cse_lp.rs b/crates/polars-plan/src/plans/optimizer/cse/cse_lp.rs index 931a905b16ee..4be865fec34e 100644 --- a/crates/polars-plan/src/plans/optimizer/cse/cse_lp.rs +++ b/crates/polars-plan/src/plans/optimizer/cse/cse_lp.rs @@ -359,7 +359,7 @@ pub(crate) fn elim_cmn_subplans( let mut id_array = Default::default(); with_ir_arena(lp_arena, expr_arena, |arena| { - let lp_node = IRNode::new(root); + let lp_node = IRNode::new_mutate(root); let mut visitor = LpIdentifierVisitor::new(&mut sp_count, &mut id_array); lp_node.visit(&mut visitor, arena).map(|_| ()).unwrap(); diff --git a/crates/polars-plan/src/plans/visitor/lp.rs b/crates/polars-plan/src/plans/visitor/lp.rs index 69dd47444da7..87b22a627bd7 100644 --- a/crates/polars-plan/src/plans/visitor/lp.rs +++ b/crates/polars-plan/src/plans/visitor/lp.rs @@ -65,7 +65,8 @@ impl TreeWalker for IRNode { self.to_alp(&arena.0).copy_inputs(&mut scratch); for &node in scratch.as_slice() { - let lp_node = IRNode::new(node); + let mut lp_node = IRNode::new(node); + lp_node.mutate = self.mutate; match op(&lp_node, arena)? { // let the recursion continue VisitRecursion::Continue | VisitRecursion::Skip => {}, @@ -90,7 +91,8 @@ impl TreeWalker for IRNode { // rewrite the nodes for node in &mut inputs { - let lp_node = IRNode::new(*node); + let mut lp_node = IRNode::new(*node); + lp_node.mutate = self.mutate; *node = op(lp_node, arena)?.node; } let lp = arena.0.get(self.node);