Skip to content

Commit

Permalink
refactor: Don't mutate arena by default in Rewriting Visitor
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Feb 13, 2025
1 parent ed19c73 commit b85a2cd
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 4 deletions.
1 change: 1 addition & 0 deletions crates/polars-plan/src/plans/optimizer/cse/cse_lp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,7 @@ pub(crate) fn elim_cmn_subplans(

lp_node.visit(&mut visitor, arena).map(|_| ()).unwrap();

let lp_node = IRNode::new_mutate(root);
let mut rewriter = CommonSubPlanRewriter::new(&sp_count, &id_array);
lp_node.rewrite(&mut rewriter, arena).unwrap();

Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/plans/optimizer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ More information on the new streaming engine: https://github.com/pola-rs/polars/
#[cfg(feature = "cse")]
if comm_subexpr_elim && !members.has_ext_context {
let mut optimizer = CommonSubExprOptimizer::new();
let alp_node = IRNode::new(lp_top);
let alp_node = IRNode::new_mutate(lp_top);

lp_top = try_with_ir_arena(lp_arena, expr_arena, |arena| {
let rewritten = alp_node.rewrite(&mut optimizer, arena)?;
Expand Down
21 changes: 18 additions & 3 deletions crates/polars-plan/src/plans/visitor/lp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,21 @@ use crate::prelude::*;
#[derive(Copy, Clone, Debug)]
pub struct IRNode {
node: Node,
// Whether it may mutate the arena on rewrite.
// If set the Rewriting Treewalker will mutate the arena.
mutate: bool,
}

impl IRNode {
pub fn new(node: Node) -> Self {
Self { node }
Self {
node,
mutate: false,
}
}

pub fn new_mutate(node: Node) -> Self {
Self { node, mutate: true }
}

pub fn node(&self) -> Node {
Expand Down Expand Up @@ -85,8 +95,13 @@ impl TreeWalker for IRNode {
}
let lp = arena.0.get(self.node);
let lp = lp.with_exprs_and_input(exprs, inputs);
arena.0.replace(self.node, lp);
Ok(self)
if self.mutate {
arena.0.replace(self.node, lp);
Ok(self)
} else {
let node = arena.0.add(lp);
Ok(IRNode::new_mutate(node))
}
}
}

Expand Down

0 comments on commit b85a2cd

Please sign in to comment.