GH-49534: [R] Implement dplyr recode_values(), replace_values(), and replace_when() (#49536)

### Rationale for this change

Implement new dplyr functions

### What changes are included in this PR?

Implement them

### Are these changes tested?

Yeah

### Are there any user-facing changes?

Moar functions

### AI Use

Code generated using Claude, with plenty of input from me. I've gone through it in detail and refactored lots, but it needs a last pass before it's ready for review.

* GitHub Issue: #49534

Lead-authored-by: Nic Crane <thisisnic@gmail.com>
Co-authored-by: Jonathan Keane <jkeane@gmail.com>
Signed-off-by: Nic Crane <thisisnic@gmail.com>

Author: thisisnic

r/NAMESPACE

@@ -432,6 +432,7 @@ importFrom(bit64,str.integer64)
 importFrom(glue,glue)
 importFrom(methods,as)
 importFrom(purrr,as_mapper)
+importFrom(purrr,compact)
 importFrom(purrr,flatten)
 importFrom(purrr,imap)
 importFrom(purrr,imap_chr)

r/R/arrow-package.R

@@ -18,7 +18,7 @@
 #' @importFrom stats quantile median na.omit na.exclude na.pass na.fail
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map2_chr map_dbl map_dfr map_int map_lgl keep imap imap_chr
-#' @importFrom purrr flatten reduce walk
+#' @importFrom purrr compact flatten reduce walk
 #' @importFrom assertthat assert_that is.string
 #' @importFrom rlang list2 %||% is_false abort dots_n warn enquo quo_is_null enquos is_integerish quos quo
 #' @importFrom rlang eval_tidy new_data_mask syms env new_environment env_bind set_names exec

r/R/dplyr-funcs-conditional.R

@@ -15,6 +15,169 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' Parse logical condition formulas
+#'
+#' Converts condition ~ value formulas into Arrow expressions. Unlike
+#' [parse_value_mapping()], the LHS must be a logical expression (not a value
+#' to match against).
+#'
+#' @param formulas A list of two-sided formulas where LHS is a logical condition
+#'   and RHS is the value to use when TRUE (e.g., `x > 5 ~ "high"`).
+#' @param mask The data mask for evaluating formula expressions.
+#'
+#' @return A list with `query` (list of logical expressions) and `value`
+#'   (list of replacement expressions).
+#'
+#' @keywords internal
+#' @noRd
+parse_condition_formulas <- function(formulas, mask) {
+  fn <- call_name(rlang::caller_call())
+  # Compact NULL entries (allows conditional formulas like: if (cond) x ~ y)
+  formulas <- compact(formulas)
+  n <- length(formulas)
+  query <- vector("list", n)
+  value <- vector("list", n)
+  # Process each formula: condition ~ value
+  for (i in seq_len(n)) {
+    f <- formulas[[i]]
+    if (!is_formula(f, lhs = TRUE)) {
+      validation_error(paste0("Each argument to ", fn, "() must be a two-sided formula"))
+    }
+    # f[[2]] is LHS (logical condition), f[[3]] is RHS (value when TRUE)
+    query[[i]] <- arrow_eval(f[[2]], mask)
+    value[[i]] <- arrow_eval(f[[3]], mask)
+    # Validate LHS is logical (unlike parse_value_mapping which does equality matching)
+    if (!call_binding("is.logical", query[[i]])) {
+      validation_error(paste0("Left side of each formula in ", fn, "() must be a logical expression"))
+    }
+  }
+  list(query = query, value = value)
+}
+
+#' Create case_when Expression from query/value lists
+#' @param query List of logical Arrow Expressions.
+#' @param value List of value Arrow Expressions.
+#' @return An Arrow Expression representing the case_when.
+#' @keywords internal
+#' @noRd
+build_case_when_expr <- function(query, value) {
+  Expression$create(
+    "case_when",
+    args = c(
+      Expression$create(
+        "make_struct",
+        args = query,
+        options = list(field_names = as.character(seq_along(query)))
+      ),
+      value
+    )
+  )
+}
+
+#' Build a match expression for x against a value (scalar, NA, or vector).
+#' @param x Arrow Expression for the column to match against.
+#' @param match_value Value to match - R scalar, vector, or NA. Expressions
+#'   are compared with equality.
+#' @return Arrow Expression that is TRUE when x matches match_value.
+#' @keywords internal
+#' @noRd
+build_match_expr <- function(x, match_value) {
+  # Expressions or length-1 non-NA: use equality directly
+  if (inherits(match_value, "Expression") || length(match_value) == 1 && !is.na(match_value)) {
+    return(x == match_value)
+  }
+
+  # R scalar NA requires is.na() since x == NA returns NA in Arrow
+  if (length(match_value) == 1) {
+    return(call_binding("is.na", x))
+  }
+
+  # R vector: use %in%, handling NA separately if present
+  has_na <- any(is.na(match_value))
+  non_na_values <- match_value[!is.na(match_value)]
+
+  if (length(non_na_values) == 0) {
+    call_binding("is.na", x)
+  } else if (has_na) {
+    call_binding("%in%", x, non_na_values) | call_binding("is.na", x)
+  } else {
+    call_binding("%in%", x, match_value)
+  }
+}
+
+#' Build query/value lists from parallel from/to vectors.
+#' NA values in `from` use is.na() for matching.
+#' @param x Arrow Expression for the column to match against.
+#' @param from Vector of values to match.
+#' @param to Vector of replacement values (recycled to length of `from`).
+#' @return list(query, value) for use with build_case_when_expr().
+#' @keywords internal
+#' @noRd
+parse_from_to_mapping <- function(x, from, to) {
+  n <- length(from)
+  to <- vctrs::vec_recycle(to, n)
+  query <- map(from, ~ build_match_expr(x, .x))
+  value <- map(to, Expression$scalar)
+  list(query = query, value = value)
+}
+
+#' Build query/value lists from value ~ replacement formulas.
+#' NA values on LHS use is.na() for matching.
+#' @param x Arrow Expression for the column to match against.
+#' @param formulas List of two-sided formulas (value ~ replacement).
+#' @param mask Data mask for evaluating formula expressions.
+#' @param fn Calling function name (for error messages).
+#' @return list(query, value) for use with build_case_when_expr().
+#' @keywords internal
+#' @noRd
+parse_formula_mapping <- function(x, formulas, mask, fn) {
+  # Compact NULL entries (allows conditional formulas like: if (cond) x ~ y)
+  formulas <- compact(formulas)
+  n <- length(formulas)
+  query <- vector("list", n)
+  value <- vector("list", n)
+  for (i in seq_len(n)) {
+    f <- formulas[[i]]
+    if (!is_formula(f, lhs = TRUE)) {
+      validation_error(paste0("Each argument to ", fn, "() must be a two-sided formula"))
+    }
+    # f[[2]] is LHS (value to match), f[[3]] is RHS (replacement)
+    lhs <- arrow_eval(f[[2]], mask)
+    query[[i]] <- build_match_expr(x, lhs)
+    value[[i]] <- arrow_eval(f[[3]], mask)
+  }
+  list(query = query, value = value)
+}
+
+#' Dispatch to formula or from/to parser based on which args are provided.
+#' Returns list(query, value) or NULL if no mappings.
+#' @param x Arrow Expression for the column to match against.
+#' @param formulas List of two-sided formulas (value ~ replacement).
+#' @param from Vector of values to match (alternative to formulas).
+#' @param to Vector of replacement values (used with `from`).
+#' @param mask The data mask for evaluating formula expressions.
+#' @keywords internal
+#' @noRd
+parse_value_mapping <- function(x, formulas = list(), from = NULL, to = NULL, mask) {
+  fn <- call_name(rlang::caller_call())
+  # Mutually exclusive interfaces
+  if (length(formulas) > 0 && !is.null(from)) {
+    validation_error(paste0("Can't use both `...` and `from`/`to` in ", fn, "()"))
+  }
+
+  if (length(formulas) > 0) {
+    parse_formula_mapping(x, formulas, mask, fn)
+  } else if (!is.null(from)) {
+    if (is.null(to)) {
+      validation_error("`to` must be provided when using `from`")
+    }
+    parse_from_to_mapping(x, from, to)
+  } else {
+    # No mappings provided
+    NULL
+  }
+}
+
 register_bindings_conditional <- function() {
   register_binding("%in%", function(x, table) {
     # We use `is_in` here, unlike with Arrays, which use `is_in_meta_binary`
@@ -133,44 +296,79 @@ register_bindings_conditional <- function() {
       }
 
       formulas <- list2(...)
-      n <- length(formulas)
-      if (n == 0) {
+      if (length(formulas) == 0) {
         validation_error("No cases provided")
       }
-      query <- vector("list", n)
-      value <- vector("list", n)
-      mask <- caller_env()
-      for (i in seq_len(n)) {
-        f <- formulas[[i]]
-        if (!inherits(f, "formula")) {
-          validation_error("Each argument to case_when() must be a two-sided formula")
-        }
-        query[[i]] <- arrow_eval(f[[2]], mask)
-        value[[i]] <- arrow_eval(f[[3]], mask)
-        if (!call_binding("is.logical", query[[i]])) {
-          validation_error("Left side of each formula in case_when() must be a logical expression")
-        }
-      }
+      parsed <- parse_condition_formulas(formulas, caller_env())
+      query <- parsed$query
+      value <- parsed$value
       if (!is.null(.default)) {
         if (length(.default) != 1) {
-          validation_error(paste0("`.default` must have size 1, not size ", length(.default), "."))
+          arrow_not_supported("`.default` must be size 1; vectors of length > 1")
         }
-
-        query[n + 1] <- TRUE
-        value[n + 1] <- .default
+        n <- length(query)
+        query[[n + 1]] <- TRUE
+        value[[n + 1]] <- .default
       }
-      Expression$create(
-        "case_when",
-        args = c(
-          Expression$create(
-            "make_struct",
-            args = query,
-            options = list(field_names = as.character(seq_along(query)))
-          ),
-          value
-        )
-      )
+      build_case_when_expr(query, value)
     },
     notes = "`.ptype` and `.size` arguments not supported"
   )
+
+  register_binding("dplyr::replace_when", function(x, ...) {
+    formulas <- list2(...)
+    if (length(formulas) == 0) {
+      return(x)
+    }
+    parsed <- parse_condition_formulas(formulas, caller_env())
+    query <- parsed$query
+    value <- parsed$value
+    n <- length(query)
+    query[[n + 1]] <- TRUE
+    value[[n + 1]] <- x
+    build_case_when_expr(query, value)
+  })
+
+  register_binding("dplyr::replace_values", function(x, ..., from = NULL, to = NULL) {
+    parsed <- parse_value_mapping(x, list2(...), from, to, caller_env())
+    if (is.null(parsed)) {
+      return(x)
+    }
+    query <- parsed$query
+    value <- parsed$value
+    n <- length(query)
+    query[[n + 1]] <- TRUE
+    value[[n + 1]] <- x
+    build_case_when_expr(query, value)
+  })
+
+  register_binding(
+    "dplyr::recode_values",
+    function(x, ..., from = NULL, to = NULL, default = NULL, unmatched = "default", ptype = NULL) {
+      if (!is.null(ptype)) {
+        arrow_not_supported("`recode_values()` with `ptype` specified")
+      }
+      if (unmatched != "default") {
+        arrow_not_supported('`recode_values()` with `unmatched` other than "default"')
+      }
+
+      parsed <- parse_value_mapping(x, list2(...), from, to, caller_env())
+      if (is.null(parsed)) {
+        validation_error("`...` can't be empty")
+      }
+      query <- parsed$query
+      value <- parsed$value
+
+      if (!is.null(default)) {
+        if (length(default) != 1) {
+          arrow_not_supported("`default` must be size 1; vectors of length > 1")
+        }
+        n <- length(query)
+        query[[n + 1]] <- TRUE
+        value[[n + 1]] <- Expression$scalar(default)
+      }
+      build_case_when_expr(query, value)
+    },
+    notes = "`ptype` argument and `unmatched = \"error\"` not supported"
+  )
 }

r/R/dplyr-funcs-doc.R

@@ -21,7 +21,7 @@
 #'
 #' The `arrow` package contains methods for 38 `dplyr` table functions, many of
 #' which are "verbs" that do transformations to one or more tables.
-#' The package also has mappings of 226 R functions to the corresponding
+#' The package also has mappings of 229 R functions to the corresponding
 #' functions in the Arrow compute library. These allow you to write code inside
 #' of `dplyr` methods that call R functions, including many in packages like
 #' `stringr` and `lubridate`, and they will get translated to Arrow and run
@@ -214,6 +214,9 @@
 #' * [`if_else()`][dplyr::if_else()]
 #' * [`n()`][dplyr::n()]
 #' * [`n_distinct()`][dplyr::n_distinct()]
+#' * [`recode_values()`][dplyr::recode_values()]: `ptype` argument and `unmatched = "error"` not supported
+#' * [`replace_values()`][dplyr::replace_values()]
+#' * [`replace_when()`][dplyr::replace_when()]
 #' * [`when_all()`][dplyr::when_all()]
 #' * [`when_any()`][dplyr::when_any()]
 #'