Repository: pola-rs/polars
File: crates/polars-python/src/lazyframe/visit.rs (Rust)

Code before:
    fn from(value: &ExprIR) -> Self {
        Self {
            node: value.node().0,
            output_name: value.output_name().to_string(),
        }
    }
}

type Version = (u16, u16);

#[pyclass]
pub struct NodeTraverser {
    root: Node,
    lp_arena: Arc<Mutex<Arena<IR>>>,
    expr_arena: Arc<Mutex<Arena<AExpr>>>,
    scratch: Vec<Node>,
    expr_scratch: Vec<ExprIR>,
    expr_mapping: Option<Vec<Node>>,
}

impl NodeTraverser {
    // Versioning for IR, (major, minor)
    // Increment major on breaking changes to the IR (e.g. renaming
    // fields, reordering tuples), minor on backwards compatible
    // changes (e.g. exposing a new expression node).
    const VERSION: Version = (9, 1);

    pub fn new(root: Node, lp_arena: Arena<IR>, expr_arena: Arena<AExpr>) -> Self {
        Self {
            root,
            lp_arena: Arc::new(Mutex::new(lp_arena)),
            expr_arena: Arc::new(Mutex::new(expr_arena)),
            scratch: vec![],
            expr_scratch: vec![],
            expr_mapping: None,
        }
    }

    #[allow(clippy::type_complexity)]
    pub fn get_arenas(&self) -> (Arc<Mutex<Arena<IR>>>, Arc<Mutex<Arena<AExpr>>>) {
        (self.lp_arena.clone(), self.expr_arena.clone())
    }

    fn fill_inputs(&mut self) {
        let lp_arena = self.lp_arena.lock().unwrap();
        let this_node = lp_arena.get(self.root);
        self.scratch.clear();
        this_node.copy_inputs(&mut self.scratch);
    }

    fn fill_expressions(&mut self) {

Review comment (performance):
*nit*: Since this changes the `Cache` node, please do a major version bump here.

Code after:
    fn from(value: &ExprIR) -> Self {
        Self {
            node: value.node().0,
            output_name: value.output_name().to_string(),
        }
    }
}

type Version = (u16, u16);

#[pyclass]
pub struct NodeTraverser {
    root: Node,
    lp_arena: Arc<Mutex<Arena<IR>>>,
    expr_arena: Arc<Mutex<Arena<AExpr>>>,
    scratch: Vec<Node>,
    expr_scratch: Vec<ExprIR>,
    expr_mapping: Option<Vec<Node>>,
}

impl NodeTraverser {
    // Versioning for IR, (major, minor)
    // Increment major on breaking changes to the IR (e.g. renaming
    // fields, reordering tuples), minor on backwards compatible
    // changes (e.g. exposing a new expression node).
    const VERSION: Version = (10, 0);

    pub fn new(root: Node, lp_arena: Arena<IR>, expr_arena: Arena<AExpr>) -> Self {
        Self {
            root,
            lp_arena: Arc::new(Mutex::new(lp_arena)),
            expr_arena: Arc::new(Mutex::new(expr_arena)),
            scratch: vec![],
            expr_scratch: vec![],
            expr_mapping: None,
        }
    }

    #[allow(clippy::type_complexity)]
    pub fn get_arenas(&self) -> (Arc<Mutex<Arena<IR>>>, Arc<Mutex<Arena<AExpr>>>) {
        (self.lp_arena.clone(), self.expr_arena.clone())
    }

    fn fill_inputs(&mut self) {
        let lp_arena = self.lp_arena.lock().unwrap();
        let this_node = lp_arena.get(self.root);
        self.scratch.clear();
        this_node.copy_inputs(&mut self.scratch);
    }

    fn fill_expressions(&mut self) {