Skip to content

Commit

Permalink
feat(python): Expose temporal function expression ops to expr ir (#16546
Browse files Browse the repository at this point in the history 
)
  • Loading branch information
@wence-
wence- authored May 28, 2024
1 parent d53bcb2 commit 52c919f
Show file tree
Hide file tree
Showing 3 changed files with 164 additions and 18 deletions.
2 changes: 1 addition & 1 deletion crates/polars-plan/src/dsl/function_expr/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ pub(super) use self::business::BusinessFunction;
#[cfg(feature = "dtype-categorical")]
pub use self::cat::CategoricalFunction;
#[cfg(feature = "temporal")]
pub(super) use self::datetime::TemporalFunction;
pub use self::datetime::TemporalFunction;
pub use self::pow::PowFunction;
#[cfg(feature = "range")]
pub(super) use self::range::RangeFunction;
Expand Down
179 changes: 162 additions & 17 deletions py-polars/src/lazyframe/visitor/expr_nodes.rs
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
use polars::datatypes::TimeUnit;
use polars_core::series::IsSorted;
use polars_core::utils::arrow::legacy::kernels::NonExistent;
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;
use polars_plan::dsl::{BooleanFunction, StringFunction};
use polars_plan::dsl::{BooleanFunction, StringFunction, TemporalFunction};
use polars_plan::prelude::{
AExpr, FunctionExpr, GroupbyOptions, IRAggExpr, LiteralValue, Operator, PowFunction,
WindowMapping, WindowType,
Expand Down Expand Up @@ -191,6 +193,66 @@ impl PyBooleanFunction {
}
}

#[pyclass(name = "TemporalFunction")]
#[derive(Copy, Clone)]
pub enum PyTemporalFunction {
Millennium,
Century,
Year,
IsLeapYear,
IsoYear,
Quarter,
Month,
Week,
WeekDay,
Day,
OrdinalDay,
Time,
Date,
Datetime,
Duration,
Hour,
Minute,
Second,
Millisecond,
Microsecond,
Nanosecond,
TotalDays,
TotalHours,
TotalMinutes,
TotalSeconds,
TotalMilliseconds,
TotalMicroseconds,
TotalNanoseconds,
ToString,
CastTimeUnit,
WithTimeUnit,
ConvertTimeZone,
TimeStamp,
Truncate,
MonthStart,
MonthEnd,
BaseUtcOffset,
DSTOffset,
Round,
ReplaceTimeZone,
Combine,
DatetimeFunction,
}

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

impl IntoPy<PyObject> for Wrap<TimeUnit> {
fn into_py(self, py: Python<'_>) -> PyObject {
self.0.to_ascii().into_py(py)
}
}

#[pyclass]
pub struct BinaryExpr {
#[pyo3(get)]
Expand Down Expand Up @@ -560,7 +622,7 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
options: py.None(),
},
IRAggExpr::NUnique(n) => Agg {
name: "nunique".to_object(py),
name: "n_unique".to_object(py),
arguments: n.0,
options: py.None(),
},
Expand Down Expand Up @@ -788,8 +850,99 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
FunctionExpr::StructExpr(_) => {
return Err(PyNotImplementedError::new_err("struct expr"))
},
FunctionExpr::TemporalExpr(_) => {
return Err(PyNotImplementedError::new_err("temporal expr"))
FunctionExpr::TemporalExpr(fun) => match fun {
TemporalFunction::Millennium => (PyTemporalFunction::Millennium,).into_py(py),
TemporalFunction::Century => (PyTemporalFunction::Century,).into_py(py),
TemporalFunction::Year => (PyTemporalFunction::Year,).into_py(py),
TemporalFunction::IsLeapYear => (PyTemporalFunction::IsLeapYear,).into_py(py),
TemporalFunction::IsoYear => (PyTemporalFunction::IsoYear,).into_py(py),
TemporalFunction::Quarter => (PyTemporalFunction::Quarter,).into_py(py),
TemporalFunction::Month => (PyTemporalFunction::Month,).into_py(py),
TemporalFunction::Week => (PyTemporalFunction::Week,).into_py(py),
TemporalFunction::WeekDay => (PyTemporalFunction::WeekDay,).into_py(py),
TemporalFunction::Day => (PyTemporalFunction::Day,).into_py(py),
TemporalFunction::OrdinalDay => (PyTemporalFunction::OrdinalDay,).into_py(py),
TemporalFunction::Time => (PyTemporalFunction::Time,).into_py(py),
TemporalFunction::Date => (PyTemporalFunction::Date,).into_py(py),
TemporalFunction::Datetime => (PyTemporalFunction::Datetime,).into_py(py),
TemporalFunction::Duration(time_unit) => {
(PyTemporalFunction::Duration, Wrap(*time_unit)).into_py(py)
},
TemporalFunction::Hour => (PyTemporalFunction::Hour,).into_py(py),
TemporalFunction::Minute => (PyTemporalFunction::Minute,).into_py(py),
TemporalFunction::Second => (PyTemporalFunction::Second,).into_py(py),
TemporalFunction::Millisecond => (PyTemporalFunction::Millisecond,).into_py(py),
TemporalFunction::Microsecond => (PyTemporalFunction::Microsecond,).into_py(py),
TemporalFunction::Nanosecond => (PyTemporalFunction::Nanosecond,).into_py(py),
TemporalFunction::TotalDays => (PyTemporalFunction::TotalDays,).into_py(py),
TemporalFunction::TotalHours => (PyTemporalFunction::TotalHours,).into_py(py),
TemporalFunction::TotalMinutes => {
(PyTemporalFunction::TotalMinutes,).into_py(py)
},
TemporalFunction::TotalSeconds => {
(PyTemporalFunction::TotalSeconds,).into_py(py)
},
TemporalFunction::TotalMilliseconds => {
(PyTemporalFunction::TotalMilliseconds,).into_py(py)
},
TemporalFunction::TotalMicroseconds => {
(PyTemporalFunction::TotalMicroseconds,).into_py(py)
},
TemporalFunction::TotalNanoseconds => {
(PyTemporalFunction::TotalNanoseconds,).into_py(py)
},
TemporalFunction::ToString(format) => {
(PyTemporalFunction::ToString, format).into_py(py)
},
TemporalFunction::CastTimeUnit(time_unit) => {
(PyTemporalFunction::CastTimeUnit, Wrap(*time_unit)).into_py(py)
},
TemporalFunction::WithTimeUnit(time_unit) => {
(PyTemporalFunction::WithTimeUnit, Wrap(*time_unit)).into_py(py)
},
TemporalFunction::ConvertTimeZone(time_zone) => {
(PyTemporalFunction::ConvertTimeZone, time_zone).into_py(py)
},
TemporalFunction::TimeStamp(time_unit) => {
(PyTemporalFunction::TimeStamp, Wrap(*time_unit)).into_py(py)
},
TemporalFunction::Truncate(bucket) => {
(PyTemporalFunction::Truncate, bucket).into_py(py)
},
TemporalFunction::MonthStart => (PyTemporalFunction::MonthStart,).into_py(py),
TemporalFunction::MonthEnd => (PyTemporalFunction::MonthEnd,).into_py(py),
TemporalFunction::BaseUtcOffset => {
(PyTemporalFunction::BaseUtcOffset,).into_py(py)
},
TemporalFunction::DSTOffset => (PyTemporalFunction::DSTOffset,).into_py(py),
TemporalFunction::Round(bucket) => {
(PyTemporalFunction::Round, bucket).into_py(py)
},
TemporalFunction::ReplaceTimeZone(time_zone, non_existent) => (
PyTemporalFunction::ReplaceTimeZone,
time_zone
.as_ref()
.map_or_else(|| py.None(), |s| s.to_object(py)),
match non_existent {
NonExistent::Null => "nullify",
NonExistent::Raise => "raise",
},
)
.into_py(py),
TemporalFunction::Combine(time_unit) => {
(PyTemporalFunction::Combine, Wrap(*time_unit)).into_py(py)
},
TemporalFunction::DatetimeFunction {
time_unit,
time_zone,
} => (
PyTemporalFunction::DatetimeFunction,
Wrap(*time_unit),
time_zone
.as_ref()
.map_or_else(|| py.None(), |s| s.to_object(py)),
)
.into_py(py),
},
FunctionExpr::Boolean(boolfun) => match boolfun {
BooleanFunction::Any { ignore_nulls } => {
Expand Down Expand Up @@ -980,20 +1133,12 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
FunctionExpr::Log1p => return Err(PyNotImplementedError::new_err("log1p")),
FunctionExpr::Exp => return Err(PyNotImplementedError::new_err("exp")),
FunctionExpr::Unique(_) => return Err(PyNotImplementedError::new_err("unique")),
FunctionExpr::Round { decimals: _ } => {
return Err(PyNotImplementedError::new_err("round"))
},
FunctionExpr::RoundSF { digits: _ } => {
return Err(PyNotImplementedError::new_err("round sf"))
},
FunctionExpr::Round { decimals } => ("round", decimals).to_object(py),
FunctionExpr::RoundSF { digits } => ("round_sig_figs", digits).to_object(py),
FunctionExpr::Floor => ("floor",).to_object(py),
FunctionExpr::Ceil => ("ceil",).to_object(py),
FunctionExpr::UpperBound => {
return Err(PyNotImplementedError::new_err("upper bound"))
},
FunctionExpr::LowerBound => {
return Err(PyNotImplementedError::new_err("lower bound"))
},
FunctionExpr::UpperBound => ("upper_bound",).to_object(py),
FunctionExpr::LowerBound => ("lower_bound",).to_object(py),
FunctionExpr::Fused(_) => return Err(PyNotImplementedError::new_err("fused")),
FunctionExpr::ConcatExpr(_) => {
return Err(PyNotImplementedError::new_err("concat expr"))
Expand Down Expand Up @@ -1060,7 +1205,7 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
return Err(PyNotImplementedError::new_err("fill null with strategy"))
},
FunctionExpr::GatherEvery { n, offset } => {
("strided_slice", offset, n).to_object(py)
("gather_every", offset, n).to_object(py)
},
FunctionExpr::Reinterpret(_) => {
return Err(PyNotImplementedError::new_err("reinterpret"))
Expand Down
1 change: 1 addition & 0 deletions py-polars/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ fn _expr_nodes(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
m.add_class::<PyOperator>().unwrap();
m.add_class::<PyStringFunction>().unwrap();
m.add_class::<PyBooleanFunction>().unwrap();
m.add_class::<PyTemporalFunction>().unwrap();
// Options
m.add_class::<PyWindowMapping>().unwrap();
m.add_class::<PyRollingGroupOptions>().unwrap();
Expand Down

0 comments on commit 52c919f

Please sign in to comment.