Refresh coalesce() on top of vctrs::vec_case_when() (#7725)

Author: DavisVaughan

NEWS.md

@@ -1,8 +1,11 @@
 # dplyr (development version)
 
-* `if_else()` no longer allows `condition` to be a logical array. It must be a logical vector with no `dim` attribute (#7723).
+* The following vector functions have gotten significantly faster and use much less memory due to a rewrite in C via vctrs (#7723):
+
+  * `if_else()`
+  * `coalesce()`
 
-* `if_else()` has gotten significantly faster and uses much less memory due to a rewrite in C via `vctrs::vec_if_else()` (#7723).
+* `if_else()` no longer allows `condition` to be a logical array. It must be a logical vector with no `dim` attribute (#7723).
 
 * Passing `size` to `if_else()` is now deprecated. The output size is always taken from the `condition` (#7722).
 

R/coalesce.R

@@ -45,36 +45,125 @@
 coalesce <- function(..., .ptype = NULL, .size = NULL) {
   args <- list2(...)
 
-  if (vec_any_missing(args)) {
-    # Drop `NULL`s
-    not_missing <- !vec_detect_missing(args)
-    args <- vec_slice(args, not_missing)
-  }
-
   if (length(args) == 0L) {
     abort("`...` can't be empty.")
   }
+  if (vec_all_missing(args)) {
+    abort("`...` must contain at least 1 non-`NULL` value.")
+  }
+
+  # We do vector, type, and size checks up front before dropping any `NULL`
+  # values or extracting out a `default` to ensure that any errors report
+  # the correct index
+  list_check_all_vectors(args, allow_null = TRUE, arg = "")
+
+  .ptype <- vec_ptype_common(!!!args, .ptype = .ptype)
+  args <- vec_cast_common(!!!args, .to = .ptype)
+
+  if (is_null(.size)) {
+    .size <- vec_size_common(!!!args)
+  } else {
+    # Check recyclability, but delay actual recycling
+    list_check_all_recyclable(args, .size, allow_null = TRUE, arg = "")
+  }
 
-  # Recycle early so logical conditions computed below will be the same length,
-  # as required by `vec_case_when()`
-  args <- vec_recycle_common(!!!args, .size = .size)
+  # From this point on we don't expect any errors
 
-  # Name early to get correct indexing in `vec_case_when()` error messages
-  names <- names2(args)
-  names <- names_as_error_names(names)
-  args <- set_names(args, names)
+  args <- convert_from_coalesce_to_case_when(args, .size)
+  values <- args$values
+  default <- args$default
 
-  conditions <- map(args, function(arg) {
-    !vec_detect_missing(arg)
+  cases <- map(values, function(value) {
+    !vec_detect_missing(value)
   })
 
-  vec_case_when(
-    conditions = conditions,
-    values = args,
-    conditions_arg = "",
-    values_arg = "",
+  vctrs::vec_case_when(
+    cases = cases,
+    values = values,
+    default = default,
     ptype = .ptype,
-    size = .size,
-    call = current_env()
+    size = .size
   )
 }
+
+# Goal is to convert from `...` of `coalesce()` to `values` and `default`
+# of `vec_case_when()`
+#
+# Recognize that these are equivalent:
+#
+# ```
+# coalesce(x, y)
+# case_when(!vec_detect_missing(x) ~ x, !vec_detect_missing(y) ~ y)
+#
+# coalesce(x, y_with_no_missings)
+# case_when(!vec_detect_missing(x) ~ x, .default = y_with_no_missings)
+#
+# coalesce(x, NULL, y, 0)
+# case_when(!vec_detect_missing(x) ~ x, !vec_detect_missing(y) ~ y, .default = 0)
+# ```
+#
+# Note how the last element can be used as `default` if it doesn't contain any
+# missing values. This is a very nice optimization since `vec_case_when()`
+# doesn't need to recycle that value, and efficiently computes its output
+# locations!
+#
+# Note how `NULL`s are dropped during the conversion.
+convert_from_coalesce_to_case_when <- function(args, size) {
+  if (vec_any_missing(args)) {
+    # Drop `NULL`
+    args <- vec_slice(args, vec_detect_complete(args))
+  }
+
+  args_size <- length(args)
+
+  if (args_size == 0L) {
+    abort("Checked for at least 1 non-`NULL` value earlier", .internal = TRUE)
+  }
+
+  # Try to promote the `last` element of `args` to `default`
+  #
+  # For the 99% case of `coalesce(x, 0)`, this:
+  # - Avoids recycling `0` to size `size`.
+  # - Avoids computing `!vec_detect_missing()` on that recycled `0`.
+  #
+  # Can only do this if the `last` element doesn't contain missing values
+  # due to how names are handled. We don't want to take the name from any `NA`
+  # element, which is what would happen if we promoted the whole `y` vector here
+  # to `default`.
+  #
+  # ```
+  # x <- c(a = NA, b = 2)
+  # y <- c(c = NA, d = 4)
+  #
+  # coalesce(x, y)
+  # # Want c(NA, b = 2)
+  # # Not c(c = NA, b = 2)
+  #
+  # # Compare to
+  # case_when(!vec_detect_missing(x) ~ x, !vec_detect_missing(y) ~ y)
+  # case_when(!vec_detect_missing(x) ~ x, .default = y)
+  # ```
+  last <- args[[args_size]]
+
+  if (vec_any_missing(last)) {
+    default <- NULL
+  } else {
+    default <- last
+    args <- args[-args_size]
+  }
+
+  # Most of the time this recycle is a no-op. Two cases where it isn't:
+  # - `coalesce(x, 0, 1)`, where `1` becomes `default` but we still have a
+  #   scalar `0`.
+  # - `coalesce(x, NA)`, where `NA` can't be promoted, so we have a scalar `NA`.
+  args <- vec_recycle_common(!!!args, .size = size)
+
+  list(values = args, default = default)
+}
+
+vec_all_missing <- function(x) {
+  if (!vec_any_missing(x)) {
+    return(FALSE)
+  }
+  sum(vec_detect_missing(x)) == vec_size(x)
+}

tests/testthat/_snaps/coalesce.md

@@ -1,15 +1,16 @@
 # coalesce() gives meaningful error messages
 
     Code
-      (expect_error(coalesce(1:2, 1:3)))
-    Output
-      <error/vctrs_error_incompatible_size>
+      coalesce(1:2, 1:3)
+    Condition
       Error in `coalesce()`:
       ! Can't recycle `..1` (size 2) to match `..2` (size 3).
+
+---
+
     Code
-      (expect_error(coalesce(1:2, letters[1:2])))
-    Output
-      <error/vctrs_error_ptype2>
+      coalesce(1:2, letters[1:2])
+    Condition
       Error in `coalesce()`:
       ! Can't combine `..1` <integer> and `..2` <character>.
 
@@ -21,21 +22,21 @@
       Error in `coalesce()`:
       ! Can't recycle `..2` (size 2) to size 1.
 
-# must have at least one non-`NULL` vector
+# can't be empty
 
     Code
       coalesce()
     Condition
       Error in `coalesce()`:
       ! `...` can't be empty.
 
----
+# must have at least one non-`NULL` vector
 
     Code
       coalesce(NULL, NULL)
     Condition
       Error in `coalesce()`:
-      ! `...` can't be empty.
+      ! `...` must contain at least 1 non-`NULL` value.
 
 # inputs must be vectors
 
@@ -61,3 +62,51 @@
       Error in `coalesce()`:
       ! Can't combine `..1` <double> and `y` <character>.
 
+---
+
+    Code
+      coalesce(1:2, 1:3)
+    Condition
+      Error in `coalesce()`:
+      ! Can't recycle `..1` (size 2) to match `..2` (size 3).
+
+---
+
+    Code
+      coalesce(1:2, y = 1:3)
+    Condition
+      Error in `coalesce()`:
+      ! Can't recycle `..1` (size 2) to match `y` (size 3).
+
+---
+
+    Code
+      coalesce(1, NULL, "x")
+    Condition
+      Error in `coalesce()`:
+      ! Can't combine `..1` <double> and `..3` <character>.
+
+---
+
+    Code
+      coalesce(1, NULL, y = "x")
+    Condition
+      Error in `coalesce()`:
+      ! Can't combine `..1` <double> and `y` <character>.
+
+---
+
+    Code
+      coalesce(1:2, NULL, 1:3)
+    Condition
+      Error in `coalesce()`:
+      ! Can't recycle `..1` (size 2) to match `..3` (size 3).
+
+---
+
+    Code
+      coalesce(1:2, NULL, y = 1:3)
+    Condition
+      Error in `coalesce()`:
+      ! Can't recycle `..1` (size 2) to match `y` (size 3).
+