diff --git a/crates/polars-plan/src/dsl/function_expr/mod.rs b/crates/polars-plan/src/dsl/function_expr/mod.rs index 8f649a6535a2..b46f71a77fa4 100644 --- a/crates/polars-plan/src/dsl/function_expr/mod.rs +++ b/crates/polars-plan/src/dsl/function_expr/mod.rs @@ -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; diff --git a/py-polars/src/lazyframe/visitor/expr_nodes.rs b/py-polars/src/lazyframe/visitor/expr_nodes.rs index 51e5dcd459cc..04102812d3e2 100644 --- a/py-polars/src/lazyframe/visitor/expr_nodes.rs +++ b/py-polars/src/lazyframe/visitor/expr_nodes.rs @@ -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, @@ -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 for Wrap { + fn into_py(self, py: Python<'_>) -> PyObject { + self.0.to_ascii().into_py(py) + } +} + #[pyclass] pub struct BinaryExpr { #[pyo3(get)] @@ -560,7 +622,7 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult { options: py.None(), }, IRAggExpr::NUnique(n) => Agg { - name: "nunique".to_object(py), + name: "n_unique".to_object(py), arguments: n.0, options: py.None(), }, @@ -788,8 +850,99 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult { 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 } => { @@ -980,20 +1133,12 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult { 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")) @@ -1060,7 +1205,7 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult { 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")) diff --git a/py-polars/src/lib.rs b/py-polars/src/lib.rs index 1a4aebb878d9..bf41d6846aad 100644 --- a/py-polars/src/lib.rs +++ b/py-polars/src/lib.rs @@ -107,6 +107,7 @@ fn _expr_nodes(_py: Python, m: &Bound) -> PyResult<()> { m.add_class::().unwrap(); m.add_class::().unwrap(); m.add_class::().unwrap(); + m.add_class::().unwrap(); // Options m.add_class::().unwrap(); m.add_class::().unwrap();