如何高效从6万个kmer计数TXT文件构建目标Pandas DataFrame

提问场景 #

免责声明：这是首次发帖，若不符合社区规范还请见谅。

• 基础技术栈：Python3、Jupyter Notebooks、Pandas
• 前置工作：通过KMC kmer计数器完成60000个DNA序列的kmer计数，目标是将kmer计数作为Bag Of Words模型的特征，输入下游机器学习算法。

数据结构说明 #

每一个kmer计数结果存储为独立TXT文件，文件内以空格分隔两列内容：kmer字符串、对应计数值，示例结构：

AAAAAC     2
AAAAAG     6
AAAAAT     2
AAAACC     4
AAAACG     2
AAAACT     3
AAAAGA     5

全量数据集共包含6万个同类文件。

目标输出 #

生成单个宽表DataFrame，每行对应1个DNA序列的全量kmer计数，每列对应1种唯一kmer，单元格值为该kmer在对应序列中的计数值。

现有方案的性能问题 #

最初实现的循环逐文件读取、逐行追加的代码可正常运行，但性能极差：仅导入100个文件就耗时58秒，推算全量6万文件的处理时间将超过10小时，完全不具备实操性。原始实现代码如下：

import time
import pandas as pd

countsPath = r'D:\DataSet\MULTI\bow\6mer'
start = time.time()
df = pd.DataFrame()

for i in range(0, 60000):
    sample = pd.read_fwf(
        countsPath + r'\kmers-' + str(k) +'-seqNb-'+ str(i) + '.txt',
        sep=" ",
        header=None
    ).T
    new_header = sample.iloc[0]
    sample = sample[1:]
    sample.columns = new_header
    
    df = df.append(sample, ignore_index=True)

end = time.time()
print(f"Runtime of the program is {end - start} secs")
display(df)

此前尝试过使用Dask DataFrames、Dask Bag加速流程，但读取后的结果是所有文件的内容按顺序堆叠在单列中，无法实现“每个文件对应单独一行”的宽表结构。 #

高效实现方案 #

原有代码的核心性能瓶颈 #

• 逐次调用pd.append()是最大的性能损耗点：该方法每次执行都会全量复制已存在的DataFrame内容，6万次循环带来的内存复制量随循环次数线性增长，占整体耗时的90%以上
• 单文件解析的冗余开销：每次循环调用pd.read_fwf()的高层接口，附带大量不必要的元数据校验、类型自动推断等操作，解析效率远低于原生文本读取
• 无必要的格式转换：原始文件是两列结构，原代码通过转置调整结构的操作完全可以省略，额外增加了计算量

单线程优化实现（10倍以上提速） #

利用6-mer维度固定的特点（6-mer总共有4^6=4096种可能值，全数据集的列集合完全固定，不存在动态新增列的情况），跳过pandas的逐文件高层解析，直接使用原生Python读取文件、解析为字典，批量收集后一次性生成DataFrame：

import os
import time
import pandas as pd

countsPath = r'D:\DataSet\MULTI\bow\6mer'
k = 6
all_rows = []

start = time.time()
for i in range(60000):
    file_path = os.path.join(countsPath, f'kmers-{k}-seqNb-{i}.txt')
    row = {}
    with open(file_path, 'r') as f:
        for line in f:
            kmer, count = line.strip().split()
            row[kmer] = int(count)
    all_rows.append(row)

# 一次性把所有行转为DataFrame，缺失值填充为0
df = pd.DataFrame(all_rows).fillna(0)
end = time.time()
print(f"总耗时: {end - start:.2f} 秒")

该实现完全移除了循环内的DataFrame追加操作，最终的DataFrame构建是一次性内存分配，没有重复复制的开销。通常单线程运行该代码处理6万文件的总耗时可控制在1-2分钟以内。

并行加速实现（最高30倍提速） #

如果对速度要求更高，可使用多进程并行读取文件，充分利用磁盘IO和CPU空闲资源：

from concurrent.futures import ProcessPoolExecutor
import os
import pandas as pd

def parse_single_file(file_info):
    k, countsPath, i = file_info
    file_path = os.path.join(countsPath, f'kmers-{k}-seqNb-{i}.txt')
    row = {}
    with open(file_path, 'r') as f:
        for line in f:
            kmer, count = line.strip().split()
            row[kmer] = int(count)
    return row

if __name__ == '__main__':
    k = 6
    countsPath = r'D:\DataSet\MULTI\bow\6mer'
    tasks = [(k, countsPath, i) for i in range(60000)]
    with ProcessPoolExecutor(max_workers=8) as executor:
        all_rows = list(executor.map(parse_single_file, tasks))
    df = pd.DataFrame(all_rows).fillna(0)

在普通SSD环境下，该并行方案的处理耗时可进一步压缩到10-30秒。

Dask正确用法（适配超大规模数据集） #

如果需要处理内存无法容纳的更大规模数据集，Dask的正确使用方式是按文件分块读取，逐块解析为字典结构，而非直接将文本作为单列读取：

import dask.bag as db
import os

k = 6
countsPath = r'D:\DataSet\MULTI\bow\6mer'
file_list = [os.path.join(countsPath, f'kmers-{k}-seqNb-{i}.txt') for i in range(60000)]

def parse_file(path):
    row = {}
    with open(path, 'r') as f:
        for line in f:
            kmer, count = line.strip().split()
            row[kmer] = int(count)
    return row

# 按文件加载内容、逐文件解析为字典，直接生成对齐列名的宽表
b = db.from_sequence(file_list)
df = b.map(parse_file).to_dataframe().fillna(0).compute()

该方法不会出现“内容堆叠在单列”的问题，每一个文件的内容都会被解析为独立的字典结构，最终拼接为对齐所有kmer列名的标准宽表结构。 #

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