Implement $struct$with_fields()

NEWS.md

@@ -29,6 +29,7 @@
   like `pl$Struct(a = pl$Boolean)` (#1053).
 - In `$all()` and `$any()`, the argument `drop_nulls` is renamed `ignore_nulls`,
   and this argument must be named (#1050).
+- New method `$struct$with_fields()` (#1109).
 
 ## Polars R Package 0.16.4
 

R/expr__struct.R

@@ -21,7 +21,8 @@
 #'   pl$col("struct_col")$struct$field("ddd")
 #' )
 ExprStruct_field = function(name) {
-  .pr$Expr$struct_field_by_name(self, result(name)) |> unwrap("in struct$field:")
+  .pr$Expr$struct_field_by_name(self, name) |>
+    unwrap("in $struct$field():")
 }
 
 
@@ -47,5 +48,37 @@ ExprStruct_field = function(name) {
 #' )
 #' df$unnest()
 ExprStruct_rename_fields = function(names) {
-  .pr$Expr$struct_rename_fields(self, result(names)) |> unwrap("in struct$rename_fields:")
+  .pr$Expr$struct_rename_fields(self, names) |> unwrap("in $struct$rename_fields:")
+}
+
+#' Add or overwrite fields of this struct
+#'
+#' This is similar to [`with_columns`][DataFrame_with_columns] on
+#' [`DataFrame`][RPolarsDataFrame].
+#'
+#' @param ... Field(s) to add. Accepts expression input. Strings are parsed as
+#' column names, other non-expression inputs are parsed as literals.
+#'
+#' @return An [`Expr`][RPolarsExpr] of data type Struct.
+#'
+#' @examples
+#' df = pl$DataFrame(x = c(1, 4, 9), y = c(4, 9, 16), multiply = c(10, 2, 3))$
+#'   with_columns(coords = pl$struct(c("x", "y")))$
+#'   select("coords", "multiply")
+#'
+#' df
+#'
+#' df = df$with_columns(
+#'   pl$col("coords")$struct$with_fields(
+#'     pl$col("coords")$struct$field("x")$sqrt(),
+#'     y_mul = pl$col("coords")$struct$field("y") * pl$col("multiply")
+#'   )
+#' )
+#'
+#' df
+#'
+#' df$unnest("coords")
+ExprStruct_with_fields = function(...) {
+  .pr$Expr$struct_with_fields(self, unpack_list(..., .context = "in $struct$with_fields()")) |>
+    unwrap("in $struct$with_fields:")
 }

R/extendr-wrappers.R

@@ -1104,6 +1104,8 @@ RPolarsExpr$struct_field_by_name <- function(name) .Call(wrap__RPolarsExpr__stru
 
 RPolarsExpr$struct_rename_fields <- function(names) .Call(wrap__RPolarsExpr__struct_rename_fields, self, names)
 
+RPolarsExpr$struct_with_fields <- function(fields) .Call(wrap__RPolarsExpr__struct_with_fields, self, fields)
+
 RPolarsExpr$meta_pop <- function() .Call(wrap__RPolarsExpr__meta_pop, self)
 
 RPolarsExpr$meta_eq <- function(other) .Call(wrap__RPolarsExpr__meta_eq, self, other)

src/rust/src/lazy/dsl.rs

@@ -2515,6 +2515,11 @@ impl RPolarsExpr {
         Ok(self.0.clone().struct_().rename_fields(string_vec).into())
     }
 
+    fn struct_with_fields(&self, fields: Robj) -> RResult<Self> {
+        let fields = robj_to!(VecPLExprColNamed, fields)?;
+        Ok(self.0.clone().struct_().with_fields(fields).into())
+    }
+
     //placed in py-polars/src/lazy/meta.rs, however extendr do not support
     //multiple export impl.
     fn meta_pop(&self) -> List {

tests/testthat/_snaps/after-wrappers.md

@@ -436,14 +436,14 @@
       [311] "str_to_lowercase"           "str_to_time"               
       [313] "str_to_titlecase"           "str_to_uppercase"          
       [315] "str_zfill"                  "struct_field_by_name"      
-      [317] "struct_rename_fields"       "sub"                       
-      [319] "sum"                        "tail"                      
-      [321] "tan"                        "tanh"                      
-      [323] "to_physical"                "top_k"                     
-      [325] "unique"                     "unique_counts"             
-      [327] "unique_stable"              "upper_bound"               
-      [329] "value_counts"               "var"                       
-      [331] "xor"                       
+      [317] "struct_rename_fields"       "struct_with_fields"        
+      [319] "sub"                        "sum"                       
+      [321] "tail"                       "tan"                       
+      [323] "tanh"                       "to_physical"               
+      [325] "top_k"                      "unique"                    
+      [327] "unique_counts"              "unique_stable"             
+      [329] "upper_bound"                "value_counts"              
+      [331] "var"                        "xor"                       
 
 # public and private methods of each class When
 

tests/testthat/test-expr_struct.R

@@ -57,6 +57,24 @@ test_that("expr struct$rename_fields", {
   err_state = result(pl$col("")$struct$rename_fields(42))
   expect_grepl_error(unwrap(err_state), "str")
   expect_grepl_error(unwrap(err_state), "\\[names\\]")
-  expect_grepl_error(unwrap(err_state), "in struct\\$rename_fields:")
+  expect_grepl_error(unwrap(err_state), "in \\$struct\\$rename_fields:")
   expect_grepl_error(unwrap(err_state), "42.0")
 })
+
+test_that("$struct$with_fields", {
+  df = pl$DataFrame(x = c(1, 4, 9), y = c(4, 9, 16), multiply = c(10, 2, 3))$
+    with_columns(coords = pl$struct(c("x", "y")))$
+    select("coords", "multiply")
+
+  out = df$select(
+    pl$col("coords")$struct$with_fields(
+      pl$col("coords")$struct$field("x")$sqrt(),
+      y_mul = pl$col("coords")$struct$field("y") * pl$col("multiply")
+    )
+  )$unnest("coords")
+
+  expect_identical(
+    out$to_data_frame(),
+    data.frame(x = c(1, 2, 3), y = c(4, 9, 16), y_mul = c(40, 18, 48))
+  )
+})