基于data.table重构Epi包ccwc()函数的内存优化咨询

算力阿峰

2026-6-17

优化data.table实现巢式病例对照（ccwc）的内存效率

我正在用data.table重构Epi包的ccwc()函数，用于从队列数据集构建巢式病例对照数据集。该函数支持定义时间尺度（entry/exit）、对照样本量（示例为10）、匹配变量（示例为job），并生成Set列标识病例-对照组（相同退出日期的病例归为同一Set）。目前核心逻辑已实现，但构建对照数据集时，非等值连接会导致数据集大幅扩容，后续需用sample()缩减，希望获得内存效率优化方案，最好无需拆分病例和对照数据集，直接在原diet数据集上操作。

原代码参考

数据与包加载

#load packages 
library(Epi)
library(data.table)
library(survival)
    
#data
data(diet)
setDT(diet)
setkey(diet,id)

#Rename time columns
diet$entry<-as.numeric(diet$doe-diet$dob)
diet$exit<-as.numeric(diet$dox-diet$dob)

Epi包ccwc()实现

dietcc <- ccwc( entry, exit, chd, controls=10, data=diet,
                include=energy, match=job)

现有data.table尝试（内存瓶颈在非等值连接）

#Select cases and define risk sets
case<-setorder(diet[chd==1], exit,job)[,
                .(id,exitcase=exit,job, event=1,set=rleid(exit,job))][,
                 mult:=.N, by=set]

#Select controls merging them with the cases
control<-diet[case[!duplicated(set,mult),.(id,exitcase,job,mult,set)], 
              .(x.entry, x.exit, x.job, x.id,i.id, mult=i.mult, i.exitcase, i.job),
              on = .( entry <= exitcase,exit > exitcase,job = job),
              by=.EACHI, nomatch=0][,
              .SD[sample(.N, min(10*mult,.N))],by = i.id][ ,set:=rleid(i.id)][,
               .(id=x.id,job=i.job,exitcase=i.exitcase,set,event=0,mult)]

#append controls and cases together
dtcc<-rbind(case,control)
setkey(dtcc,id)

#Include additional predictor (in this case "energy")
dtcc<-dtcc[diet[.(id,energy)], energy := i.energy]

模型验证（结果一致）

#full cohort
summary(coxph(Surv(as.numeric(doe),as.numeric(dox),chd) ~ scale(energy) + strata(job), data=diet))

#ccwc solution
summary(clogit(formula = Fail ~ scale(energy) + strata(Set), data = dietcc))

#data.table solution
summary(clogit(formula = event ~ scale(energy) + strata(set), data = dtcc))

内存优化方案：分组内直接抽样，避免全量非等值连接

核心思路是跳过全量非等值连接步骤，直接在每个风险集（按exit+job分组）内筛选符合条件的对照并抽样，全程在原数据集上操作，大幅减少中间数据量。

优化后代码

library(Epi)
library(data.table)
library(survival)

# 数据准备：原地转换时间变量，避免拷贝
data(diet)
setDT(diet)
diet[, `:=`(
  entry = as.numeric(doe - dob),
  exit = as.numeric(dox - dob),
  is_case = chd == 1  # 标记病例
)]

# 生成风险集标识：同exit+job的病例归为同一set
diet[, set := rleid(exit, job), by = .(exit, job)]

# 按风险集分组，为每个病例抽样对照并合并结果
dtcc <- diet[, {
  # 当前风险集的退出时间（组内所有病例/对照的exit一致）
  current_exit = exit[1]
  # 当前组的病例ID列表
  case_ids = id[is_case]
  # 筛选当前组的合格对照：非病例，且进入队列时间<=当前退出时间，退出时间>当前退出时间
  eligible_controls = id[!is_case & entry <= current_exit & exit > current_exit]
  
  # 计算需要抽样的对照总数：每个病例10个，不超过可用对照数
  n_required = length(case_ids) * 10
  sampled_controls = sample(eligible_controls, min(n_required, length(eligible_controls)), replace = FALSE)
  
  # 构建病例行
  case_data = .SD[is_case, .(id, job, exitcase = exit, set, event = 1, energy)]
  # 构建对照行，复用当前风险集的exit和set
  control_data = .SD[id %in% sampled_controls, .(id, job, exitcase = current_exit, set, event = 0, energy)]
  
  # 合并病例与对照
  rbind(case_data, control_data)
}, by = .(set, job, exit)]  # 按风险集维度分组

# 模型验证（与原结果一致）
summary(clogit(event ~ scale(energy) + strata(set), data = dtcc))

关键优化点

避免全量连接：直接在分组内筛选合格对照并抽样，不会生成中间超大数据集，内存占用仅为最终结果的量级。
原地操作原数据集：无需拆分病例/对照子集，所有计算在diet上完成，减少内存拷贝。
减少中间变量：分组计算中直接生成最终所需的病例-对照数据，不需要单独存储case、control等临时对象。
逻辑对齐原ccwc：风险集定义、抽样规则与Epi包ccwc()完全一致，模型结果可匹配。

进阶内存优化建议

若数据集极大，加载时用fread()指定列类型（如colClasses），避免自动转换带来的内存浪费。
将set、job、exit转换为整数类型（如果适用），提升data.table分组操作的效率。
若允许重复抽样（巢式病例对照一般不建议），可设置replace=TRUE，极端情况下进一步减少内存波动。

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