Implement $dt$is_leap_year() (#1077)

pola-rs · May 5, 2024 · adcc144 · adcc144
1 parent 8f41b12
commit adcc144
Show file tree

Hide file tree

Showing 7 changed files with 197 additions and 188 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -14,6 +14,7 @@
 - New methods `$str$head()` and `$str$tail()` (#1074).
 - New S3 methods `nanoarrow::as_nanoarrow_array_stream()` and `nanoarrow::infer_nanoarrow_schema()`
   for `RPolarsSeries` (#1076).
+- New method `$dt$is_leap_year()` (#1077).
 
 ## Polars R Package 0.16.3
 

diff --git a/R/expr__datetime.R b/R/expr__datetime.R
@@ -35,7 +35,7 @@
 #' df
 ExprDT_truncate = function(every, offset = NULL) {
   .pr$Expr$dt_truncate(self, every, offset) |>
-    unwrap("in dt$truncate()")
+    unwrap("in $dt$truncate()")
 }
 
 #' Round datetime
@@ -83,7 +83,7 @@ ExprDT_truncate = function(every, offset = NULL) {
 #' df
 ExprDT_round = function(every, offset = NULL) {
   .pr$Expr$dt_round(self, every, offset) |>
-    unwrap("in dt$round()")
+    unwrap("in $dt$round()")
 }
 
 # ExprDT_combine = function(self, tm: time | pli.RPolarsExpr, tu: TimeUnit = "us") -> pli.RPolarsExpr:
@@ -547,7 +547,7 @@ ExprDT_epoch = function(tu = c("us", "ns", "ms", "s", "d")) {
     or_else = Err(
       paste("tu must be one of 'ns', 'us', 'ms', 's', 'd', got", str_string(tu))
     )
-  ) |> map_err(\(err) paste("in dt$epoch:", err))
+  ) |> map_err(\(err) paste("in $dt$epoch:", err))
 
   unwrap(expr_result)
 }
@@ -575,7 +575,7 @@ ExprDT_epoch = function(tu = c("us", "ns", "ms", "s", "d")) {
 #' )
 ExprDT_timestamp = function(tu = c("ns", "us", "ms")) {
   .pr$Expr$timestamp(self, tu[1]) |>
-    map_err(\(err) paste("in dt$timestamp:", err)) |>
+    map_err(\(err) paste("in $dt$timestamp:", err)) |>
     unwrap()
 }
 
@@ -604,7 +604,7 @@ ExprDT_timestamp = function(tu = c("ns", "us", "ms")) {
 #' )
 ExprDT_with_time_unit = function(tu = c("ns", "us", "ms")) {
   .pr$Expr$dt_with_time_unit(self, tu[1]) |>
-    map_err(\(err) paste("in dt$with_time_unit:", err)) |>
+    map_err(\(err) paste("in $dt$with_time_unit:", err)) |>
     unwrap()
 }
 
@@ -634,7 +634,7 @@ ExprDT_with_time_unit = function(tu = c("ns", "us", "ms")) {
 #' )
 ExprDT_cast_time_unit = function(tu = c("ns", "us", "ms")) {
   .pr$Expr$dt_cast_time_unit(self, tu[1]) |>
-    map_err(\(err) paste("in dt$cast_time_unit:", err)) |>
+    map_err(\(err) paste("in $dt$cast_time_unit:", err)) |>
     unwrap()
 }
 
@@ -661,7 +661,7 @@ ExprDT_cast_time_unit = function(tu = c("ns", "us", "ms")) {
 ExprDT_convert_time_zone = function(time_zone) {
   check_tz_to_result(time_zone) |>
     and_then(\(valid_tz) .pr$Expr$dt_convert_time_zone(self, valid_tz)) |>
-    map_err(\(err) paste("in dt$convert_time_zone:", err)) |>
+    map_err(\(err) paste("in $dt$convert_time_zone:", err)) |>
     unwrap("in $convert_time_zone():")
 }
 
@@ -947,5 +947,19 @@ ExprDT_offset_by = function(by) {
 #' df$with_columns(times = pl$col("dates")$dt$time())
 ExprDT_time = function() {
   .pr$Expr$dt_time(self) |>
-    unwrap("in dt$time()")
+    unwrap("in $dt$time()")
+}
+
+#' Determine whether the year of the underlying date is a leap year
+#'
+#' @return An Expr of datatype Bool
+#'
+#' @examples
+#' df = pl$DataFrame(date = as.Date(c("2000-01-01", "2001-01-01", "2002-01-01")))
+#'
+#' df$with_columns(
+#'   leap_year = pl$col("date")$dt$is_leap_year()
+#' )
+ExprDT_is_leap_year = function() {
+  .pr$Expr$dt_is_leap_year(self)
 }
diff --git a/R/extendr-wrappers.R b/R/extendr-wrappers.R
@@ -852,6 +852,8 @@ RPolarsExpr$dt_total_nanoseconds <- function() .Call(wrap__RPolarsExpr__dt_total
 
 RPolarsExpr$dt_offset_by <- function(by) .Call(wrap__RPolarsExpr__dt_offset_by, self, by)
 
+RPolarsExpr$dt_is_leap_year <- function() .Call(wrap__RPolarsExpr__dt_is_leap_year, self)
+
 RPolarsExpr$repeat_by <- function(by) .Call(wrap__RPolarsExpr__repeat_by, self, by)
 
 RPolarsExpr$log10 <- function() .Call(wrap__RPolarsExpr__log10, self)

diff --git a/man/ExprDT_is_leap_year.Rd b/man/ExprDT_is_leap_year.Rd
diff --git a/src/rust/src/lazy/dsl.rs b/src/rust/src/lazy/dsl.rs
@@ -16,8 +16,7 @@ use crate::CONFIG;
 use extendr_api::{extendr, prelude::*, rprintln, Deref, DerefMut};
 use pl::PolarsError as pl_error;
 use pl::{
-    Duration, DurationMethods, IntoSeries, RollingGroupOptions, SetOperation, StringNameSpaceImpl,
-    TemporalMethods,
+    Duration, IntoSeries, RollingGroupOptions, SetOperation, StringNameSpaceImpl, TemporalMethods,
 };
 use polars::lazy::dsl;
 use polars::prelude as pl;
@@ -1469,80 +1468,35 @@ impl RPolarsExpr {
     }
 
     pub fn dt_total_days(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.days().into_series())),
-                pl::GetOutput::from_type(pl::DataType::Int64),
-            )
-            .into())
+        Ok(self.0.clone().dt().total_days().into())
     }
     pub fn dt_total_hours(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.hours().into_series())),
-                pl::GetOutput::from_type(pl::DataType::Int64),
-            )
-            .into())
+        Ok(self.0.clone().dt().total_hours().into())
     }
     pub fn dt_total_minutes(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.minutes().into_series())),
-                pl::GetOutput::from_type(pl::DataType::Int64),
-            )
-            .into())
+        Ok(self.0.clone().dt().total_minutes().into())
     }
     pub fn dt_total_seconds(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.seconds().into_series())),
-                pl::GetOutput::from_type(pl::DataType::Int64),
-            )
-            .into())
+        Ok(self.0.clone().dt().total_seconds().into())
     }
     pub fn dt_total_milliseconds(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.milliseconds().into_series())),
-                pl::GetOutput::from_type(pl::DataType::Int64),
-            )
-            .into())
+        Ok(self.0.clone().dt().total_milliseconds().into())
     }
     pub fn dt_total_microseconds(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.microseconds().into_series())),
-                pl::GetOutput::from_type(pl::DataType::Int64),
-            )
-            .into())
+        Ok(self.0.clone().dt().total_microseconds().into())
     }
     pub fn dt_total_nanoseconds(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.nanoseconds().into_series())),
-                pl::GetOutput::from_type(pl::DataType::Int64),
-            )
-            .into())
+        Ok(self.0.clone().dt().total_nanoseconds().into())
     }
 
     pub fn dt_offset_by(&self, by: Robj) -> RResult<Self> {
         Ok(self.clone().0.dt().offset_by(robj_to!(PLExpr, by)?).into())
     }
 
+    pub fn dt_is_leap_year(&self) -> Self {
+        self.clone().0.dt().is_leap_year().into()
+    }
+
     pub fn repeat_by(&self, by: &RPolarsExpr) -> Self {
         self.clone().0.repeat_by(by.0.clone()).into()
     }