如何加速R中大样本暴露人群匹配函数？附现有低效代码

问题背景 #

现有两份数据：

• 10万行的暴露人群队列
• 300万行、覆盖5年周期的普通人群队列

需求是给每个暴露个体从普通人群中随机选5个匹配对象，满足：

• 出生年月完全一致
• 排除自身ID
• 排除处于7个不可用时间段内的个体（若时间段为空则不排除）
• 允许不同暴露个体重复选取同一普通人群个体（带替换）
• 结果存入新的data.table

此前尝试matchit包无法满足需求，自行编写的for循环耗时6-7天，需要优化速度，同时实现实时保存结果，避免中断后从头重跑。

核心优化策略 #

1. 预建索引，避免逐行全表扫描 #

普通人群队列有300万行，原代码每次循环都全表扫描筛选出生年月匹配的个体，这是最大的性能瓶颈。用data.table的setkey给birthyear和birthmonth建索引，后续查询同出生年月的个体时直接定位分组，速度提升几个数量级。

library(data.table)
# 给普通人群队列按出生年月设置主键索引
setkey(unexposed.cohort, birthyear, birthmonth)

2. 分组批量处理，减少重复操作 #

把暴露队列按birthyear+birthmonth分组，同组的暴露个体共享同一潜在匹配池，不用重复查询普通人群的同组数据。比如10万暴露个体如果分成1000个出生年月组，就只需要查询1000次潜在池，而非10万次。

3. 实时保存+断点续跑 #

放弃每次循环rbind到大data.table的方式，改为每处理一个个体就把结果追加写入文件，同时记录已处理的ID，下次启动时可以跳过已完成的个体，避免中断前功尽弃。

4. 简化不可用时间段判断逻辑 #

把重复的时间段判断逻辑整理成更简洁的写法，减少代码冗余，同时不影响逻辑正确性。

完整优化代码 #

find_matches_fast <- function(exposed.cohort, unexposed.cohort, save_path = "matches_progress.csv", batch_size = 1000, resume = FALSE) {
  # 断点续跑：读取已处理的匹配ID，过滤已完成的暴露个体
  processed_ids <- character(0)
  if (resume && file.exists(save_path)) {
    existing_data <- fread(save_path)
    processed_ids <- unique(existing_data$matchID)
    exposed.cohort <- exposed.cohort[!ID %in% processed_ids]
    cat("已跳过", length(processed_ids), "个已处理个体\n")
  }
  
  # 按出生年月分组处理暴露队列，减少重复查询潜在匹配池
  exposed_groups <- exposed.cohort[, .(ids = list(ID), exposure_dates = list(exposuredate)), 
                                  by = .(birthyear, birthmonth)]
  
  # 遍历每个出生年月组
  for (g_idx in seq_len(nrow(exposed_groups))) {
    current_group <- exposed_groups[g_idx]
    # 快速获取同出生年月的普通人群个体（利用索引，O(log n)速度）
    potential_pool <- unexposed.cohort[.(current_group$birthyear, current_group$birthmonth)]
    
    # 处理当前组内的每个暴露个体
    for (i in seq_along(current_group$ids[[1]])) {
      exp_id <- current_group$ids[[1]][i]
      exp_date <- current_group$exposure_dates[[1]][i]
      
      # 筛选符合条件的匹配对象
      valid_matches <- potential_pool[
        IDVariable != exp_id &
        # 7个不可用时间段判断，示例3个，其余4个直接复制扩展即可
        (!(exp_date >= unavailable_start1 & exp_date <= unavailable_end1) | (is.na(unavailable_start1) & is.na(unavailable_end1))) &
        (!(exp_date >= unavailable_start2 & exp_date <= unavailable_end2) | (is.na(unavailable_start2) & is.na(unavailable_end2))) &
        (!(exp_date >= unavailable_start3 & exp_date <= unavailable_end3) | (is.na(unavailable_start3) & is.na(unavailable_end3)))
        # 追加剩余4个时间段的判断逻辑
      ]
      
      # 随机采样5个（不足则取全部）
      selected_matches <- if (nrow(valid_matches) >= 5) {
        valid_matches[sample(.N, 5)]
      } else {
        valid_matches
      }
      
      # 添加匹配标记（关联到对应的暴露个体ID）
      selected_matches[, matchID := exp_id]
      
      # 实时追加写入文件，避免数据丢失
      fwrite(selected_matches, save_path, append = TRUE, col.names = !file.exists(save_path))
    }
    
    # 打印进度，方便跟踪
    cat("已完成出生年月", current_group$birthyear, "-", current_group$birthmonth, "组的匹配\n")
  }
  
  # 读取最终完整匹配结果
  final_matches <- fread(save_path)
  return(final_matches)
}

额外提速建议 #

• 并行处理：如果机器有多个CPU核心，可使用furrr包将分组处理部分改为并行执行，进一步压缩时间。
• 内存瘦身：只保留普通人群队列中匹配需要的列（比如ID、出生年月、不可用时间段列），减少内存占用，提升运算速度。
• 采样优化：如果某出生年月组的潜在匹配池极大，可提前对该组做一次随机排序，后续直接取前5个，避免每次重复采样（注意：因为允许重复选取，所以无需去重）。

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