fix: Respect rewriting flag in Node rewriter (#21516)

pola-rs · Feb 28, 2025 · 69612d4 · 69612d4
1 parent e801b83
commit 69612d4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 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
@@ -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
@@ -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);