如何在data.table中用含if与shift的函数实现列条件赋值

问题背景 #

想通过data.table工具，用包含if语句的函数基于两列值生成新列的条件输出，以此学习data.table用法，但尝试的代码报错，希望了解正确的函数实现方式。

尝试的代码 #

library(data.table)

# 数据
DT <- data.table(V1 = sample(LETTERS[1:3], 20, replace = TRUE),
                 V2 = sample(1:5, 20, replace = TRUE))

# 自定义函数
fun1 <- function(x, y){
  if(x == "C" & shift(y,  type = "lead") > y){
    return("Greater")
  } else if(x == "C" & shift(y,  type = "lead") < y){
    return("Lesser")
  } else{
    return(NA)
  }
}

# 函数调用
DT.v1 <- DT[, V3 := mapply(fun1, x = V1, y = V2)]

报错信息 #

Error in if (x == "C" & shift(y, type = "lead") > y) { : 
  missing value where TRUE/FALSE needed

已有解决方案 #

猜测错误源于最后一次迭代中shift(y, type = "lead")返回NA，添加!is.na(shift(y, type = "lead"))条件后错误消失，但仅生成NULL值。目前已通过两种方式实现预期输出：

方案1：链式赋值

DT.v2 <- DT[V1 == "C" & shift(V2, type = "lead") > V2, V3 := "Greater"][
            V1 == "C" & shift(V2, type = "lead") < V2, V3 := "Lesser"]

方案2：嵌套ifelse

DT.v3 <- DT[, V3 := ifelse(V1 == "C" & shift(V2, type = "lead") > V2, "Greater",
                           ifelse(V1 == "C" & shift(V2, type = "lead") < V2, "Lesser", NA))]

问题原因与正确实现方式 #

核心问题 #

用mapply调用fun1时，每个y参数是单个数值而非整个V2列，shift(y, type = "lead")对单个值调用无意义，直接返回NA，导致if条件判断时出现缺失值触发错误。另外，shift是data.table的向量级函数，需对整列操作，而非逐元素传入函数。

正确实现方式 #

思路1：向量级函数（推荐，符合data.table设计逻辑）

把函数改成接收整列向量，内部用data.table::fcase（高效清晰的多条件判断函数，需data.table 1.12.0+版本）实现向量化判断，直接在data.table的j表达式中调用：

fun1 <- function(x_vec, y_vec){
  # 先计算整列的lead shift结果
  y_lead <- shift(y_vec, type = "lead")
  # 用fcase做多条件判断
  fcase(
    x_vec == "C" & y_lead > y_vec, "Greater",
    x_vec == "C" & y_lead < y_vec, "Lesser",
    default = NA_character_
  )
}

# 直接调用函数，无需mapply
DT[, V3 := fun1(V1, V2)]

思路2：逐元素if逻辑（适合理解基础逻辑）

若坚持用逐元素的if语句，需提前把shift后的结果生成新列，再将该列与V1、V2一起传入逐元素处理的函数：

# 先生成lead列
DT[, V2_lead := shift(V2, type = "lead")]

# 逐元素处理的函数
fun2 <- function(x, y, y_lead){
  if(is.na(y_lead)){
    return(NA)
  } else if(x == "C" & y_lead > y){
    return("Greater")
  } else if(x == "C" & y_lead < y){
    return("Lesser")
  } else{
    return(NA)
  }
}

# 用mapply调用逐元素函数
DT[, V3 := mapply(fun2, x = V1, y = V2, y_lead = V2_lead)]
# 或者用Vectorize包装函数实现向量化调用
DT[, V3 := Vectorize(fun2)(V1, V2, V2_lead)]

关键要点 #

• data.table的优势是向量级操作，尽量避免逐元素循环（如mapply），否则会降低效率；
• shift是对整列向量操作的函数，不能在逐元素的函数内部调用；
• fcase比嵌套ifelse更高效易读，是data.table推荐的多条件判断方式。

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