diff --git a/crates/polars-python/src/lazyframe/visit.rs b/crates/polars-python/src/lazyframe/visit.rs index 726b5e7debd4..32c8d3d23b7d 100644 --- a/crates/polars-python/src/lazyframe/visit.rs +++ b/crates/polars-python/src/lazyframe/visit.rs @@ -57,7 +57,7 @@ impl NodeTraverser { // 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 = (2, 1); + const VERSION: Version = (2, 2); pub fn new(root: Node, lp_arena: Arena, expr_arena: Arena) -> Self { Self { diff --git a/crates/polars-python/src/lazyframe/visitor/expr_nodes.rs b/crates/polars-python/src/lazyframe/visitor/expr_nodes.rs index a05ea6891e15..e8832e9b5488 100644 --- a/crates/polars-python/src/lazyframe/visitor/expr_nodes.rs +++ b/crates/polars-python/src/lazyframe/visitor/expr_nodes.rs @@ -1,4 +1,6 @@ use polars::datatypes::TimeUnit; +#[cfg(feature = "iejoin")] +use polars::prelude::InequalityOperator; use polars::series::ops::NullBehavior; use polars_core::prelude::{NonExistent, QuantileInterpolOptions}; use polars_core::series::IsSorted; @@ -114,6 +116,19 @@ impl IntoPy for Wrap { } } +#[cfg(feature = "iejoin")] +impl IntoPy for Wrap { + fn into_py(self, py: Python<'_>) -> PyObject { + match self.0 { + InequalityOperator::Lt => PyOperator::Lt, + InequalityOperator::LtEq => PyOperator::LtEq, + InequalityOperator::Gt => PyOperator::Gt, + InequalityOperator::GtEq => PyOperator::GtEq, + } + .into_py(py) + } +} + #[pyclass(name = "StringFunction")] #[derive(Copy, Clone)] pub enum PyStringFunction { diff --git a/crates/polars-python/src/lazyframe/visitor/nodes.rs b/crates/polars-python/src/lazyframe/visitor/nodes.rs index d8dbb71281bc..ae805e7d0ff0 100644 --- a/crates/polars-python/src/lazyframe/visitor/nodes.rs +++ b/crates/polars-python/src/lazyframe/visitor/nodes.rs @@ -470,26 +470,40 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult { input_right: input_right.0, left_on: left_on.iter().map(|e| e.into()).collect(), right_on: right_on.iter().map(|e| e.into()).collect(), - options: ( - match options.args.how { - JoinType::Left => "left", - JoinType::Right => "right", - JoinType::Inner => "inner", - JoinType::Full => "full", - #[cfg(feature = "asof_join")] - JoinType::AsOf(_) => return Err(PyNotImplementedError::new_err("asof join")), - JoinType::Cross => "cross", - JoinType::Semi => "leftsemi", - JoinType::Anti => "leftanti", - #[cfg(feature = "iejoin")] - JoinType::IEJoin(_) => return Err(PyNotImplementedError::new_err("IEJoin")), - }, - options.args.join_nulls, - options.args.slice, - options.args.suffix.as_deref(), - options.args.coalesce.coalesce(&options.args.how), - ) - .to_object(py), + options: { + let how = &options.args.how; + + ( + match how { + JoinType::Left => "left".to_object(py), + JoinType::Right => "right".to_object(py), + JoinType::Inner => "inner".to_object(py), + JoinType::Full => "full".to_object(py), + #[cfg(feature = "asof_join")] + JoinType::AsOf(_) => { + return Err(PyNotImplementedError::new_err("asof join")) + }, + JoinType::Cross => "cross".to_object(py), + JoinType::Semi => "leftsemi".to_object(py), + JoinType::Anti => "leftanti".to_object(py), + #[cfg(feature = "iejoin")] + JoinType::IEJoin(ie_options) => ( + "inequality".to_object(py), + crate::Wrap(ie_options.operator1).into_py(py), + ie_options + .operator2 + .as_ref() + .map_or_else(|| py.None(), |op| crate::Wrap(*op).into_py(py)), + ) + .into_py(py), + }, + options.args.join_nulls, + options.args.slice, + options.args.suffix.as_deref(), + options.args.coalesce.coalesce(how), + ) + .to_object(py) + }, } .into_py(py), IR::HStack {