move logic to rust

pola-rs · Jan 18, 2023 · 89015fd · 89015fd
1 parent d2c9f3d
commit 89015fd
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 9 deletions.
diff --git 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
@@ -34,6 +34,7 @@ pub enum TemporalFunction {
         closed: ClosedWindow,
         tz: Option<TimeZone>,
     },
+    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
@@ -459,6 +459,7 @@ impl From<TemporalFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
             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
@@ -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
@@ -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<Series> {
         )),
     }
 }
+
+pub(super) fn combine(s: &[Series], tu: TimeUnit) -> PolarsResult<Series> {
+    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
@@ -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
@@ -1029,6 +1029,10 @@ impl PyExpr {
         self.inner.clone().dt().round(every, offset).into()
     }
 
+    pub fn dt_combine(&self, time: PyExpr, tu: Wrap<TimeUnit>) -> PyExpr {
+        self.inner.clone().dt().combine(time.inner, tu.0).into()
+    }
+
     pub fn rolling_apply(
         &self,
         py: Python,