feat: <Series>$struct$fields (#1002)

pola-rs · Apr 1, 2024 · 5565192 · 5565192
1 parent 0ac6f5d
commit 5565192
Show file tree

Hide file tree

Showing 7 changed files with 83 additions and 13 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -110,6 +110,7 @@
 - New argument `n` in `$str$replace()` (#987).
 - Method `$over()` gains an argument `mapping_strategy` (#984, #988).
 - New method `$item()` for `DataFrame` and `Series` (#992).
+- New active binding `<Series>$struct$fields` (#1002).
 
 ### Bug fixes
 

diff --git a/R/extendr-wrappers.R b/R/extendr-wrappers.R
@@ -1270,10 +1270,10 @@ RPolarsSeries$shape <- function() .Call(wrap__RPolarsSeries__shape, self)
 
 RPolarsSeries$len <- function() .Call(wrap__RPolarsSeries__len, self)
 
-RPolarsSeries$n_chunks <- function() .Call(wrap__RPolarsSeries__n_chunks, self)
-
 RPolarsSeries$chunk_lengths <- function() .Call(wrap__RPolarsSeries__chunk_lengths, self)
 
+RPolarsSeries$n_chunks <- function() .Call(wrap__RPolarsSeries__n_chunks, self)
+
 RPolarsSeries$alias <- function(name) .Call(wrap__RPolarsSeries__alias, self, name)
 
 RPolarsSeries$all <- function() .Call(wrap__RPolarsSeries__all, self)
@@ -1314,6 +1314,8 @@ RPolarsSeries$to_frame <- function() .Call(wrap__RPolarsSeries__to_frame, self)
 
 RPolarsSeries$set_sorted_mut <- function(descending) invisible(.Call(wrap__RPolarsSeries__set_sorted_mut, self, descending))
 
+RPolarsSeries$struct_fields <- function() .Call(wrap__RPolarsSeries__struct_fields, self)
+
 RPolarsSeries$from_arrow_array_stream_str <- function(name, robj_str) .Call(wrap__RPolarsSeries__from_arrow_array_stream_str, name, robj_str)
 
 RPolarsSeries$from_arrow_array_robj <- function(name, array) .Call(wrap__RPolarsSeries__from_arrow_array_robj, name, array)

diff --git a/R/series__series.R b/R/series__series.R
@@ -86,7 +86,11 @@
 #'
 #' ## struct
 #'
-#' `$struct` stores all struct related methods.
+#' `$struct` stores all struct related methods and active bindings.
+#'
+#' Active bindings specific to Series:
+#'
+#' - `$struct$fields`: Returns a character vector of the fields in the struct.
 #'
 #' @inheritSection DataFrame_class Conversion to R data types considerations
 #' @keywords Series
@@ -115,6 +119,9 @@
 #' s
 #' s$dt$day()
 #'
+#' # Other active bindings in subnamespaces
+#' as_polars_series(data.frame(a = 1:2, b = 3:4))$struct$fields
+#'
 #' # show all available methods for Series
 #' pl$show_all_public_methods("RPolarsSeries")
 #'
@@ -195,26 +202,51 @@ add_expr_methods_to_series = function() {
 ## Sub-namespaces
 
 #' Make sub namespace of Series from Expr sub namespace
+#' @param ... Addtional funtions to add to the namespace
 #' @noRd
-series_make_sub_ns = function(pl_series, .expr_make_sub_ns_fn) {
-  df = pl_series$to_frame()
+series_make_sub_ns = function(pl_series, .expr_make_sub_ns_fn, ...) {
   # Override `self` in `$.RPolarsExpr`
   self = pl$col(pl_series$name) # nolint: object_usage_linter
 
   fns = .expr_make_sub_ns_fn(pl$col(pl_series$name))
-  lapply(fns, \(f) {
+  new_fns = lapply(fns, \(f) {
     environment(f) = parent.frame(2L)
     new_f = function() {
       expr = do.call(f, as.list(match.call()[-1]), envir = parent.frame())
       pcase(
-        inherits(expr, "RPolarsExpr"), df$select(expr)$to_series(0),
+        inherits(expr, "RPolarsExpr"), pl_series$to_frame()$select(expr)$to_series(0),
         or_else = expr
       )
     }
 
     formals(new_f) = formals(f)
     new_f
   })
+
+  if (!missing(...)) {
+    additional_fns = list(...) |>
+      lapply(\(f) {
+        environment(f) = parent.frame(2L)
+        new_f = function() {
+          do.call(f, as.list(match.call()[-1]), envir = parent.frame())
+        }
+
+        formals(new_f) = formals(f)
+        class(new_f) = class(f)
+        new_f
+      })
+
+    new_fns = c(additional_fns, new_fns)
+  }
+
+  new_fns |>
+    lapply(\(f) {
+      if (inherits(f, "property")) {
+        f()
+      } else {
+        f
+      }
+    })
 }
 
 Series_arr = method_as_active_binding(\() series_make_sub_ns(self, expr_arr_make_sub_ns))
@@ -229,7 +261,17 @@ Series_list = method_as_active_binding(\() series_make_sub_ns(self, expr_list_ma
 
 Series_str = method_as_active_binding(\() series_make_sub_ns(self, expr_str_make_sub_ns))
 
-Series_struct = method_as_active_binding(\() series_make_sub_ns(self, expr_struct_make_sub_ns))
+Series_struct = method_as_active_binding(
+  \() {
+    pl_series = NULL # Workaround for R CMD check `Undefined global functions or variables` error
+    series_make_sub_ns(
+      self, expr_struct_make_sub_ns,
+      fields = method_as_active_binding(function() {
+        unwrap(.pr$Series$struct_fields(pl_series), "in $struct$fields:")
+      })
+    )
+  }
+)
 
 
 # TODO: change the arguments in 0.17.0

diff --git a/man/Series_class.Rd b/man/Series_class.Rd
diff --git a/src/rust/src/series.rs b/src/rust/src/series.rs
@@ -500,6 +500,11 @@ impl RPolarsSeries {
         };
     }
 
+    fn struct_fields(&self) -> RResult<Vec<&str>> {
+        let ca = self.0.struct_().map_err(polars_to_rpolars_err)?;
+        Ok(ca.fields().iter().map(|s| s.name()).collect())
+    }
+
     pub fn from_arrow_array_stream_str(name: Robj, robj_str: Robj) -> RResult<Robj> {
         let name = robj_to!(str, name)?;
         let s = crate::arrow_interop::to_rust::arrow_stream_to_series_internal(robj_str)?

diff --git a/tests/testthat/_snaps/after-wrappers.md b/tests/testthat/_snaps/after-wrappers.md
@@ -713,10 +713,11 @@
       [35] "rename_mut"                  "rep"                        
       [37] "set_sorted_mut"              "shape"                      
       [39] "sleep"                       "sort_mut"                   
-      [41] "std"                         "sub"                        
-      [43] "sum"                         "to_fmt_char"                
-      [45] "to_frame"                    "to_r"                       
-      [47] "value_counts"                "var"                        
+      [41] "std"                         "struct_fields"              
+      [43] "sub"                         "sum"                        
+      [45] "to_fmt_char"                 "to_frame"                   
+      [47] "to_r"                        "value_counts"               
+      [49] "var"                        
 
 # public and private methods of each class RThreadHandle
 

diff --git a/tests/testthat/test-series-sub-namespace.R b/tests/testthat/test-series-sub-namespace.R
@@ -99,3 +99,14 @@ test_that("Method dispatch Expr -> Series works in functions", {
       is.na()
   )
 })
+
+test_that("$struct$fields", {
+  expect_identical(
+    as_polars_series(data.frame(a = 1, b = 2))$struct$fields,
+    c("a", "b")
+  )
+  expect_grepl_error(
+    as_polars_series(1:3)$struct$fields,
+    "data types don't match"
+  )
+})