data.table中如何将.SD语法改写为for循环结合set函数的语法？

你原来的代码逻辑是按id+cof分组，对每个x开头的列，统计该列值为1的行中code的唯一值数量，再乘以当前组固定的cof值得到聚合结果。用for+set的等效实现如下：

方法1：生成和原foo.out完全一致的聚合结果 #

library(data.table)
set.seed(123)
# 原示例数据构造
foo.in <- data.table(id = rep(LETTERS[1:10], each = 10),
                     code = sample(c(111, 222, 333, 444), 100, replace = TRUE),
                     cof = rep(runif(10, 1, 10), each = 10),
                     x1 = sample(0:1, 100, replace = TRUE),
                     x2 = sample(0:1, 100, replace = TRUE),
                     x3 = sample(0:1, 100, replace = TRUE),
                     x4 = sample(0:1, 100, replace = TRUE),
                     x5 = sample(0:1, 100, replace = TRUE),
                     x6 = sample(0:1, 100, replace = TRUE))

# 第一步：提取所有唯一分组
foo.out2 <- unique(foo.in[, .(id, cof)])
cols <- paste0("x", 1:6)
# 初始化待填充的列
foo.out2[, (cols) := 0]

# 循环遍历每个x列计算聚合值
for (col in cols) {
  # 按分组计算当前列的结果
  cur_res <- foo.in[, .(val = uniqueN(code[get(col) == 1]) * first(cof)), by = .(id, cof)]$val
  # 用set赋值到结果表
  set(foo.out2, j = col, value = cur_res)
}

# 验证结果和原实现完全一致
all.equal(foo.out2, foo.out, check.attributes = FALSE)
# 输出为 TRUE 说明结果匹配

方法2：如果需要把聚合结果回填到原表的每一行 #

如果你的需求是在原foo.in中新增列存储对应组的统计值，可以用下面的写法：

cols_x <- paste0("x", 1:6)
new_cols <- paste0("stat_", cols_x)
# 初始化新列
foo.in[, (new_cols) := 0]
# 生成分组ID
foo.in[, grp_id := .GRP, by = .(id, cof)]
max_grp <- max(foo.in$grp_id)

for (col in cols_x) {
  grp_res <- numeric(max_grp)
  # 按分组计算当前列的统计值
  for (g in 1:max_grp) {
    idx <- foo.in$grp_id == g
    grp_res[g] <- uniqueN(foo.in$code[idx][foo.in[[col]][idx] == 1]) * foo.in$cof[idx][1]
  }
  # 批量赋值到原表
  set(foo.in, j = paste0("stat_", col), value = grp_res[foo.in$grp_id])
}

性能提示 #

对于超大规模数据集，如果分组数不多、x列数量大的场景，第二种写法可以避免多次构造.SD子集，内存占用会更优；如果分组数极多，建议先做预聚合再关联，性能会比逐组循环更好。你可以根据自己的数据规模用microbenchmark包做基准测试对比两种写法的差异。

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