Implement $str$extract_groups()

NEWS.md

@@ -93,6 +93,7 @@
 - New functions `pl$date_ranges()` and `pl$datetime_ranges()` (#962).
 - Export the `Duration` datatype (#955).
 - New functions `pl$int_range()` and `pl$int_ranges()` (#968).
+- New string method `$str$extract_groups()` (#979).
 
 ### Bug fixes
 

R/expr__string.R

@@ -914,3 +914,42 @@ ExprStr_replace_many = function(patterns, replace_with, ascii_case_insensitive =
   .pr$Expr$str_replace_many(self, patterns, replace_with, ascii_case_insensitive) |>
     unwrap("in str$replace_many():")
 }
+
+
+#' Extract all capture groups for the given regex pattern
+#'
+#' @inheritParams ExprStr_contains
+#'
+#' @details
+#' All group names are strings. If your pattern contains unnamed groups, their
+#' numerical position is converted to a string. See examples.
+#'
+#' @return Expr of data type [Struct][DataType_Struct] with fields of data type
+#' `String`.
+#'
+#' @examples
+#' df = pl$DataFrame(
+#'   url = c(
+#'     "http://vote.com/ballon_dor?candidate=messi&ref=python",
+#'     "http://vote.com/ballon_dor?candidate=weghorst&ref=polars",
+#'     "http://vote.com/ballon_dor?error=404&ref=rust"
+#'   )
+#' )
+#'
+#' pattern = r"(candidate=(?<candidate>\w+)&ref=(?<ref>\w+))"
+#'
+#' df$with_columns(
+#'   captures = pl$col("url")$str$extract_groups(pattern)
+#' )$unnest("captures")
+#'
+#' # If the groups are unnamed, their numerical position (as a string) is used:
+#'
+#' pattern = r"(candidate=(\w+)&ref=(\w+))"
+#'
+#' df$with_columns(
+#'   captures = pl$col("url")$str$extract_groups(pattern)
+#' )$unnest("captures")
+ExprStr_extract_groups = function(pattern) {
+  .pr$Expr$str_extract_groups(self, pattern) |>
+    unwrap("in str$extract_groups():")
+}

R/extendr-wrappers.R

@@ -992,6 +992,8 @@ RPolarsExpr$str_extract <- function(pattern, group_index) .Call(wrap__RPolarsExp
 
 RPolarsExpr$str_extract_all <- function(pattern) .Call(wrap__RPolarsExpr__str_extract_all, self, pattern)
 
+RPolarsExpr$str_extract_groups <- function(pattern) .Call(wrap__RPolarsExpr__str_extract_groups, self, pattern)
+
 RPolarsExpr$str_count_matches <- function(pattern, literal) .Call(wrap__RPolarsExpr__str_count_matches, self, pattern, literal)
 
 RPolarsExpr$str_to_date <- function(format, strict, exact, cache) .Call(wrap__RPolarsExpr__str_to_date, self, format, strict, exact, cache)

src/rust/Cargo.toml

@@ -90,6 +90,7 @@ features = [
   #"dtype-struct",
   "dtype-time",
   "dynamic_group_by",
+  "extract_groups",
   "extract_jsonpath",
   "ewma",
   "find_many",

src/rust/src/lazy/dsl.rs

@@ -2142,6 +2142,10 @@ impl RPolarsExpr {
         self.0.clone().str().extract_all(pattern.0.clone()).into()
     }
 
+    pub fn str_extract_groups(&self, pattern: &str) -> RResult<Self> {
+        Ok(self.0.clone().str().extract_groups(pattern)?.into())
+    }
+
     pub fn str_count_matches(&self, pattern: Robj, literal: Robj) -> RResult<Self> {
         Ok(self
             .0

tests/testthat/_snaps/after-wrappers.md

@@ -406,29 +406,29 @@
       [267] "str_contains_any"           "str_count_matches"         
       [269] "str_ends_with"              "str_explode"               
       [271] "str_extract"                "str_extract_all"           
-      [273] "str_hex_decode"             "str_hex_encode"            
-      [275] "str_json_decode"            "str_json_path_match"       
-      [277] "str_len_bytes"              "str_len_chars"             
-      [279] "str_pad_end"                "str_pad_start"             
-      [281] "str_parse_int"              "str_replace"               
-      [283] "str_replace_all"            "str_replace_many"          
-      [285] "str_reverse"                "str_slice"                 
-      [287] "str_split"                  "str_split_exact"           
-      [289] "str_splitn"                 "str_starts_with"           
-      [291] "str_strip_chars"            "str_strip_chars_end"       
-      [293] "str_strip_chars_start"      "str_to_date"               
-      [295] "str_to_datetime"            "str_to_lowercase"          
-      [297] "str_to_time"                "str_to_titlecase"          
-      [299] "str_to_uppercase"           "str_zfill"                 
-      [301] "struct_field_by_name"       "struct_rename_fields"      
-      [303] "sub"                        "sum"                       
-      [305] "tail"                       "tan"                       
-      [307] "tanh"                       "timestamp"                 
-      [309] "to_physical"                "top_k"                     
-      [311] "unique"                     "unique_counts"             
-      [313] "unique_stable"              "upper_bound"               
-      [315] "value_counts"               "var"                       
-      [317] "xor"                       
+      [273] "str_extract_groups"         "str_hex_decode"            
+      [275] "str_hex_encode"             "str_json_decode"           
+      [277] "str_json_path_match"        "str_len_bytes"             
+      [279] "str_len_chars"              "str_pad_end"               
+      [281] "str_pad_start"              "str_parse_int"             
+      [283] "str_replace"                "str_replace_all"           
+      [285] "str_replace_many"           "str_reverse"               
+      [287] "str_slice"                  "str_split"                 
+      [289] "str_split_exact"            "str_splitn"                
+      [291] "str_starts_with"            "str_strip_chars"           
+      [293] "str_strip_chars_end"        "str_strip_chars_start"     
+      [295] "str_to_date"                "str_to_datetime"           
+      [297] "str_to_lowercase"           "str_to_time"               
+      [299] "str_to_titlecase"           "str_to_uppercase"          
+      [301] "str_zfill"                  "struct_field_by_name"      
+      [303] "struct_rename_fields"       "sub"                       
+      [305] "sum"                        "tail"                      
+      [307] "tan"                        "tanh"                      
+      [309] "timestamp"                  "to_physical"               
+      [311] "top_k"                      "unique"                    
+      [313] "unique_counts"              "unique_stable"             
+      [315] "upper_bound"                "value_counts"              
+      [317] "var"                        "xor"                       
 
 # public and private methods of each class When
 

tests/testthat/test-expr_string.R

@@ -795,3 +795,40 @@ patrick::with_parameters_test_that(
   },
   .cases = make_datetime_format_cases()
 )
+
+test_that("str$extract_groups() works", {
+  df = pl$DataFrame(
+    url = c(
+      "http://vote.com/ballon_dor?candidate=messi&ref=python",
+      "http://vote.com/ballon_dor?candidate=weghorst&ref=polars",
+      "http://vote.com/ballon_dor?error=404&ref=rust"
+    )
+  )
+
+  # named patterns
+  pattern = r"(candidate=(?<candidate>\w+)&ref=(?<ref>\w+))"
+  expect_identical(
+    df$select(
+      captures = pl$col("url")$str$extract_groups(pattern)
+    )$unnest("captures")$to_list(),
+    list(candidate = c("messi", "weghorst", NA), ref = c("python", "polars", NA))
+  )
+
+  # unnamed patterns
+  pattern = r"(candidate=(\w+)&ref=(\w+))"
+  expect_identical(
+    df$select(
+      captures = pl$col("url")$str$extract_groups(pattern)
+    )$unnest("captures")$to_list(),
+    list("1" = c("messi", "weghorst", NA), "2" = c("python", "polars", NA))
+  )
+
+  # empty pattern
+  pattern = ""
+  expect_identical(
+    df$select(
+      captures = pl$col("url")$str$extract_groups(pattern)
+    )$unnest("captures")$to_list(),
+    list(url = NULL)
+  )
+})