多gzip组学文件Dask全外连接内存溢出问题及优化咨询

问题描述 #

我正在处理1000+个总计约1GB的制表符分隔.gz格式组学数据集文件，每个文件含~300k行，结构为pos列及对应患者的count1、count2列，且各文件pos列无必然重叠。目标是将所有文件全外连接为统一数据表，示例结构如下：

我尝试过用Dask树归并方式实现全外连接、DuckDB的嵌套全外连接SQL语句、Polars结合functools.reduce，但即便在256GB内存环境下仍频繁出现内存溢出。现咨询：

1. 是否有更优实现方式？比如用Pandas逐次合并两文件并写入磁盘？
2. Dask代码内存溢出的原因是什么？如何在不增加RAM的情况下优化？
3. 还有哪些简化问题的思路？比如构建唯一pos列后如何关联数据？

附相关Dask核心代码：

if __name__ == "__main__":
    local_cluster = LocalCluster(
        dashboard_address=":8786",
        local_directory="/dev/shm/tmpdir",
        n_workers=2,
        threads_per_worker=6,
        memory_limit="70GB",
    )
    client = Client(local_cluster)
    # 读取文件逻辑...
    merged_df = tree_reduction(dfs)

tree_reduction函数：

def tree_reduction(dfs):
    if len(dfs) == 1:
        return dfs[0]
    elif len(dfs) == 2:
        return dd.merge(dfs[0], dfs[1], how="outer", left_index=True, right_index=True)
    else:
        midpoint = len(dfs) // 2
        left = tree_reduction(dfs[:midpoint])
        right = tree_reduction(dfs[midpoint:])
        return dd.merge(left, right, how="outer", left_index=True, right_index=True)

解决方案 #

问题1：是否有更优实现方式？比如用Pandas逐次合并两文件并写入磁盘？ #

• 直接用Pandas内存逐次合并不可行：随着合并次数增加，中间表的pos基数会持续膨胀，最终还是会触发内存溢出。
• 更优的磁盘分块合并方案：
  1. 遍历所有文件，提取并去重所有pos值，排序后分割为多个小分片（如每个分片含10w个pos），保存为分片索引文件。
  2. 针对每个分片索引，遍历所有原始文件，提取该分片内的pos数据，合并为对应分片的完整数据块后写入磁盘。
  3. 最后将所有分片数据块拼接为最终表。
• DuckDB优化写法：之前的嵌套全外连接效率极低，改用UNION ALL+分组聚合的方式，避免全外连接的内存开销：

  CREATE TABLE merged AS
  SELECT
    pos,
    MAX(CASE WHEN source = 'ABC' THEN count1 END) AS ABC_count1,
    MAX(CASE WHEN source = 'ABC' THEN count2 END) AS ABC_count2,
    MAX(CASE WHEN source = 'DEF' THEN count1 END) AS DEF_count1,
    MAX(CASE WHEN source = 'DEF' THEN count2 END) AS DEF_count2
    -- 其他患者列以此类推
  FROM (
    SELECT 'ABC' AS source, pos, count1, count2 FROM read_csv('ABC.gz', sep='\t')
    UNION ALL
    SELECT 'DEF' AS source, pos, count1, count2 FROM read_csv('DEF.gz', sep='\t')
    -- 其他文件依次UNION ALL
  )
  GROUP BY pos;
  

问题2：Dask代码内存溢出的原因是什么？如何在不增加RAM的情况下优化？ #

• 内存溢出原因：

  1. 树归并的中间合并步骤会生成包含全量pos的中间表，Dask默认会缓存这些中间结果，后期中间表的列数、行数快速膨胀，超出单worker内存限制。
  2. 单worker设置70GB内存上限，但树归并时左右分支的合并会让worker同时处理多个大分片，实际内存占用远超设定值。
  3. 基于索引的外连接需要对两边索引做全局排序或哈希分区，过程中产生大量临时数据，容易触发内存溢出。

• 优化方案：

  1. 禁用中间结果缓存：合并时添加persist=False，或调用client.cancel()清理中间结果，避免内存堆积。
  2. 调整分区策略：
     • 读取文件时手动设置更大的分区大小（如100MB/分区），减少分区数量；
     • 合并前对所有dd.DataFrame按pos做哈希分区（df = df.set_index('pos').repartition(npartitions=32)），让相同pos的数据落在同一分区，降低跨分区合并开销。
  3. 修改树归并逻辑：每合并2-3个文件就将结果持久化到磁盘（用to_parquet），读取磁盘文件继续合并，避免内存中堆积大中间表。
  4. 降低worker内存压力：减少n_workers（如设为1）或降低threads_per_worker，避免多线程竞争内存。

问题3：还有哪些简化问题的思路？比如构建唯一pos列后如何关联数据？ #

• 思路1：宽表转长表再聚合（最推荐）
  彻底避免全外连接，通过分组聚合实现宽表转换：

  1. 遍历所有文件，读取时为每个文件添加patient标识列（如文件名前缀），将宽表转为长表格式：pos, patient, count_type, value（count_type区分count1/count2）。
  2. 收集所有长表数据，按pos和patient分组，将count1/count2转回列，最终按pos聚合得到宽表。
     Polars实现代码示例：

  import polars as pl
  from pathlib import Path
  
  def process_file(file_path):
      patient = Path(file_path).stem.split('.')[0]
      return pl.read_csv(file_path, sep='\t') \
               .with_columns(pl.lit(patient).alias('patient')) \
               .melt(id_vars=['pos', 'patient'], value_vars=['count1', 'count2'], 
                     variable_name='count_type', value_name='value')
  
  all_files = list(Path('./data').glob('*.gz'))
  merged_long = pl.concat([process_file(f) for f in all_files], rechunk=True)
  merged_wide = merged_long.pivot(
      index='pos',
      columns=['patient', 'count_type'],
      values='value'
  ).fill_null(None)
  merged_wide.write_csv('merged_result.csv')
  

• 思路2：预构建全局pos索引

  1. 遍历所有文件，提取并去重所有pos值，排序后保存为全局索引文件。
  2. 对每个原始文件，与全局索引做左连接补充缺失pos的NA值，并重命名count列为患者专属名称。
  3. 最后将所有处理后的文件按pos列对齐，合并为宽表。
     该思路适合pos基数不大的场景，缺点是需要多一次遍历提取pos的IO操作，但避免了多次全外连接的内存开销。

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