如何在R中处理超内存Parquet数据集并完成分组汇总？

针对Parquet大数据集在R中处理时的内存溢出问题，结合你的业务需求，以下是几种高效的解决方案：

1. 基于Arrow的分阶段处理（优先推荐） #

先利用Arrow的磁盘计算能力完成全量分组聚合，再分批处理宽表转换，避免一次性加载全量数据到内存：

步骤1：全量聚合后写入分区Parquet #

先在Arrow层面完成分组汇总（无需加载数据到内存），并按gene1/gene2分区存储，大幅提升后续读取效率：

library(arrow)
library(dplyr)
library(tidyr)

# 打开原始Parquet数据集
raw_data <- open_dataset("path/to/parquet_files")

# 执行全量分组聚合，写入分区Parquet
agg_data <- raw_data %>%
  group_by(sample_id, gene1, gene2) %>%
  summarise(
    mean_importance = mean(importance),
    mean_count = mean(n_count),
    .groups = "drop"
  ) %>%
  filter(gene1 != gene2) %>%
  write_dataset(
    path = "aggregated_gene_data",
    format = "parquet",
    partitioning = c("gene1", "gene2")
  )

步骤2：分批处理宽表转换 #

基于分区后的聚合数据，每次仅加载单个gene1/gene2组合的数据到内存做pivot_wider：

# 打开聚合后的分区数据集
agg_dataset <- open_dataset("aggregated_gene_data")

# 获取所有唯一的基因组合（数据量小，直接加载到内存）
unique_gene_pairs <- agg_dataset %>%
  distinct(gene1, gene2) %>%
  collect()

# 循环处理每个基因组合
result_list <- lapply(1:nrow(unique_gene_pairs), function(i) {
  g1 <- unique_gene_pairs$gene1[i]
  g2 <- unique_gene_pairs$gene2[i]
  
  # 仅读取对应分区的数据，加载到内存后做宽表转换
  agg_dataset %>%
    filter(gene1 == g1, gene2 == g2) %>%
    collect() %>%
    pivot_wider(
      names_from = "sample_id",
      values_from = c("mean_importance", "mean_count"),
      names_sep = "__"
    )
})

# 合并所有结果
final_result <- bind_rows(result_list)

2. 用DuckDB直接SQL处理（零内存压力） #

DuckDB支持直接查询Parquet文件，通过SQL完成分组聚合和透视操作，全程在磁盘上计算，无需加载全量数据到R内存：

library(duckdb)
library(DBI)

# 建立DuckDB连接
con <- dbConnect(duckdb())

# 注册Parquet文件为数据库视图
dbExecute(con, "CREATE OR REPLACE VIEW gene_data AS SELECT * FROM parquet_scan('path/to/parquet_files/*.parquet')")

# 执行分组聚合+透视的SQL查询
final_result <- dbGetQuery(con, "
WITH aggregated AS (
  SELECT
    sample_id,
    gene1,
    gene2,
    AVG(importance) AS mean_importance,
    AVG(n_count) AS mean_count
  FROM gene_data
  WHERE gene1 != gene2
  GROUP BY sample_id, gene1, gene2
)
SELECT *
FROM aggregated
PIVOT (
  MAX(mean_importance) AS mean_importance,
  MAX(mean_count) AS mean_count
  FOR sample_id IN (SELECT DISTINCT sample_id FROM aggregated)
)
")

# 关闭连接并释放资源
dbDisconnect(con, shutdown = TRUE)

3. 用data.table优化内存效率 #

data.table的dcast（对应pivot_wider）内存效率远高于tidyr，可配合Arrow读取Parquet文件：

library(data.table)
library(arrow)

# 分批读取Parquet文件并合并为data.table（若内存仍不足，可分批次聚合）
file_list <- list.files("path/to/parquet_files", full.names = TRUE)
dt_combined <- rbindlist(lapply(file_list, function(f) as.data.table(read_parquet(f))))

# 分组聚合
agg_dt <- dt_combined[gene1 != gene2, .(
  mean_importance = mean(importance),
  mean_count = mean(n_count)
), by = .(sample_id, gene1, gene2)]

# 执行宽表转换
final_dt <- dcast(
  agg_dt,
  gene1 + gene2 ~ sample_id,
  value.var = c("mean_importance", "mean_count"),
  sep = "__"
)

4. 改进原有并行方案 #

若坚持使用mclapply，需修正代码错误并优化读取逻辑：

library(parallel)
library(arrow)
library(dplyr)
library(tidyr)

pq_dataset <- open_dataset("path/to/parquet_files")

# 获取唯一基因组合
grp_list <- pq_dataset %>%
  distinct(gene1, gene2) %>%
  filter(gene1 != gene2) %>%
  collect()

# 并行处理（修正filter逻辑，避免重复扫描）
res <- mclapply(1:nrow(grp_list), function(i) {
  g1 <- grp_list$gene1[i]
  g2 <- grp_list$gene2[i]
  
  pq_dataset %>%
    filter(gene1 == g1, gene2 == g2) %>%
    group_by(sample_id, gene1, gene2) %>%
    summarise(
      mean_importance = mean(importance),
      mean_count = mean(n_count),
      .groups = "drop"
    ) %>%
    collect() %>%
    pivot_wider(
      names_from = "sample_id",
      values_from = c("mean_importance", "mean_count"),
      names_sep = "__"
    )
}, mc.cores = 20) # 避免用满60核导致磁盘IO竞争

final_result <- bind_rows(res)

注：并行核数不宜过高，否则会引发磁盘IO瓶颈，建议根据磁盘IO能力调整（如10-20核）。

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