分组DataFrame中使用case_when报错：条件全求值引发问题

原始数据 #

创建数据框的代码：

df <- data.frame(group   = c(1, 1, 1, 2, 2, 2),
                 var     = c(1, 2, 3, 2, 3, 4),
                 value   = c(1, 2, 3, 4, 5, 6),
                 ranking = c(1, 2, 3, 1, 2, 3))

原始输出：

group var value ranking
1     1   1     1       1
2     1   2     2       2
3     1   3     3       3
4     2   2     4       1
5     2   3     5       2
6     2   4     6       3

需求 #

按var分组处理ranking列：

• 若分组内存在group == 1的行，该分组所有行的ranking统一使用group 1的ranking值；
• 若分组只有group == 2的行，将原有ranking值加指定数值（此处为3）。

预期输出：

group var value ranking
1     1   1     1       1
2     1   2     2       2
3     1   3     3       3
4     2   2     4       2
5     2   3     5       3
6     2   4     6       6

错误代码及报错信息 #

尝试的代码：

df |> 
  group_by(var) |> 
  mutate(ranking = case_when(n() == 2   ~ ranking[group == 1],
                             group == 1 ~ ranking,
                             group == 2 ~ ranking + 3))

报错：

Error in `mutate()`:
ℹ In argument: `ranking = case_when(...)`.
ℹ In group 4: `var = 4`.
Caused by error:
! `ranking` must be size 1, not 0.

问题原因 #

case_when会对所有分支的表达式进行求值，即便当前行不满足该分支的条件。在var=4的分组中没有group==1的行，ranking[group == 1]会返回长度为0的向量，无法匹配当前行的长度要求，导致报错。

解决方案 #

方法一：先提取分组内group1的ranking值 #

先在分组内获取group1的ranking值（不存在则为NA），再通过case_when判断赋值：

library(dplyr)

df |> 
  group_by(var) |> 
  mutate(
    group1_rank = first(ranking[group == 1]),
    ranking = case_when(
      !is.na(group1_rank) ~ group1_rank,
      group == 2 ~ ranking + 3,
      TRUE ~ ranking
    )
  ) |> 
  select(-group1_rank)

方法二：用any()判断分组是否存在group1 #

直接通过any(group == 1)判断分组是否包含group1，统一赋值：

library(dplyr)

df |> 
  group_by(var) |> 
  mutate(
    ranking = ifelse(any(group == 1), 
                     rep(ranking[group == 1], n()), 
                     ranking + 3)
  )

两种方法都能得到符合预期的输出结果。

内容的提问来源于stack exchange，提问作者deschen