Returning a tibble: how to vectorize with case_when?

Question

I have a function which returns a tibble. It runs OK, but I want to vectorize it.

library(tidyverse)

tibTest <- tibble(argX = 1:4, argY = 7:4)

square_it <- function(xx, yy) {
  if(xx >= 4){
    tibble(x = NA, y = NA)
  } else if(xx == 3){
    tibble(x = as.integer(), y = as.integer())
  } else if (xx == 2){
    tibble(x = xx^2 - 1, y = yy^2 -1)
  } else {
    tibble(x = xx^2, y = yy^2)
  }
}

It runs OK in a mutate when I call it with map2, giving me the result I wanted:

tibTest %>%
  mutate(sq = map2(argX, argY, square_it)) %>%
  unnest()
## A tibble: 3 x 4
#     argX  argY     x     y
#    <int> <int> <dbl> <dbl>
# 1     1     7     1    49
# 2     2     6     3    35
# 3     4     4    NA    NA

My first attempt to vectorize it failed, and I can see why - I can't return a vector of tibbles.

square_it2 <- function(xx, yy){
  case_when(
    x >= 4 ~ tibble(x = NA, y = NA),
    x == 3 ~ tibble(x = as.integer(), y = as.integer()),
    x == 2 ~ tibble(x = xx^2 - 1, y = yy^2 -1),
    TRUE   ~ tibble(x = xx^2,     y = yy^2)
  )
}
# square_it2(4, 2)  # FAILS

My next attempt runs OK on a simple input. I can return a list of tibbles, and that's what I want for the unnest

square_it3 <- function(xx, yy){
  case_when(
    xx >= 4 ~ list(tibble(x = NA, y = NA)),
    xx == 3 ~ list(tibble(x = as.integer(), y = as.integer())),
    xx == 2 ~ list(tibble(x = xx^2 - 1, y = yy^2 -1)),
    TRUE   ~ list(tibble(x = xx^2,     y = yy^2))
  )
}
square_it3(4, 2)
# [[1]]
# # A tibble: 1 x 2
# x     y    
# <lgl> <lgl>
#   1 NA    NA   

But when I call it in a mutate, it doesn't give me the result I had with square_it. I can sort of see what's wrong. In the xx == 2 clause, xx acts as an atomic value of 2. But in building the tibble, xx is a length-4 vector.

tibTest %>%
  mutate(sq =  square_it3(argX, argY)) %>%
  unnest()
# # A tibble: 9 x 4
#    argX  argY     x     y
#    <int> <int> <dbl> <dbl>
# 1     1     7     1    49
# 2     1     7     4    36
# 3     1     7     9    25
# 4     1     7    16    16
# 5     2     6     0    48
# 6     2     6     3    35
# 7     2     6     8    24
# 8     2     6    15    15
# 9     4     4    NA    NA

How do I get the same result as I did with square_it, but from a vectorized function using case_when ?

Answer 1

We define row_case_when which has a similar formula interface as case_when except it has a first argument of .data, acts by row and expects that the value of each leg to be a data frame. It returns a data.frame/tibble. Wrapping in a list, rowwise and unnest are not needed.

case_when2 <- function (.data, ...) {
    fs <- dplyr:::compact_null(rlang:::list2(...))
    n <- length(fs)
    if (n == 0) {
        abort("No cases provided")
    }
    query <- vector("list", n)
    value <- vector("list", n)
    default_env <- rlang:::caller_env()
    quos_pairs <- purrr::map2(fs, seq_along(fs), dplyr:::validate_formula,
        rlang:::default_env, rlang:::current_env())
    for (i in seq_len(n)) {
        pair <- quos_pairs[[i]]
        query[[i]] <- rlang::eval_tidy(pair$lhs, data = .data, env = default_env)
        value[[i]] <- rlang::eval_tidy(pair$rhs, data = .data, env = default_env)
        if (!is.logical(query[[i]])) {
            abort_case_when_logical(pair$lhs, i, query[[i]])
        }
        if (query[[i]]) return(value[[i]])
    }
}

row_case_when <- function(.data, ...) {
  .data %>% 
    group_by(.group = 1:n(), !!!.data) %>%
    do(case_when2(., ...)) %>%
    mutate %>%
    ungroup %>%
    select(-.group)
}

Test run

It is used like this:

library(dplyr)

tibTest <- tibble(argX = 1:4, argY = 7:4) # test data from question

tibTest %>%
  row_case_when(argX >= 4 ~ tibble(x = NA, y = NA),
    argX == 3 ~ tibble(x = as.integer(), y = as.integer()),
    argX == 2 ~ tibble(x = argX^2 - 1, y = argY^2 -1),
    TRUE   ~ tibble(x = argX^2,     y = argY^2)
  )

giving:

# A tibble: 3 x 4
   argX  argY     x     y
  <int> <int> <dbl> <dbl>
1     1     7     1    49
2     2     6     3    35
3     4     4    NA    NA

mutate_cond and mutate_when

These are not quite the same as row_case_when since they don't run through conditions taking the first true one but by using mutually exclusive conditions they can be used for certain aspects of this problem. They do not handle changing the number of rows in the result but we can use dplyr::filter to remove rows for a particular condition.

mutate_cond defined in dplyr mutate/replace several columns on a subset of rows is like mutate except the second argument is a condition and the subsequent arguments are applied only to rows for which that condition is TRUE.

mutate_when defined in dplyr mutate/replace several columns on a subset of rows is similar to case_when except it applies to rows, the replacement values are provided in a list and alternate arguments are conditions and lists. Also all legs are always run applying the replacement values to the rows satisfying the conditions (as opposed to, for each row, performing the replacement on just the first true leg). To get a similar effect to row_case_when be sure that the conditions are mutually exclusive.

# mutate_cond example
tibTest %>%
  filter(argX != 3) %>%
  mutate(x = NA_integer_, y = NA_integer_) %>%
  mutate_cond(argX == 2, x = argX^2 - 1L, y = argY^2 - 1L) %>%
  mutate_cond(argX < 2, x = argX^2, y = argY^2)

# mutate_when example
tibTest %>%
  filter(argX != 3) %>%
  mutate_when(TRUE, list(x = NA_integer_, y = NA_integer_),
              argX == 2, list(x = argX^2 - 1L, y = argY^2 - 1L), 
              argX < 2, list(x = argX^2, y = argY^2))