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

Code before:

    /// Get expression dtype of expr_node, the schema used is that of the current root node
    fn get_dtype<'py>(&self, expr_node: usize, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let expr_node = Node(expr_node);
        let lp_arena = self.lp_arena.lock().unwrap();
        let schema = lp_arena.get(self.root).schema(&lp_arena);
        let expr_arena = self.expr_arena.lock().unwrap();
        let field = expr_arena
            .get(expr_node)
            .to_field(&schema, &expr_arena)
            .map_err(PyPolarsErr::from)?;
        Wrap(field.dtype).into_pyobject(py)
    }

    /// Set the current node in the plan.
    fn set_node(&mut self, node: usize) {
        self.root = Node(node);
    }

    /// Get the current node in the plan.
    fn get_node(&mut self) -> usize {
        self.root.0
    }

    /// Set a python UDF that will replace the subtree location with this function src.
    fn set_udf(&mut self, function: PyObject) {
        let mut lp_arena = self.lp_arena.lock().unwrap();
        let schema = lp_arena.get(self.root).schema(&lp_arena).into_owned();
        let ir = IR::PythonScan {
            options: PythonOptions {
                scan_fn: Some(function.into()),
                schema,
                output_schema: None,
                with_columns: None,
                python_source: PythonScanSource::Cuda,
                predicate: Default::default(),
                n_rows: None,
                validate_schema: false,
            },
        };
        lp_arena.replace(self.root, ir);
    }

    fn view_current_node(&self, py: Python<'_>) -> PyResult<PyObject> {
        let lp_arena = self.lp_arena.lock().unwrap();
        let lp_node = lp_arena.get(self.root);
        nodes::into_py(py, lp_node)
    }

    fn view_expression(&self, py: Python<'_>, node: usize) -> PyResult<PyObject> {
        let expr_arena = self.expr_arena.lock().unwrap();

Review comment (suggestion):
Used by polars GPU, conservative default to`false` for now. It is also set here in the Rust-side to avoid breaking the Python `set_udf()` function call


Code after:

    /// Get expression dtype of expr_node, the schema used is that of the current root node
    fn get_dtype<'py>(&self, expr_node: usize, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let expr_node = Node(expr_node);
        let lp_arena = self.lp_arena.lock().unwrap();
        let schema = lp_arena.get(self.root).schema(&lp_arena);
        let expr_arena = self.expr_arena.lock().unwrap();
        let field = expr_arena
            .get(expr_node)
            .to_field(&schema, &expr_arena)
            .map_err(PyPolarsErr::from)?;
        Wrap(field.dtype).into_pyobject(py)
    }

    /// Set the current node in the plan.
    fn set_node(&mut self, node: usize) {
        self.root = Node(node);
    }

    /// Get the current node in the plan.
    fn get_node(&mut self) -> usize {
        self.root.0
    }

    /// Set a python UDF that will replace the subtree location with this function src.
    #[pyo3(signature = (function, is_pure = false))]
    fn set_udf(&mut self, function: PyObject, is_pure: bool) {
        let mut lp_arena = self.lp_arena.lock().unwrap();
        let schema = lp_arena.get(self.root).schema(&lp_arena).into_owned();
        let ir = IR::PythonScan {
            options: PythonOptions {
                scan_fn: Some(function.into()),
                schema,
                output_schema: None,
                with_columns: None,
                python_source: PythonScanSource::Cuda,
                predicate: Default::default(),
                n_rows: None,
                validate_schema: false,
                is_pure,
            },
        };
        lp_arena.replace(self.root, ir);
    }

    fn view_current_node(&self, py: Python<'_>) -> PyResult<PyObject> {
        let lp_arena = self.lp_arena.lock().unwrap();
        let lp_node = lp_arena.get(self.root);
        nodes::into_py(py, lp_node)
    }