Skip to content

Commit

Permalink
feat(python): expose BooleanFunction in expr IR (#16355)
Browse files Browse the repository at this point in the history
  • Loading branch information
@wence-
wence- committed May 21, 2024
1 parent 7b155c4 commit 6289dc3
Show file tree
Hide file tree
Showing 2 changed files with 125 additions and 57 deletions.
180 changes: 123 additions & 57 deletions py-polars/src/lazyframe/visitor/expr_nodes.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use polars_core::series::IsSorted;
use polars_ops::prelude::ClosedInterval;
use polars_plan::dsl::function_expr::rolling::RollingFunction;
use polars_plan::dsl::function_expr::rolling_by::RollingFunctionBy;
use polars_plan::dsl::function_expr::trigonometry::TrigonometricFunction;
Expand Down Expand Up @@ -35,7 +36,29 @@ pub struct Literal {
dtype: PyObject,
}

#[pyclass]
#[pyclass(name = "ClosedInterval")]
#[derive(Copy, Clone)]
pub enum PyClosedInterval {
Both,
Left,
Right,
None,
}

impl IntoPy<PyObject> for Wrap<ClosedInterval> {
fn into_py(self, py: Python<'_>) -> PyObject {
match self.0 {
ClosedInterval::Both => PyClosedInterval::Both,
ClosedInterval::Left => PyClosedInterval::Left,
ClosedInterval::Right => PyClosedInterval::Right,
ClosedInterval::None => PyClosedInterval::None,
}
.into_py(py)
}
}

#[pyclass(name = "Operator")]
#[derive(Copy, Clone)]
pub enum PyOperator {
Eq,
EqValidity,
Expand All @@ -59,7 +82,43 @@ pub enum PyOperator {
LogicalOr,
}

#[pymethods]
impl PyOperator {
fn __hash__(&self) -> isize {
*self as isize
}
}

impl IntoPy<PyObject> for Wrap<Operator> {
fn into_py(self, py: Python<'_>) -> PyObject {
match self.0 {
Operator::Eq => PyOperator::Eq,
Operator::EqValidity => PyOperator::EqValidity,
Operator::NotEq => PyOperator::NotEq,
Operator::NotEqValidity => PyOperator::NotEqValidity,
Operator::Lt => PyOperator::Lt,
Operator::LtEq => PyOperator::LtEq,
Operator::Gt => PyOperator::Gt,
Operator::GtEq => PyOperator::GtEq,
Operator::Plus => PyOperator::Plus,
Operator::Minus => PyOperator::Minus,
Operator::Multiply => PyOperator::Multiply,
Operator::Divide => PyOperator::Divide,
Operator::TrueDivide => PyOperator::TrueDivide,
Operator::FloorDivide => PyOperator::FloorDivide,
Operator::Modulus => PyOperator::Modulus,
Operator::And => PyOperator::And,
Operator::Or => PyOperator::Or,
Operator::Xor => PyOperator::Xor,
Operator::LogicalAnd => PyOperator::LogicalAnd,
Operator::LogicalOr => PyOperator::LogicalOr,
}
.into_py(py)
}
}

#[pyclass(name = "StringFunction")]
#[derive(Copy, Clone)]
pub enum PyStringFunction {
ConcatHorizontal,
ConcatVertical,
Expand Down Expand Up @@ -107,59 +166,38 @@ pub enum PyStringFunction {
}

#[pymethods]
impl PyOperator {
fn __hash__(&self) -> u64 {
use PyOperator::*;
match self {
Eq => Eq as u64,
EqValidity => EqValidity as u64,
NotEq => NotEq as u64,
NotEqValidity => NotEqValidity as u64,
Lt => Lt as u64,
LtEq => LtEq as u64,
Gt => Gt as u64,
GtEq => GtEq as u64,
Plus => Plus as u64,
Minus => Minus as u64,
Multiply => Multiply as u64,
Divide => Divide as u64,
TrueDivide => TrueDivide as u64,
FloorDivide => FloorDivide as u64,
Modulus => Modulus as u64,
And => And as u64,
Or => Or as u64,
Xor => Xor as u64,
LogicalAnd => LogicalAnd as u64,
LogicalOr => LogicalOr as u64,
}
impl PyStringFunction {
fn __hash__(&self) -> isize {
*self as isize
}
}

impl IntoPy<PyObject> for Wrap<Operator> {
fn into_py(self, py: Python<'_>) -> PyObject {
match self.0 {
Operator::Eq => PyOperator::Eq,
Operator::EqValidity => PyOperator::EqValidity,
Operator::NotEq => PyOperator::NotEq,
Operator::NotEqValidity => PyOperator::NotEqValidity,
Operator::Lt => PyOperator::Lt,
Operator::LtEq => PyOperator::LtEq,
Operator::Gt => PyOperator::Gt,
Operator::GtEq => PyOperator::GtEq,
Operator::Plus => PyOperator::Plus,
Operator::Minus => PyOperator::Minus,
Operator::Multiply => PyOperator::Multiply,
Operator::Divide => PyOperator::Divide,
Operator::TrueDivide => PyOperator::TrueDivide,
Operator::FloorDivide => PyOperator::FloorDivide,
Operator::Modulus => PyOperator::Modulus,
Operator::And => PyOperator::And,
Operator::Or => PyOperator::Or,
Operator::Xor => PyOperator::Xor,
Operator::LogicalAnd => PyOperator::LogicalAnd,
Operator::LogicalOr => PyOperator::LogicalOr,
}
.into_py(py)
#[pyclass(name = "BooleanFunction")]
#[derive(Copy, Clone)]
pub enum PyBooleanFunction {
Any,
All,
IsNull,
IsNotNull,
IsFinite,
IsInfinite,
IsNan,
IsNotNan,
IsFirstDistinct,
IsLastDistinct,
IsUnique,
IsDuplicated,
IsBetween,
IsIn,
AllHorizontal,
AnyHorizontal,
Not,
}

#[pymethods]
impl PyBooleanFunction {
fn __hash__(&self) -> isize {
*self as isize
}
}

Expand Down Expand Up @@ -265,7 +303,7 @@ pub struct Window {
options: PyObject,
}

#[pyclass]
#[pyclass(name = "WindowMapping")]
pub struct PyWindowMapping {
inner: WindowMapping,
}
Expand All @@ -283,7 +321,7 @@ impl PyWindowMapping {
}
}

#[pyclass]
#[pyclass(name = "RollingGroupOptions")]
pub struct PyRollingGroupOptions {
inner: RollingGroupOptions,
}
Expand Down Expand Up @@ -338,7 +376,7 @@ impl PyRollingGroupOptions {
}
}

#[pyclass]
#[pyclass(name = "GroupbyOptions")]
pub struct PyGroupbyOptions {
inner: GroupbyOptions,
}
Expand Down Expand Up @@ -772,9 +810,37 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
return Err(PyNotImplementedError::new_err("temporal expr"))
},
FunctionExpr::Boolean(boolfun) => match boolfun {
BooleanFunction::IsNull => ("is_null",).to_object(py),
BooleanFunction::IsNotNull => ("is_not_null",).to_object(py),
_ => return Err(PyNotImplementedError::new_err("boolean expr")),
BooleanFunction::Any { ignore_nulls } => {
(PyBooleanFunction::Any, *ignore_nulls).into_py(py)
},
BooleanFunction::All { ignore_nulls } => {
(PyBooleanFunction::All, *ignore_nulls).into_py(py)
},
BooleanFunction::IsNull => (PyBooleanFunction::IsNull,).into_py(py),
BooleanFunction::IsNotNull => (PyBooleanFunction::IsNotNull,).into_py(py),
BooleanFunction::IsFinite => (PyBooleanFunction::IsFinite,).into_py(py),
BooleanFunction::IsInfinite => (PyBooleanFunction::IsInfinite,).into_py(py),
BooleanFunction::IsNan => (PyBooleanFunction::IsNan,).into_py(py),
BooleanFunction::IsNotNan => (PyBooleanFunction::IsNotNan,).into_py(py),
BooleanFunction::IsFirstDistinct => {
(PyBooleanFunction::IsFirstDistinct,).into_py(py)
},
BooleanFunction::IsLastDistinct => {
(PyBooleanFunction::IsLastDistinct,).into_py(py)
},
BooleanFunction::IsUnique => (PyBooleanFunction::IsUnique,).into_py(py),
BooleanFunction::IsDuplicated => (PyBooleanFunction::IsDuplicated,).into_py(py),
BooleanFunction::IsBetween { closed } => {
(PyBooleanFunction::IsBetween, Wrap(*closed)).into_py(py)
},
BooleanFunction::IsIn => (PyBooleanFunction::IsIn,).into_py(py),
BooleanFunction::AllHorizontal => {
(PyBooleanFunction::AllHorizontal,).into_py(py)
},
BooleanFunction::AnyHorizontal => {
(PyBooleanFunction::AnyHorizontal,).into_py(py)
},
BooleanFunction::Not => (PyBooleanFunction::Not,).into_py(py),
},
FunctionExpr::Abs => ("abs",).to_object(py),
FunctionExpr::Hist { .. } => return Err(PyNotImplementedError::new_err("hist")),
Expand Down
2 changes: 2 additions & 0 deletions py-polars/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,8 @@ fn _expr_nodes(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
m.add_class::<Window>().unwrap();
m.add_class::<PyOperator>().unwrap();
m.add_class::<PyStringFunction>().unwrap();
m.add_class::<PyBooleanFunction>().unwrap();
m.add_class::<PyClosedInterval>().unwrap();
// Options
m.add_class::<PyWindowMapping>().unwrap();
m.add_class::<PyRollingGroupOptions>().unwrap();
Expand Down

0 comments on commit 6289dc3

Please sign in to comment.