Skip to content

Commit 5e9cf5c

Browse files
authored
feat(rust): Expose IRFunctionExpr::Fused in the python node visitor (#27889)
1 parent a65aa6e commit 5e9cf5c

2 files changed

Lines changed: 11 additions & 4 deletions

File tree

crates/polars-python/src/lazyframe/visit.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ impl NodeTraverser {
5858
// Increment major on breaking changes to the IR (e.g. renaming
5959
// fields, reordering tuples), minor on backwards compatible
6060
// changes (e.g. exposing a new expression node).
61-
const VERSION: Version = (14, 2);
61+
const VERSION: Version = (14, 3);
6262

6363
pub fn new(root: Node, lp_arena: Arena<IR>, expr_arena: Arena<AExpr>) -> Self {
6464
Self {

crates/polars-python/src/lazyframe/visitor/expr_nodes.rs

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ use polars_ops::series::InterpolationMethod;
1010
#[cfg(feature = "search_sorted")]
1111
use polars_ops::series::SearchSortedSide;
1212
use polars_plan::plans::{
13-
DynLiteralValue, IRBooleanFunction, IRFunctionExpr, IRPowFunction, IRRollingFunction,
14-
IRRollingFunctionBy, IRStringFunction, IRStructFunction, IRTemporalFunction,
13+
DynLiteralValue, FusedOperator, IRBooleanFunction, IRFunctionExpr, IRPowFunction,
14+
IRRollingFunction, IRRollingFunctionBy, IRStringFunction, IRStructFunction, IRTemporalFunction,
1515
};
1616
use polars_plan::prelude::{
1717
AExpr, GroupbyOptions, IRAggExpr, LiteralValue, Operator, WindowMapping,
@@ -1427,7 +1427,14 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<Py<PyAny>> {
14271427
IRFunctionExpr::Truncate { decimals } => ("truncate", decimals).into_py_any(py),
14281428
IRFunctionExpr::Floor => ("floor",).into_py_any(py),
14291429
IRFunctionExpr::Ceil => ("ceil",).into_py_any(py),
1430-
IRFunctionExpr::Fused(_) => return Err(PyNotImplementedError::new_err("fused")),
1430+
IRFunctionExpr::Fused(op) => {
1431+
let op_name = match op {
1432+
FusedOperator::MultiplyAdd => "fma",
1433+
FusedOperator::SubMultiply => "fsm",
1434+
FusedOperator::MultiplySub => "fms",
1435+
};
1436+
("fused", op_name).into_py_any(py)
1437+
},
14311438
IRFunctionExpr::ConcatExpr { .. } => {
14321439
return Err(PyNotImplementedError::new_err("concat expr"));
14331440
},

0 commit comments

Comments
 (0)