基于指定列分组高效为data.table添加计算列的优化需求

测试数据生成 #

library(data.table)
rowlength<-1e4;
set.seed(1);
toydata<-data.table(Char1=replicate(rowlength, {paste0(sample(c(letters), round(runif(1, min = 1, max = 5)), replace = TRUE), collapse = "")}),
Char2=replicate(rowlength, {paste0(sample(c(letters, c(1:5)), round(runif(1, min = 1, max = 5)), replace = TRUE), collapse = "")}),
Int1=rlnorm(rowlength,meanlog = 10, sdlog = 2),
Int2=rlnorm(rowlength,meanlog = 10, sdlog = 2),
Int3=round(runif(rowlength,min = 1,max = 1000))*1e4,
Int4=round(runif(rowlength,min = 1,max = 1000))*1e4);

原实现的问题 #

原代码用by = toydata[,.SD, .SDcols = pickcol]作为分组键，这种方式会为每个分组生成子表，在500万行的大数据量下会产生极大的内存和计算开销，导致运行极慢。需求明确：

• 按外部传入的pickcol（支持1个或多个列）分组
• 每组仅首行计算sum(pmin(Int4, pmax(0, Int1-Int3)))，其余行设为0
• 只能用data.table包实现

优化后的代码 #

pickcol<-c("Char2","Char1");

# 先初始化整列为0，避免重复生成0向量
toydata[, needthis := 0]
# 提取每组首行的索引，直接赋值计算结果
toydata[toydata[, .I[1], by = pickcol]$V1, needthis := sum(pmin(Int4, pmax(0, Int1-Int3))), by = pickcol]

优化说明 #

1. 先统一初始化needthis为0，省去原代码中每组生成rep(0, .N-1)的重复操作
2. 用.I[1]直接获取每组的首行索引，通过索引定位赋值是data.table的高效操作方式
3. 分组键直接使用by = pickcol，无需生成子表，大幅降低内存占用和计算时间

验证输出 #

执行以下代码验证结果：

summary(toydata$needthis[toydata$needthis != 0])

得到结果：

Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
   2930  115218  579795 1056456 1249714 7390000 

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