From 89015fd65857a77e39e5f9253ec1a0db054c5f31 Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Wed, 18 Jan 2023 09:54:01 +0100 Subject: [PATCH] move logic to rust --- polars/polars-lazy/polars-plan/src/dsl/dt.rs | 8 ++++++++ .../src/dsl/function_expr/datetime.rs | 2 ++ .../polars-plan/src/dsl/function_expr/mod.rs | 1 + .../polars-plan/src/dsl/function_expr/schema.rs | 1 + .../src/dsl/function_expr/temporal.rs | 12 +++++++++++- py-polars/polars/internals/expr/datetime.py | 16 ++++++++-------- py-polars/src/lazy/dsl.rs | 4 ++++ 7 files changed, 35 insertions(+), 9 deletions(-) diff --git a/polars/polars-lazy/polars-plan/src/dsl/dt.rs b/polars/polars-lazy/polars-plan/src/dsl/dt.rs index 458b5bb7a6e30..30e2130b675b2 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/dt.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/dt.rs @@ -222,4 +222,12 @@ impl DateLikeNameSpace { tz, ))) } + + pub fn combine(self, time: Expr, tu: TimeUnit) -> Expr { + self.0.map_many_private( + FunctionExpr::TemporalExpr(TemporalFunction::Combine(tu)), + &[time], + false, + ) + } } diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs index 06043b6ea79b7..5a10d46c98f73 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs @@ -34,6 +34,7 @@ pub enum TemporalFunction { closed: ClosedWindow, tz: Option, }, + Combine(TimeUnit), } impl Display for TemporalFunction { @@ -62,6 +63,7 @@ impl Display for TemporalFunction { #[cfg(feature = "timezones")] TzLocalize(_) => "tz_localize", DateRange { .. } => return write!(f, "date_range"), + Combine(_) => "combine", }; write!(f, "dt.{s}") } diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs index ac1ab55f10eca..f27cbf38c58dd 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs @@ -459,6 +459,7 @@ impl From for SpecialEq> { CastTimezone(tz) => map!(datetime::cast_timezone, &tz), #[cfg(feature = "timezones")] TzLocalize(tz) => map!(datetime::tz_localize, &tz), + Combine(tu) => map_as_slice!(temporal::combine, tu), DateRange { name, every, diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs index 37a84814b5cc9..360a9585387da 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs @@ -159,6 +159,7 @@ impl FunctionExpr { #[cfg(feature = "timezones")] CastTimezone(tz) | TzLocalize(tz) => return cast_tz(tz), DateRange { .. } => return super_type(), + Combine(tu) => DataType::Datetime(*tu, None), }; with_dtype(dtype) } diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs index 6956ecee39237..3d1ac013e538f 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs @@ -1,7 +1,6 @@ #[cfg(feature = "date_offset")] use polars_time::prelude::*; -#[cfg(feature = "date_offset")] use super::*; #[cfg(feature = "date_offset")] @@ -33,3 +32,14 @@ pub(super) fn date_offset(s: Series, offset: Duration) -> PolarsResult { )), } } + +pub(super) fn combine(s: &[Series], tu: TimeUnit) -> PolarsResult { + let date = &s[0]; + let time = &s[1]; + + let date = date.cast(&DataType::Date)?; + let datetime = date.cast(&DataType::Datetime(tu, None)).unwrap(); + + let duration = time.cast(&DataType::Duration(tu))?; + Ok(datetime + duration) +} diff --git a/py-polars/polars/internals/expr/datetime.py b/py-polars/polars/internals/expr/datetime.py index 73590545f19d8..e1265dc24c52b 100644 --- a/py-polars/polars/internals/expr/datetime.py +++ b/py-polars/polars/internals/expr/datetime.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING import polars.internals as pli -from polars.datatypes import DTYPE_TEMPORAL_UNITS, Date, Datetime, Duration, Int32 +from polars.datatypes import DTYPE_TEMPORAL_UNITS, Date, Int32 from polars.utils import _timedelta_to_pl_duration if TYPE_CHECKING: @@ -253,7 +253,7 @@ def round( ) ) - def combine(self, tm: time | pli.Expr) -> pli.Expr: + def combine(self, tm: time | pli.Expr, tu: TimeUnit = "us") -> pli.Expr: """ Create a naive Datetime from an existing Date/Datetime expression and a Time. @@ -264,6 +264,8 @@ def combine(self, tm: time | pli.Expr) -> pli.Expr: ---------- tm A python time literal or polars expression/column that resolves to a time. + tu : {'ns', 'us', 'ms'} + Time unit. Examples -------- @@ -309,10 +311,8 @@ def combine(self, tm: time | pli.Expr) -> pli.Expr: raise TypeError( f"Expected 'tm' to be a python time or polars expression, found {tm!r}" ) - duration = pli.expr_to_lit_or_expr(tm).cast(Duration) - return pli.wrap_expr( - self._pyexpr.cast(Date, True).cast(Datetime, True) + duration._pyexpr - ) + tm = pli.expr_to_lit_or_expr(tm) + return pli.wrap_expr(self._pyexpr.dt_combine(tm._pyexpr, tu)) def strftime(self, fmt: str) -> pli.Expr: """ @@ -905,7 +905,7 @@ def epoch(self, tu: EpochTimeUnit = "us") -> pli.Expr: Parameters ---------- - tu : {'us', 'ns', 'ms', 's', 'd'} + tu : {'ns', 'us', 'ms', 's', 'd'} Time unit. Examples @@ -950,7 +950,7 @@ def timestamp(self, tu: TimeUnit = "us") -> pli.Expr: Parameters ---------- - tu : {'us', 'ns', 'ms'} + tu : {'ns', 'us', 'ms'} Time unit. Examples diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs index 27d5bc4c81fd6..c6350e4a8b7f8 100644 --- a/py-polars/src/lazy/dsl.rs +++ b/py-polars/src/lazy/dsl.rs @@ -1029,6 +1029,10 @@ impl PyExpr { self.inner.clone().dt().round(every, offset).into() } + pub fn dt_combine(&self, time: PyExpr, tu: Wrap) -> PyExpr { + self.inner.clone().dt().combine(time.inner, tu.0).into() + } + pub fn rolling_apply( &self, py: Python,