r/data.table按首组返回值设置列类型 首组返回NA时类型错误如何解决

问题原因 #

R 中默认的NA是逻辑类型，data.table 按组使用:=赋值时，会以第一个分组的返回值类型确定新建列的整体类型。当第一个分组返回默认逻辑型NA时，整列会被定为逻辑型，后续分组返回的数值会被强制转换为逻辑值（非0数值转逻辑值均为TRUE），就会出现你遇到的异常。

不需要提前预设列类型的解决方案 #

方案1：修改自定义函数，返回对应类型的缺失值（最优） #

R 提供了对应不同数据类型的专用缺失值：整型用NA_integer_、浮点型用NA_real_、字符型用NA_character_，只需要把函数中返回的普通NA替换为和正常返回值类型一致的专用缺失值即可：

library(data.table)

myfun <- function(x) {
    if(x == 0) {
        # 正常返回值为x*2属于浮点型，所以用浮点型缺失值
        # 如果需要返回整型，可改为return(NA_integer_)，同时正常返回改为return(as.integer(x*2))
        return(NA_real_)
    } else {
        return(x*2)
    }
}

DT <- data.table(x= c(0, 1, 2, 3), y= LETTERS[1:4])
DT[, x2 := myfun(x), by= y]

运行后无警告，x2列结果为NA,2,4,6，类型符合预期。

方案2：赋值时强制统一返回值类型 #

如果不想修改原函数，可以在赋值时直接对返回结果做类型强制转换：

# 原函数保持不变
myfun <- function(x) {
    if(x == 0) {
        return(NA)
    } else {
        return(x*2)
    }
}

DT <- data.table(x= c(0, 1, 2, 3), y= LETTERS[1:4])
# 赋值时统一转为浮点型，也可根据需求改为as.integer
DT[, x2 := as.numeric(myfun(x)), by= y]

同样可以得到正确结果，不需要提前创建列。

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