使用data.table按组计算多变量非缺失观测值结果异常的问题排查

我用data.table按group分组计算指定变量的非缺失观测值数量，编写了如下代码：

library(data.table)
dat <- fread("col1 col2 col3 group
       1    2    4    A
       3    2    2    A
       1    NA   1    B
       3    2    1    B")

vars_of_interest <- c("col1", "col2")
vars_of_interest_obs <- paste0(vars_of_interest, "_obs_tot")

dat <- setDT(dat)[, (vars_of_interest_obs) := sum(!is.na(vars_of_interest)), by = c("group")]

运行后得到结果：

col1 col2 col3 group col1_obs_tot col2_obs_tot
1:    1    2    4     A            2            2
2:    3    2    2     A            2            2
3:    1   NA    1     B            2            2
4:    3    2    1     B            2            2

但col2_obs_tot的正确结果应该是：

col2_obs_tot
2
2
1
1

请问代码哪里出错了？

错误原因 #

代码中sum(!is.na(vars_of_interest))存在逻辑错误：vars_of_interest是一个字符向量（c("col1", "col2")），is.na(vars_of_interest)只会检查这个字符向量本身的元素是否为缺失值，而不会去检查数据框中对应列的缺失值。因为该字符向量内没有NA，所以sum(!is.na(vars_of_interest))始终返回2，导致所有分组的结果都错误。

修正代码 #

需要针对每个指定变量，分别计算分组内的非缺失值数量，以下是两种可行方法：

方法1：使用.SD和.SDcols

library(data.table)
dat <- fread("col1 col2 col3 group
       1    2    4    A
       3    2    2    A
       1    NA   1    B
       3    2    1    B")

vars_of_interest <- c("col1", "col2")
vars_of_interest_obs <- paste0(vars_of_interest, "_obs_tot")

dat[, (vars_of_interest_obs) := lapply(.SD, function(x) sum(!is.na(x))), 
    by = group, .SDcols = vars_of_interest]

方法2：循环处理每个变量

library(data.table)
dat <- fread("col1 col2 col3 group
       1    2    4    A
       3    2    2    A
       1    NA   1    B
       3    2    1    B")

vars_of_interest <- c("col1", "col2")

for (var in vars_of_interest) {
  dat[, paste0(var, "_obs_tot") := sum(!is.na(get(var))), by = group]
}

运行修正后的代码，得到正确结果：

col1 col2 col3 group col1_obs_tot col2_obs_tot
1:    1    2    4     A            2            2
2:    3    2    2     A            2            2
3:    1   NA    1     B            2            1
4:    3    2    1     B            2            1

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