Skip to content

Commit

Permalink
refactor(python): Expose IEJoin IR node to python (#19104)
Browse files Browse the repository at this point in the history
Co-authored-by: ritchie <[email protected]>
  • Loading branch information
wence- and ritchie46 authored Oct 7, 2024
1 parent 018dfd1 commit addaf83
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 21 deletions.
2 changes: 1 addition & 1 deletion crates/polars-python/src/lazyframe/visit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<IR>, expr_arena: Arena<AExpr>) -> Self {
Self {
Expand Down
15 changes: 15 additions & 0 deletions crates/polars-python/src/lazyframe/visitor/expr_nodes.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -114,6 +116,19 @@ impl IntoPy<PyObject> for Wrap<Operator> {
}
}

#[cfg(feature = "iejoin")]
impl IntoPy<PyObject> for Wrap<InequalityOperator> {
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 {
Expand Down
54 changes: 34 additions & 20 deletions crates/polars-python/src/lazyframe/visitor/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -470,26 +470,40 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
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 {
Expand Down

0 comments on commit addaf83

Please sign in to comment.