init

NEWS.md

@@ -38,7 +38,8 @@
   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).
+- New method `$struct$with_fields()` (#1109) and new function `pl$field()` to
+  be used in expressions in `$struct$with_fields()` (#1113).
 - New methods for `RPolarsDataType`: `$is_enum()`, `$is_categorical()`,
   `$is_known()`, `$is_string()`, `$contains_views()`, `$contains_categorical()`
   (#1112).

R/expr__struct.R

@@ -53,8 +53,9 @@ ExprStruct_rename_fields = function(names) {
 
 #' Add or overwrite fields of this struct
 #'
-#' This is similar to [`with_columns`][DataFrame_with_columns] on
-#' [`DataFrame`][RPolarsDataFrame].
+#' This is similar to [`$with_columns()`][DataFrame_with_columns] on
+#' [`DataFrame`][RPolarsDataFrame]. Use [`pl$field()`][pl_field] to quickly
+#' select a field in a `$struct$with_fields()` context.
 #'
 #' @param ... Field(s) to add. Accepts expression input. Strings are parsed as
 #' column names, other non-expression inputs are parsed as literals.
@@ -70,8 +71,8 @@ ExprStruct_rename_fields = function(names) {
 #'
 #' 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")
+#'     pl$field("x")$sqrt(),
+#'     y_mul = pl$field("y") * pl$col("multiply")
 #'   )
 #' )
 #'

R/extendr-wrappers.R

@@ -38,6 +38,8 @@ sum_horizontal <- function(dotdotdot) .Call(wrap__sum_horizontal, dotdotdot)
 
 mean_horizontal <- function(dotdotdot) .Call(wrap__mean_horizontal, dotdotdot)
 
+field <- function(names) .Call(wrap__field, names)
+
 concat_list <- function(exprs) .Call(wrap__concat_list, exprs)
 
 concat_str <- function(dotdotdot, separator, ignore_nulls) .Call(wrap__concat_str, dotdotdot, separator, ignore_nulls)

R/functions__lazy.R

@@ -1345,3 +1345,20 @@ pl_int_ranges = function(start = 0, end = NULL, step = 1, ..., dtype = pl$Int64)
   int_ranges(start, end, step, dtype) |>
     unwrap("in pl$int_ranges():")
 }
+
+
+#' Quickly select a field in a Struct
+#'
+#' This is syntactic sugar that should mostly be used in
+#' [`$struct$with_fields()`][ExprStruct_with_fields]. `pl$field("x")` is
+#' equivalent to `pl$col("my_struct")$struct$field("x")`.
+#'
+#' @param name Name of the field to select.
+#'
+#' @return An Expr with the datatype from the selected field.
+#'
+#' @inherit ExprStruct_with_fields examples
+pl_field = function(name) {
+  field(name) |>
+    unwrap("in pl$field():")
+}

src/rust/src/rlib.rs

@@ -404,6 +404,18 @@ pub fn int_ranges(start: Robj, end: Robj, step: Robj, dtype: Robj) -> RResult<RP
     Ok(result.into())
 }
 
+#[extendr]
+pub fn field(names: Robj) -> RResult<RPolarsExpr> {
+    let names = robj_to!(Vec, String, names)?;
+    Ok(pl::Expr::Field(
+        names
+            .into_iter()
+            .map(|name| pl::Arc::from(name.as_str()))
+            .collect(),
+    )
+    .into())
+}
+
 extendr_module! {
     mod rlib;
 
@@ -421,6 +433,8 @@ extendr_module! {
     fn sum_horizontal;
     fn mean_horizontal;
 
+    fn field;
+
     fn concat_list;
     fn concat_str;
 

tests/testthat/_snaps/after-wrappers.md

@@ -32,31 +32,32 @@
        [53] "deserialize_lf"            "disable_string_cache"     
        [55] "dtypes"                    "duration"                 
        [57] "element"                   "enable_string_cache"      
-       [59] "first"                     "fold"                     
-       [61] "from_epoch"                "get_global_rpool_cap"     
-       [63] "head"                      "implode"                  
-       [65] "int_range"                 "int_ranges"               
-       [67] "is_schema"                 "last"                     
-       [69] "len"                       "lit"                      
-       [71] "max"                       "max_horizontal"           
-       [73] "mean"                      "mean_horizontal"          
-       [75] "median"                    "mem_address"              
-       [77] "min"                       "min_horizontal"           
-       [79] "n_unique"                  "numeric_dtypes"           
-       [81] "raw_list"                  "read_csv"                 
-       [83] "read_ipc"                  "read_ndjson"              
-       [85] "read_parquet"              "reduce"                   
-       [87] "rolling_corr"              "rolling_cov"              
-       [89] "same_outer_dt"             "scan_csv"                 
-       [91] "scan_ipc"                  "scan_ndjson"              
-       [93] "scan_parquet"              "select"                   
-       [95] "set_global_rpool_cap"      "show_all_public_functions"
-       [97] "show_all_public_methods"   "std"                      
-       [99] "struct"                    "sum"                      
-      [101] "sum_horizontal"            "tail"                     
-      [103] "thread_pool_size"          "time"                     
-      [105] "using_string_cache"        "var"                      
-      [107] "when"                      "with_string_cache"        
+       [59] "field"                     "first"                    
+       [61] "fold"                      "from_epoch"               
+       [63] "get_global_rpool_cap"      "head"                     
+       [65] "implode"                   "int_range"                
+       [67] "int_ranges"                "is_schema"                
+       [69] "last"                      "len"                      
+       [71] "lit"                       "max"                      
+       [73] "max_horizontal"            "mean"                     
+       [75] "mean_horizontal"           "median"                   
+       [77] "mem_address"               "min"                      
+       [79] "min_horizontal"            "n_unique"                 
+       [81] "numeric_dtypes"            "raw_list"                 
+       [83] "read_csv"                  "read_ipc"                 
+       [85] "read_ndjson"               "read_parquet"             
+       [87] "reduce"                    "rolling_corr"             
+       [89] "rolling_cov"               "same_outer_dt"            
+       [91] "scan_csv"                  "scan_ipc"                 
+       [93] "scan_ndjson"               "scan_parquet"             
+       [95] "select"                    "set_global_rpool_cap"     
+       [97] "show_all_public_functions" "show_all_public_methods"  
+       [99] "std"                       "struct"                   
+      [101] "sum"                       "sum_horizontal"           
+      [103] "tail"                      "thread_pool_size"         
+      [105] "time"                      "using_string_cache"       
+      [107] "var"                       "when"                     
+      [109] "with_string_cache"        
 
 ---
 

tests/testthat/test-expr_struct.R

@@ -73,8 +73,34 @@ test_that("$struct$with_fields", {
     )
   )$unnest("coords")
 
+  # same thing but with pl$field()
+  out2 = df$select(
+    pl$col("coords")$struct$with_fields(
+      pl$field("x")$sqrt(),
+      y_mul = pl$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))
   )
+
+  expect_identical(
+    out2$to_data_frame(),
+    data.frame(x = c(1, 2, 3), y = c(4, 9, 16), y_mul = c(40, 18, 48))
+  )
+})
+
+test_that("pl$field() errors if field doesn't exist", {
+  expect_grepl_error(
+    pl$DataFrame(x = c(1, 4, 9))$
+      select(coords = pl$struct("x"))$
+      select(
+      pl$col("coords")$struct$with_fields(
+        pl$field("foobar")
+      )
+    ),
+    "field not found"
+  )
 })