如何用data.table的fread按分组读取CSV文件每组的指定行数

完全可以实现，根据你的内存条件和文件规模，可选择以下两种方案：

方案1：单列预读取法（推荐，实现简单、效率高） #

适合内存可以容纳全量SOURCE列的场景，绝大多数情况下几千万行的单列字符数据内存占用不会超过1GB，完全可以轻松运行：

• 第一步：仅读取SOURCE标识列，获取全量行的分组信息

library(data.table)
# 仅读取SOURCE列，header根据你的文件实际情况调整
dt_source <- fread("your_file.csv", select = "SOURCE", header = TRUE)

• 第二步：生成每个分组最多保留20000行的行索引

# 若需要随机采样而非取前20000行，可将1:min(20000, .N)替换为sample(.N, min(20000, .N))
keep_idx <- dt_source[, .I[1:min(20000, .N)], by = SOURCE]$V1

• 第三步：按筛选后的行索引读取目标数据
  如果你的data.table版本 >= 1.14.0，可以直接用以下代码：

# 先读取全量再按索引过滤，适合整体文件内存占用不高的情况
dt_sample <- fread("your_file.csv", header = TRUE)[keep_idx, ]

如果文件全量读取内存不足，可以借助awk预处理过滤后再传入fread，不需要加载全量文件：

# 构造awk筛选命令，自动保留表头和目标行
awk_cmd <- sprintf("awk 'NR==1 || NR in {%s}' your_file.csv", paste(keep_idx + 1, collapse = ","))
dt_sample <- fread(cmd = awk_cmd, header = TRUE)

方案2：流式逐块读取法 #

适合文件规模极大、甚至单列SOURCE都无法全部加载进内存的极端场景：

library(data.table)
# 配置参数
target_per_group <- 20000
chunk_size <- 1000000 # 每次读取的块大小，可根据内存调整
file_path <- "your_file.csv"

# 初始化变量
collected_count <- data.table(SOURCE = character(), cnt = integer())
result_dt <- data.table()
# 读取表头
col_names <- names(fread(file_path, nrows = 0))
skip_row <- 1 # 跳过表头

while(TRUE) {
  # 读取当前块
  current_chunk <- fread(file_path, skip = skip_row, nrows = chunk_size, header = FALSE, col.names = col_names)
  if(nrow(current_chunk) == 0) break
  
  # 过滤已经采集满的分组
  full_group <- collected_count[cnt >= target_per_group, SOURCE]
  current_chunk <- current_chunk[!SOURCE %in% full_group]
  if(nrow(current_chunk) == 0) {
    skip_row <- skip_row + chunk_size
    next
  }
  
  # 每组仅保留还需要的行数
  current_keep <- current_chunk[, {
    needed <- target_per_group - collected_count[SOURCE == .BY$SOURCE, cnt]
    if(length(needed) == 0) needed <- target_per_group
    .SD[1:min(needed, .N)]
  }, by = SOURCE]
  
  # 合并结果、更新计数
  result_dt <- rbind(result_dt, current_keep)
  add_cnt <- current_keep[, .N, by = SOURCE]
  collected_count <- merge(collected_count, add_cnt, by = "SOURCE", all = TRUE)
  collected_count[is.na(cnt), cnt := 0][is.na(N), N := 0][, cnt := cnt + N][, N := NULL]
  
  # 所有分组都采集完成则退出循环，不需要读完全部文件
  if(all(collected_count$cnt >= target_per_group)) break
  skip_row <- skip_row + chunk_size
}

注意：如果CSV使用非逗号分隔符，需要在fread中指定sep参数对应实际分隔符，例如制表符设置为sep = "\t"。

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