feat[rust]: add iso_year and make dt computed attributes serde (#4873)

polars/polars-lazy/src/dsl/dt.rs

@@ -1,6 +1,7 @@
 use polars_time::prelude::TemporalMethods;
 
 use super::*;
+use crate::prelude::function_expr::TemporalFunction;
 
 /// Specialized expressions for [`Series`] with dates/datetimes.
 pub struct DateLikeNameSpace(pub(crate) Expr);
@@ -83,27 +84,28 @@ impl DateLikeNameSpace {
 
     /// Get the year of a Date/Datetime
     pub fn year(self) -> Expr {
-        let function = move |s: Series| s.year().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.year")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Year))
+    }
+
+    /// Get the iso-year of a Date/Datetime.
+    /// This may not correspond with a calendar year.
+    pub fn iso_year(self) -> Expr {
+        self.0
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::IsoYear))
     }
 
     /// Get the month of a Date/Datetime
     pub fn month(self) -> Expr {
-        let function = move |s: Series| s.month().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.month")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Month))
     }
 
     /// Extract quarter from underlying NaiveDateTime representation.
     /// Quarters range from 1 to 4.
     pub fn quarter(self) -> Expr {
-        let function = move |s: Series| s.quarter().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.quarter")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Quarter))
     }
 
     /// Extract the week from the underlying Date representation.
@@ -112,74 +114,54 @@ impl DateLikeNameSpace {
     /// Returns the ISO week number starting from 1.
     /// The return value ranges from 1 to 53. (The last week of year differs by years.)
     pub fn week(self) -> Expr {
-        let function = move |s: Series| s.week().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.week")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Week))
     }
 
     /// Extract the week day from the underlying Date representation.
     /// Can be performed on Date and Datetime.
 
     /// Returns the weekday number where monday = 0 and sunday = 6
     pub fn weekday(self) -> Expr {
-        let function = move |s: Series| s.weekday().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.weekday")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::WeekDay))
     }
 
     /// Get the month of a Date/Datetime
     pub fn day(self) -> Expr {
-        let function = move |s: Series| s.day().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("day")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Day))
     }
     /// Get the ordinal_day of a Date/Datetime
     pub fn ordinal_day(self) -> Expr {
-        let function = move |s: Series| s.ordinal_day().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.ordinal_day")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::OrdinalDay))
     }
     /// Get the hour of a Datetime/Time64
     pub fn hour(self) -> Expr {
-        let function = move |s: Series| s.hour().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.hour")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Hour))
     }
     /// Get the minute of a Datetime/Time64
     pub fn minute(self) -> Expr {
-        let function = move |s: Series| s.minute().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.minute")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Minute))
     }
 
     /// Get the second of a Datetime/Time64
     pub fn second(self) -> Expr {
-        let function = move |s: Series| s.second().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.second")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::Second))
     }
     /// Get the nanosecond of a Time64
     pub fn nanosecond(self) -> Expr {
-        let function = move |s: Series| s.nanosecond().map(|ca| ca.into_series());
         self.0
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("dt.nanosecond")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::NanoSecond))
     }
 
     pub fn timestamp(self, tu: TimeUnit) -> Expr {
         self.0
-            .map(
-                move |s| s.timestamp(tu).map(|ca| ca.into_series()),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .with_fmt("dt.timestamp")
+            .map_private(FunctionExpr::TemporalExpr(TemporalFunction::TimeStamp(tu)))
     }
 
     /// Offset this `Date/Datetime` by a given offset [`Duration`].

polars/polars-lazy/src/dsl/function_expr/datetime.rs

@@ -0,0 +1,85 @@
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+
+use super::*;
+
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[derive(Clone, PartialEq, Debug, Eq, Hash)]
+pub enum TemporalFunction {
+    Year,
+    IsoYear,
+    Quarter,
+    Month,
+    Week,
+    WeekDay,
+    Day,
+    OrdinalDay,
+    Hour,
+    Minute,
+    Second,
+    NanoSecond,
+    TimeStamp(TimeUnit),
+}
+
+impl Display for TemporalFunction {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        use TemporalFunction::*;
+        let s = match self {
+            Year => "year",
+            IsoYear => "iso_year",
+            Quarter => "quarter",
+            Month => "month",
+            Week => "week",
+            WeekDay => "weekday",
+            Day => "day",
+            OrdinalDay => "ordinal_day",
+            Hour => "hour",
+            Minute => "minute",
+            Second => "second",
+            NanoSecond => "nanosecond",
+            TimeStamp(tu) => return write!(f, "dt.timestamp({})", tu),
+        };
+        write!(f, "dt.{}", s)
+    }
+}
+
+pub(super) fn year(s: &Series) -> PolarsResult<Series> {
+    s.year().map(|ca| ca.into_series())
+}
+pub(super) fn iso_year(s: &Series) -> PolarsResult<Series> {
+    s.iso_year().map(|ca| ca.into_series())
+}
+pub(super) fn month(s: &Series) -> PolarsResult<Series> {
+    s.month().map(|ca| ca.into_series())
+}
+pub(super) fn quarter(s: &Series) -> PolarsResult<Series> {
+    s.quarter().map(|ca| ca.into_series())
+}
+pub(super) fn week(s: &Series) -> PolarsResult<Series> {
+    s.week().map(|ca| ca.into_series())
+}
+pub(super) fn weekday(s: &Series) -> PolarsResult<Series> {
+    s.weekday().map(|ca| ca.into_series())
+}
+pub(super) fn day(s: &Series) -> PolarsResult<Series> {
+    s.day().map(|ca| ca.into_series())
+}
+pub(super) fn ordinal_day(s: &Series) -> PolarsResult<Series> {
+    s.ordinal_day().map(|ca| ca.into_series())
+}
+pub(super) fn hour(s: &Series) -> PolarsResult<Series> {
+    s.hour().map(|ca| ca.into_series())
+}
+pub(super) fn minute(s: &Series) -> PolarsResult<Series> {
+    s.minute().map(|ca| ca.into_series())
+}
+pub(super) fn second(s: &Series) -> PolarsResult<Series> {
+    s.second().map(|ca| ca.into_series())
+}
+pub(super) fn nanosecond(s: &Series) -> PolarsResult<Series> {
+    s.nanosecond().map(|ca| ca.into_series())
+}
+
+pub(super) fn timestamp(s: &Series, tu: TimeUnit) -> PolarsResult<Series> {
+    s.timestamp(tu).map(|ca| ca.into_series())
+}

polars/polars-lazy/src/dsl/function_expr/mod.rs

@@ -2,6 +2,8 @@
 mod arg_where;
 #[cfg(feature = "round_series")]
 mod clip;
+#[cfg(feature = "temporal")]
+mod datetime;
 mod dispatch;
 mod fill_null;
 #[cfg(feature = "is_in")]
@@ -37,6 +39,8 @@ use polars_core::prelude::*;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
+#[cfg(feature = "temporal")]
+pub(super) use self::datetime::TemporalFunction;
 pub(super) use self::nan::NanFunction;
 #[cfg(feature = "strings")]
 pub(super) use self::strings::StringFunction;
@@ -61,6 +65,8 @@ pub enum FunctionExpr {
     SearchSorted,
     #[cfg(feature = "strings")]
     StringExpr(StringFunction),
+    #[cfg(feature = "temporal")]
+    TemporalExpr(TemporalFunction),
     #[cfg(feature = "date_offset")]
     DateOffset(Duration),
     #[cfg(feature = "trigonometry")]
@@ -122,6 +128,8 @@ impl Display for FunctionExpr {
             SearchSorted => write!(f, "search_sorted"),
             #[cfg(feature = "strings")]
             StringExpr(s) => write!(f, "{}", s),
+            #[cfg(feature = "temporal")]
+            TemporalExpr(fun) => write!(f, "{}", fun),
             #[cfg(feature = "date_offset")]
             DateOffset(_) => write!(f, "dt.offset_by"),
             #[cfg(feature = "trigonometry")]
@@ -255,6 +263,8 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             }
             #[cfg(feature = "strings")]
             StringExpr(s) => s.into(),
+            #[cfg(feature = "temporal")]
+            TemporalExpr(func) => func.into(),
 
             #[cfg(feature = "date_offset")]
             DateOffset(offset) => {
@@ -366,3 +376,25 @@ impl From<StringFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
         }
     }
 }
+
+#[cfg(feature = "temporal")]
+impl From<TemporalFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
+    fn from(func: TemporalFunction) -> Self {
+        use TemporalFunction::*;
+        match func {
+            Year => map!(datetime::year),
+            IsoYear => map!(datetime::iso_year),
+            Month => map!(datetime::month),
+            Quarter => map!(datetime::quarter),
+            Week => map!(datetime::week),
+            WeekDay => map!(datetime::weekday),
+            Day => map!(datetime::day),
+            OrdinalDay => map!(datetime::ordinal_day),
+            Hour => map!(datetime::hour),
+            Minute => map!(datetime::minute),
+            Second => map!(datetime::second),
+            NanoSecond => map!(datetime::nanosecond),
+            TimeStamp(tu) => map!(datetime::timestamp, tu),
+        }
+    }
+}

polars/polars-lazy/src/dsl/function_expr/schema.rs

@@ -106,6 +106,17 @@ impl FunctionExpr {
                     Uppercase | Lowercase => with_dtype(DataType::Utf8),
                 }
             }
+            #[cfg(feature = "temporal")]
+            TemporalExpr(fun) => {
+                use TemporalFunction::*;
+                let dtype = match fun {
+                    Year | IsoYear => DataType::Int32,
+                    Month | Quarter | Week | WeekDay | Day | OrdinalDay | Hour | Minute
+                    | NanoSecond | Second => DataType::UInt32,
+                    TimeStamp(_) => DataType::Int64,
+                };
+                with_dtype(dtype)
+            }
 
             #[cfg(feature = "date_offset")]
             DateOffset(_) => same_type(),