handle list inputs

Author: thisisnic

r/R/dplyr-funcs-conditional.R

@@ -89,11 +89,15 @@ parse_from_to_mapping <- function(x, from, to) {
   query <- vector("list", n)
   value <- vector("list", n)
   for (i in seq_len(n)) {
+    from_i <- from[[i]]
     # Handle NA specially: use is.na() since x == NA returns NA in Arrow
-    if (is.na(from[[i]])) {
+    if (length(from_i) == 1 && is.na(from_i)) {
       query[[i]] <- call_binding("is.na", x)
+    } else if (length(from_i) > 1) {
+      # Multiple values: use %in% to match any
+      query[[i]] <- call_binding("%in%", x, from_i)
     } else {
-      query[[i]] <- x == from[[i]]
+      query[[i]] <- x == from_i
     }
     value[[i]] <- Expression$scalar(to[[i]])
   }

r/tests/testthat/test-dplyr-funcs-conditional.R

@@ -581,6 +581,14 @@ test_that("replace_values()", {
     tbl
   )
 
+  # from/to with list of vectors - multiple values map to single replacement
+  compare_dplyr_binding(
+    .input |>
+      mutate(result = replace_values(chr, from = list(c("a", "b"), "c"), to = c("AB", "C"))) |>
+      collect(),
+    tbl
+  )
+
   # unmatched values kept
   compare_dplyr_binding(
     .input |>
@@ -630,6 +638,11 @@ test_that("replace_values()", {
   )
 
   # validation errors
+  expect_arrow_eval_error(
+    replace_values(chr, "A"),
+    "Each argument to replace_values\\(\\) must be a two-sided formula",
+    class = "validation_error"
+  )
   expect_arrow_eval_error(
     replace_values(chr, "a" ~ "A", from = "b"),
     "Can't use both `...` and `from`/`to` in replace_values\\(\\)",
@@ -659,6 +672,14 @@ test_that("recode_values()", {
     tbl
   )
 
+  # from/to with list of vectors - multiple values map to single replacement
+  compare_dplyr_binding(
+    .input |>
+      mutate(result = recode_values(chr, from = list(c("a", "b"), "c"), to = c("AB", "C"))) |>
+      collect(),
+    tbl
+  )
+
   # custom default
   compare_dplyr_binding(
     .input |>
@@ -697,6 +718,11 @@ test_that("recode_values()", {
     "`\\.\\.\\.` can't be empty",
     class = "validation_error"
   )
+  expect_arrow_eval_error(
+    recode_values(chr, "A"),
+    "Each argument to recode_values\\(\\) must be a two-sided formula",
+    class = "validation_error"
+  )
   expect_arrow_eval_error(
     recode_values(chr, "a" ~ "A", from = "b"),
     "Can't use both `...` and `from`/`to` in recode_values\\(\\)",