Implement replace_na()

NEWS.md

@@ -5,6 +5,8 @@
   * `drop_na()` (@markfairbanks, #194)
 
   * `fill()` (@markfairbanks, #197)
+
+  * `replace_na()` (@markfairbanks, #202)
 
 # dtplyr 1.1.0
 

R/replace_na.R

@@ -0,0 +1,57 @@
+#' Replace NAs with specified values
+#'
+#' @description
+#' This is a method for the tidyr `replace_na()` generic. It is translated to
+#' [data.table::fcoalesce()].
+#'
+#' Note that unlike `tidyr::replace_na()`, `data.table::fcoalesce()` cannot
+#' replace `NULL` values in lists.
+#'
+#' @inheritParams tidyr::replace_na
+#' @param data A [lazy_dt()].
+#' @examples
+#' library(tidyr)
+#'
+#' # Replace NAs in a data frame
+#' dt <- lazy_dt(tibble(x = c(1, 2, NA), y = c("a", NA, "b")))
+#' dt %>% replace_na(list(x = 0, y = "unknown"))
+#'
+#' # Replace NAs using `dplyr::mutate()`
+#' dt %>% dplyr::mutate(x = replace_na(x, 0))
+# exported onLoad
+replace_na.dtplyr_step <- function(data, replace = list()) {
+
+  stopifnot(is.list(replace))
+  if (length(replace) == 0) {
+    return(data)
+  }
+
+  sim_data <- simulate_vars(data)
+  replace_vars <- intersect(names(replace), names(sim_data))
+
+  replace_calls <- vector("list", length(replace_vars))
+  names(replace_calls) <- replace_vars
+
+  for (i in seq_along(replace_vars)) {
+    var <- replace_vars[[i]]
+    check_replacement(replace[[i]], var)
+    replace_calls[[i]] <- call2("fcoalesce", sym(var), replace[[i]])
+  }
+
+  mutate(data, !!!replace_calls)
+}
+
+# exported onLoad
+replace_na.data.table <- function(data, replace = list()) {
+  data <- lazy_dt(data)
+  tidyr::replace_na(data, replace = replace)
+}
+
+check_replacement <- function(x, var) {
+  n <- length(x)
+  if (n == 1) {
+    return()
+  }
+
+  abort(glue::glue("Replacement for `{var}` is length {n}, not length 1"))
+}

R/tidyeval.R

@@ -15,7 +15,7 @@ dt_eval <- function(x) {
 # even when data.table isn't attached
 dt_funs <- c(
   "copy", "dcast", "nafill",
-  "fcase", "fintersect", "frank", "frankv", "fsetdiff", "funion",
+  "fcase", "fcoalesce", "fintersect", "frank", "frankv", "fsetdiff", "funion",
   "setcolorder", "setnames"
 )
 add_dt_wrappers <- function(env) {
@@ -104,7 +104,7 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
     } else if (is_symbol(x[[2]], ".env")) {
       sym(paste0("..", var))
     }
-  } else if (is_call(x, "coalesce")) {
+  } else if (is_call(x, "coalesce") || is_call(x, "replace_na")) {
     x[[1L]] <- quote(fcoalesce)
     x
   } else if (is_call(x, "case_when")) {

R/zzz.R

@@ -7,6 +7,7 @@
   register_s3_method("tidyr", "drop_na", "data.table")
   register_s3_method("tidyr", "fill", "data.table")
   register_s3_method("tidyr", "pivot_wider", "data.table")
+  register_s3_method("tidyr", "replace_na", "data.table")
 
   register_s3_method("dplyr", "filter", "dtplyr_step")
   register_s3_method("dplyr", "intersect", "dtplyr_step")
@@ -15,6 +16,7 @@
   register_s3_method("tidyr", "drop_na", "dtplyr_step")
   register_s3_method("tidyr", "fill", "dtplyr_step")
   register_s3_method("tidyr", "pivot_wider", "dtplyr_step")
+  register_s3_method("tidyr", "replace_na", "dtplyr_step")
 }
 
 register_s3_method <- function(pkg, generic, class, fun = NULL) {

tests/testthat/test-replace_na.R

@@ -0,0 +1,35 @@
+# lazy data.tables -----------------------------------------------------------
+
+test_that("empty call does nothing", {
+  tbl <- tibble(x = c(1, NA))
+  dt <- lazy_dt(tbl, "DT")
+  out <- collect(replace_na(dt))
+  expect_equal(out, tbl)
+})
+
+test_that("missing values are replaced", {
+  tbl <- tibble(x = c(1, NA))
+  dt <- lazy_dt(tbl, "DT")
+  step <- replace_na(dt, list(x = 0))
+  out <- collect(step)
+  expect_equal(show_query(step), expr(copy(DT)[, `:=`(x = fcoalesce(x, 0))]))
+  expect_equal(out$x, c(1, 0))
+})
+
+test_that("don't complain about variables that don't exist", {
+  tbl <- tibble(a = c(1, NA))
+  dt <- lazy_dt(tbl, "DT")
+  out <- collect(replace_na(dt, list(a = 100, b = 0)))
+  expect_equal(out, tibble(a = c(1, 100)))
+})
+
+# Inside mutate() -----------------------------------------------------------
+
+test_that("missing values are replaced", {
+  tbl <- tibble(x = c(1, NA))
+  dt <- lazy_dt(tbl, "DT")
+  step <- mutate(dt, x = replace_na(x, 0))
+  out <- collect(step)
+  expect_equal(show_query(step), expr(copy(DT)[, `:=`(x = fcoalesce(x, 0))]))
+  expect_equal(out$x, c(1, 0))
+})