R中data.table聚合性能优化：按组求和并保留指定行数据

优化后代码实现 #

library(data.table)

# 构造示例数据
my_data <- data.table(MY_NUM1 = 1:10,
                      MY_NUM2 = 11:20,
                      MY_CHAR1 = LETTERS[1:10],
                      MY_CHAR2 = LETTERS[11:20],
                      MY_KEY = c(rep(1:2,3), 2:3, 3:4),
                      MY_BY  = rep(letters[1:5], each = 2))

my_sum_vars <- c("MY_NUM1", "MY_NUM2")
my_by <- c("MY_BY")

# 定义非求和列集合
non_sum_cols <- setdiff(names(my_data), c(my_sum_vars, my_by))

# 给分组键建立索引（大规模数据下显著提速）
setkey(my_data, MY_BY)

# 核心高效分组逻辑
my_summary <- my_data[, {
  # 确定组内目标行：优先取MY_KEY=1的行，否则取首行
  target_idx <- if (any(MY_KEY == 1)) which(MY_KEY == 1) else 1L
  # 提取目标行的非求和列
  key_cols <- .SD[target_idx, non_sum_cols, with = FALSE]
  # 计算求和列的组内总和
  sum_cols <- lapply(.SD[, my_sum_vars, with = FALSE], sum)
  # 合并两类结果
  c(key_cols, sum_cols)
}, by = my_by]

# 查看结果
my_summary

优化点说明 #

1. 消除自定义函数开销：原代码每个分组都调用getKeyOrFirst函数，函数调用的额外开销在大规模数据下会被急剧放大。直接在分组表达式内完成逻辑，利用data.table原生向量化操作，大幅降低运行时间。
2. 减少内存冗余：通过索引定位目标行后提取对应列，而非拆分整个数据集再合并，避免了中间数据的重复存储。
3. 索引加速分组：setkey(my_data, MY_BY)为分组键建立索引，data.table会基于索引执行分组操作，比无索引的分组效率提升明显（超大规模数据集效果更显著）。
4. 精准.SD范围控制：明确指定.SD处理的列范围，减少不必要的数据加载与计算，进一步优化性能。

验证结果 #

运行上述代码后，输出与期望完全一致：

MY_BY MY_CHAR1 MY_CHAR2 MY_KEY MY_NUM1 MY_NUM2
1:     a        A        K      1       3      23
2:     b        C        M      1       7      27
3:     c        E        O      1      11      31
4:     d        G        Q      2      15      35
5:     e        I        S      3      19      39

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