Support data frames in coalesce()

Closes tidyverse#5326

NEWS.md

@@ -2,7 +2,9 @@
 
 ## Minor improvements and bug fixes
 
-* cummean() no longer has off-by-one indexing problem (#5287).
+* `coalesce()` now supports data frames correctly (#5326).
+
+* `cummean()` no longer has off-by-one indexing problem (#5287).
 
 # dplyr 1.0.0
 

R/coalesce.R

@@ -5,8 +5,9 @@
 #' which does the same thing for `NULL`s.
 #'
 #' @param ... <[`dynamic-dots`][rlang::dyn-dots]> Vectors. Inputs should be
-#'    recyclable (either be length 1 or same length as the longest vector) and
-#'    coercible to a common type.
+#'   recyclable (either be length 1 or same length as the longest vector) and
+#'   coercible to a common type. If data frames, they are coalesced column by
+#'   column.
 #' @return A vector the same length as the first `...` argument with
 #'   missing values replaced by the first non-missing value.
 #' @seealso [na_if()] to replace specified values with a `NA`.
@@ -40,9 +41,29 @@ coalesce <- function(...) {
   x <- values[[1]]
   values <- values[-1]
 
+  if (!is_null(attr(x, "dim"))) {
+    abort("Can't coalesce matrices or arrays.")
+  }
+  if (is.data.frame(x)) {
+    df_coalesce(x, values)
+  } else {
+    vec_coalesce(x, values)
+  }
+}
+
+vec_coalesce <- function(x, values) {
   for (i in seq_along(values)) {
     x_miss <- is.na(x)
     vec_slice(x, x_miss) <- vec_slice(values[[i]], x_miss)
   }
+
+  x
+}
+df_coalesce <- function(x, values) {
+  for (i in seq_along(x)) {
+    col_values <- map(values, `[[`, i)
+    x[[i]] <- vec_coalesce(x[[i]], col_values)
+  }
+
   x
 }

tests/testthat/test-coalesce.R

@@ -26,3 +26,28 @@ test_that("coalesce() gives meaningful error messages", {
     coalesce(1:2, letters[1:2])
   })
 })
+
+test_that("coalesce() supports data frames (#5326)", {
+  out <- coalesce(
+    data.frame(x = c(NA, 1)),
+    data.frame(x = 1:2)
+  )
+  expect_identical(out, data.frame(x = c(1, 1)))
+
+  df1 <- data.frame(x = c(NA, 1, NA), y = c(2, NA, NA), z = c(1:2, NA))
+  df2 <- tibble::tibble(x = 1:3, y = c(3, 4, NA), z = c(NA, NA, NA))
+  df3 <- data.frame(x = NA, y = c(30, 40, 50), z = 101:103)
+  out <- coalesce(df1, df2, df3)
+  exp <- tibble(x = c(1, 1, 3), y = c(2, 4, 50), z = c(1L, 2L, 103L))
+  expect_identical(out, exp)
+
+  expect_error(
+    coalesce(
+      data.frame(x = c(NA, 1)),
+      data.frame(x = c("a", "b"))
+    ),
+    class = "vctrs_error_incompatible_type"
+  )
+
+  expect_error(coalesce(as.matrix(mtcars), as.matrix(mtcars)), "matrices")
+})