1
#' Case When with data.table
2
#'
3
#' Does what \code{dplyr::case_when()} does, with the same syntax, but with
4
#' \code{data.table::fifelse()} under the hood
5
#'
6
#' @param ... statements of the form: \code{condition ~ label}, where the label is applied if the condition is met
7
#'
8
#' @import data.table
9
#'
10
#' @examples
11
#'
12
#' x <- rnorm(100)
13
#' dt_case_when(
14
#'   x < median(x) ~ "low",
15
#'   x >= median(x) ~ "high",
16
#'   is.na(x) ~ "other"
17
#'   )
18
#'
19
#' library(data.table)
20
#' temp <- data.table(pseudo_id = c(1, 2, 3, 4, 5),
21
#'                    x = sample(1:5, 5, replace = TRUE))
22
#' temp[, y := dt_case_when(pseudo_id == 1 ~ x * 1,
23
#'                          pseudo_id == 2 ~ x * 2,
24
#'                          pseudo_id == 3 ~ x * 3,
25
#'                          pseudo_id == 4 ~ x * 4,
26
#'                          pseudo_id == 5 ~ x * 5)]
27
#'
28
#' @export
29
dt_case_when <- function(...){
30
  # grab the dots
31 1
  dots <- list(...)
32
  # checking the dots
33 1
  .check_dots(dots)
34

35
  # extract info from dots
36 1
  n <- length(dots)
37 1
  conds <- conditions(dots)
38 1
  labels <- assigned_label(dots)
39 1
  class <- class(labels)
40

41
  # make the right NA based on assigned labels
42 1
  na_type <- na_type_fun(class)
43

44
  # create fifelse() call
45 1
  calls <- call("fifelse", conds[[n]], labels[[n]], eval(na_type))
46 1
  for (i in rev(seq_len(n))[-1]){
47 1
    calls <- call("fifelse", conds[[i]], labels[[i]], calls)
48
  }
49

50 1
  eval(calls, envir = parent.frame())
51
}
52

53

54
#' fifelse from data.table
55
#'
56
#' See \code{data.table::\link[data.table:fifelse]{fifelse()}} for details.
57
#'
58
#' @name fifelse
59
#' @keywords internal
60
#' @export
61
#' @importFrom data.table fifelse
62
NULL
63

64
# Helpers -----------------
65

66
na_type_fun <- function(class){
67 1
  switch(class,
68
         "logical"   = NA,
69
         "complex"   = NA_complex_,
70
         "character" = NA_character_,
71
         "integer"   = NA_integer_,
72
         NA_real_)
73
}
74
conditions <- function(list){
75 1
  unlist(lapply(list, function(x) x[[2]]))
76
}
77
assigned_label <- function(list){
78 1
  unlist(lapply(list, function(x) x[[3]]))
79
}
80
is_formula <- function(x){
81 1
  is.call(x) && x[[1]] == quote(`~`)
82
}
83

84
# Check functions -------------------
85

86
.check_dots <- function(dots){
87 1
  forms <- all(unlist(lapply(dots, is_formula)))
88 1
  if (!forms)
89 1
    stop("Not all arguments are formulas", call. = FALSE)
90
}

Read our documentation on viewing source code .

Loading