feat: Implement <expr>$dt$days_in_month() (#1659)

Author: etiennebacher

NEWS.md

@@ -7,6 +7,7 @@
 - `<expr>$item()` to strictly extract a single value from an expression (#1652).
 - `<expr>$arr$eval()` to run any Polars expression on all subarrays of an Array column (#1653).
 - `<expr>$name$replace()` to replace expression names using regular expressions (#1654).
+- `<expr>$dt$days_in_month()` (#1659).
 
 ## polars 1.6.0
 

R/000-wrappers.R

@@ -1649,6 +1649,12 @@ class(`PlRDataTypeExpr`) <- c("PlRDataTypeExpr__bundle", "savvy_polars__sealed")
   }
 }
 
+`PlRExpr_dt_days_in_month` <- function(self) {
+  function() {
+    .savvy_wrap_PlRExpr(.Call(savvy_PlRExpr_dt_days_in_month__impl, `self`))
+  }
+}
+
 `PlRExpr_dt_dst_offset` <- function(self) {
   function() {
     .savvy_wrap_PlRExpr(.Call(savvy_PlRExpr_dt_dst_offset__impl, `self`))
@@ -3529,6 +3535,7 @@ class(`PlRDataTypeExpr`) <- c("PlRDataTypeExpr__bundle", "savvy_polars__sealed")
   e$`dt_convert_time_zone` <- `PlRExpr_dt_convert_time_zone`(ptr)
   e$`dt_date` <- `PlRExpr_dt_date`(ptr)
   e$`dt_day` <- `PlRExpr_dt_day`(ptr)
+  e$`dt_days_in_month` <- `PlRExpr_dt_days_in_month`(ptr)
   e$`dt_dst_offset` <- `PlRExpr_dt_dst_offset`(ptr)
   e$`dt_epoch_seconds` <- `PlRExpr_dt_epoch_seconds`(ptr)
   e$`dt_hour` <- `PlRExpr_dt_hour`(ptr)

R/expr-datetime.R

@@ -1240,3 +1240,17 @@ expr_dt_replace <- function(
     )
   })
 }
+
+#' Extract the number of days in the month from the underlying Date representation.
+#'
+#' @inherit as_polars_expr return
+#' @examples
+#' df <- pl$DataFrame(date = as.Date(c("2020-01-01", "2020-02-01", "2020-03-01")))
+#'
+#' df$with_columns(
+#'   days_in_month = pl$col("date")$dt$days_in_month()
+#' )
+expr_dt_days_in_month <- function() {
+  self$`_rexpr`$dt_days_in_month() |>
+    wrap()
+}

altdoc/mkdocs.yml

@@ -633,6 +633,7 @@ nav:
             - dt_convert_time_zone: man/expr_dt_convert_time_zone.md
             - dt_date: man/expr_dt_date.md
             - dt_day: man/expr_dt_day.md
+            - dt_days_in_month: man/expr_dt_days_in_month.md
             - dt_dst_offset: man/expr_dt_dst_offset.md
             - dt_epoch: man/expr_dt_epoch.md
             - dt_hour: man/expr_dt_hour.md

man/expr_dt_days_in_month.Rd

@@ -1199,6 +1199,11 @@ SEXP savvy_PlRExpr_dt_day__impl(SEXP self__) {
     return handle_result(res);
 }
 
+SEXP savvy_PlRExpr_dt_days_in_month__impl(SEXP self__) {
+    SEXP res = savvy_PlRExpr_dt_days_in_month__ffi(self__);
+    return handle_result(res);
+}
+
 SEXP savvy_PlRExpr_dt_dst_offset__impl(SEXP self__) {
     SEXP res = savvy_PlRExpr_dt_dst_offset__ffi(self__);
     return handle_result(res);
@@ -3634,6 +3639,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"savvy_PlRExpr_dt_convert_time_zone__impl", (DL_FUNC) &savvy_PlRExpr_dt_convert_time_zone__impl, 2},
     {"savvy_PlRExpr_dt_date__impl", (DL_FUNC) &savvy_PlRExpr_dt_date__impl, 1},
     {"savvy_PlRExpr_dt_day__impl", (DL_FUNC) &savvy_PlRExpr_dt_day__impl, 1},
+    {"savvy_PlRExpr_dt_days_in_month__impl", (DL_FUNC) &savvy_PlRExpr_dt_days_in_month__impl, 1},
     {"savvy_PlRExpr_dt_dst_offset__impl", (DL_FUNC) &savvy_PlRExpr_dt_dst_offset__impl, 1},
     {"savvy_PlRExpr_dt_epoch_seconds__impl", (DL_FUNC) &savvy_PlRExpr_dt_epoch_seconds__impl, 1},
     {"savvy_PlRExpr_dt_hour__impl", (DL_FUNC) &savvy_PlRExpr_dt_hour__impl, 1},

src/init.c

@@ -245,6 +245,7 @@ SEXP savvy_PlRExpr_dt_combine__ffi(SEXP self__, SEXP c_arg__time, SEXP c_arg__ti
 SEXP savvy_PlRExpr_dt_convert_time_zone__ffi(SEXP self__, SEXP c_arg__time_zone);
 SEXP savvy_PlRExpr_dt_date__ffi(SEXP self__);
 SEXP savvy_PlRExpr_dt_day__ffi(SEXP self__);
+SEXP savvy_PlRExpr_dt_days_in_month__ffi(SEXP self__);
 SEXP savvy_PlRExpr_dt_dst_offset__ffi(SEXP self__);
 SEXP savvy_PlRExpr_dt_epoch_seconds__ffi(SEXP self__);
 SEXP savvy_PlRExpr_dt_hour__ffi(SEXP self__);

src/rust/api.h

@@ -253,4 +253,8 @@ impl PlRExpr {
             )
             .into())
     }
+
+    fn dt_days_in_month(&self) -> Result<Self> {
+        Ok(self.inner.clone().dt().days_in_month().into())
+    }
 }

src/rust/src/expr/datetime.rs

@@ -1012,3 +1012,11 @@ patrick::with_parameters_test_that(
     }
   }
 )
+
+test_that("dt$days_in_month", {
+  df <- pl$DataFrame(date = as.Date(c("2020-01-01", "2020-02-01", "2020-03-01")))
+  expect_equal(
+    df$select(pl$col("date")$dt$days_in_month()),
+    pl$DataFrame(date = c(31, 29, 31))$cast(pl$Int8)
+  )
+})