Merge branch 'main' into pl-concat-str

R/000-wrappers.R

@@ -48,6 +48,12 @@ NULL
 }
 
 
+`arg_where` <- function(`condition`) {
+  `condition` <- .savvy_extract_ptr(`condition`, "PlRExpr")
+  .savvy_wrap_PlRExpr(.Call(savvy_arg_where__impl, `condition`))
+}
+
+
 `as_struct` <- function(`exprs`) {
   .savvy_wrap_PlRExpr(.Call(savvy_as_struct__impl, `exprs`))
 }

R/expr-expr.R

@@ -1510,16 +1510,16 @@ expr__arg_unique <- function() {
     wrap()
 }
 
-# TODO-REWRITE: requires pl$arg_where()
-# #' Return indices where expression is true
-# #'
-# #' @inherit as_polars_expr return
-# #' @examples
-# #' df <- pl$DataFrame(a = c(1, 1, 2, 1))
-# #' df$select((pl$col("a") == 1)$arg_true())
-# expr__arg_true <- function() {
-# pl$arg_where(self$`_rexpr`)
-# }
+#' Return indices where expression is true
+#'
+#' @inherit as_polars_expr return
+#' @examples
+#' df <- pl$DataFrame(a = c(1, 1, 2, 1))
+#' df$select((pl$col("a") == 1)$arg_true())
+expr__arg_true <- function() {
+  arg_where(self$`_rexpr`) |> 
+    wrap()
+}
 
 #' Get the number of non-null elements in the column
 #'

R/functions-lazy.R

@@ -60,3 +60,18 @@ pl__coalesce <- function(...) {
     coalesce() |>
     wrap()
 }
+
+#' Return indices where `condition` evaluates to `TRUE`
+#'
+#' @param condition Boolean expression to evaluate.
+#' @inherit as_polars_expr return
+#'
+#' @examples
+#' df <- pl$DataFrame(a = 1:5)
+#' df$select(
+#'   pl$arg_where(pl$col("a") %% 2 == 0)
+#' )
+pl__arg_where <- function(condition) {
+  arg_where(as_polars_expr(condition)$`_rexpr`) |>
+    wrap()
+}

src/init.c

@@ -44,6 +44,11 @@ SEXP savvy_any_horizontal__impl(SEXP c_arg__exprs) {
     return handle_result(res);
 }
 
+SEXP savvy_arg_where__impl(SEXP c_arg__condition) {
+    SEXP res = savvy_arg_where__ffi(c_arg__condition);
+    return handle_result(res);
+}
+
 SEXP savvy_as_struct__impl(SEXP c_arg__exprs) {
     SEXP res = savvy_as_struct__ffi(c_arg__exprs);
     return handle_result(res);
@@ -2563,6 +2568,7 @@ SEXP savvy_PlRWhen_then__impl(SEXP self__, SEXP c_arg__statement) {
 static const R_CallMethodDef CallEntries[] = {
     {"savvy_all_horizontal__impl", (DL_FUNC) &savvy_all_horizontal__impl, 1},
     {"savvy_any_horizontal__impl", (DL_FUNC) &savvy_any_horizontal__impl, 1},
+    {"savvy_arg_where__impl", (DL_FUNC) &savvy_arg_where__impl, 1},
     {"savvy_as_struct__impl", (DL_FUNC) &savvy_as_struct__impl, 1},
     {"savvy_coalesce__impl", (DL_FUNC) &savvy_coalesce__impl, 1},
     {"savvy_col__impl", (DL_FUNC) &savvy_col__impl, 1},

src/rust/Cargo.toml

@@ -28,6 +28,7 @@ rev = "841c387d99d7024037556c4ef79d96bf2caac397"
 features = [
     "abs",
     "approx_unique",
+    "arg_where",
     "array_any_all",
     "array_count",
     "array_to_struct",

src/rust/api.h

@@ -1,5 +1,6 @@
 SEXP savvy_all_horizontal__ffi(SEXP c_arg__exprs);
 SEXP savvy_any_horizontal__ffi(SEXP c_arg__exprs);
+SEXP savvy_arg_where__ffi(SEXP c_arg__condition);
 SEXP savvy_as_struct__ffi(SEXP c_arg__exprs);
 SEXP savvy_coalesce__ffi(SEXP c_arg__exprs);
 SEXP savvy_col__ffi(SEXP c_arg__name);

src/rust/src/functions/lazy.rs

@@ -270,3 +270,8 @@ pub fn concat_str(s: ListSexp, separator: &str, ignore_nulls: bool) -> Result<Pl
     let s = <Wrap<Vec<Expr>>>::from(s).0;
     Ok(dsl::concat_str(s, separator, ignore_nulls).into())
 }
+
+#[savvy]
+pub fn arg_where(condition: PlRExpr) -> Result<PlRExpr> {
+    Ok(dsl::arg_where(condition.inner.clone()).into())
+}

tests/testthat/test-expr-expr.R

@@ -1338,6 +1338,14 @@ test_that("arg_unique", {
   )
 })
 
+test_that("arg_true", {
+  df <- pl$DataFrame(a = c(1, 1, 2, 1))
+  expect_equal(
+    df$select((pl$col("a") == 1)$arg_true()),
+    pl$DataFrame(a = c(0, 1, 3))$cast(pl$UInt32)
+  )
+})
+
 # test_that("Expr_quantile", {
 #   v <- sample(0:100)
 #   expect_equal(

tests/testthat/test-functions-lazy.R

@@ -18,3 +18,11 @@ test_that("pl$coalesce()", {
     "must be passed by position, not name"
   )
 })
+
+test_that("arg_where", {
+  df <- pl$DataFrame(a = 1:5)
+  expect_equal(
+    df$select(pl$arg_where(pl$col("a") %% 2 == 0)),
+    pl$DataFrame(a = c(1, 3))$cast(pl$UInt32)
+  )
+})