diff --git a/NEWS.md b/NEWS.md index 7e41592f4..883337632 100644 --- a/NEWS.md +++ b/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 diff --git a/R/expr__struct.R b/R/expr__struct.R index 499ec3156..5ae89b126 100644 --- a/R/expr__struct.R +++ b/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:") } diff --git a/R/extendr-wrappers.R b/R/extendr-wrappers.R index c610b3f6b..672e462ff 100644 --- a/R/extendr-wrappers.R +++ b/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) diff --git a/man/ExprStruct_with_fields.Rd b/man/ExprStruct_with_fields.Rd new file mode 100644 index 000000000..358bc20fb --- /dev/null +++ b/man/ExprStruct_with_fields.Rd @@ -0,0 +1,37 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/expr__struct.R +\name{ExprStruct_with_fields} +\alias{ExprStruct_with_fields} +\title{Add or overwrite fields of this struct} +\usage{ +ExprStruct_with_fields(...) +} +\arguments{ +\item{...}{Field(s) to add. Accepts expression input. Strings are parsed as +column names, other non-expression inputs are parsed as literals.} +} +\value{ +An \code{\link[=RPolarsExpr]{Expr}} of data type Struct. +} +\description{ +This is similar to \code{\link[=DataFrame_with_columns]{with_columns}} on +\code{\link[=RPolarsDataFrame]{DataFrame}}. +} +\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") +} diff --git a/src/rust/src/lazy/dsl.rs b/src/rust/src/lazy/dsl.rs index 7663087c1..de7f7e8e1 100644 --- a/src/rust/src/lazy/dsl.rs +++ b/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 { + 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 { diff --git a/tests/testthat/_snaps/after-wrappers.md b/tests/testthat/_snaps/after-wrappers.md index 91f50a27f..d79fb6dea 100644 --- a/tests/testthat/_snaps/after-wrappers.md +++ b/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 diff --git a/tests/testthat/test-expr_struct.R b/tests/testthat/test-expr_struct.R index 8f9b686c3..ba6e43b11 100644 --- a/tests/testthat/test-expr_struct.R +++ b/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)) + ) +})