python多进程读取大量小文件

背景

最近参加一个算法比赛，比赛的数据很大，解压后大约35G左右，由712839个小csv文件组成，每个文件大小在几kb大小不等，每个文件里包含一列、若干行数据，且每个文件的数据长度不一致，甚至会遇到空文件。不了解数据具体情况前，进行特征工程时候，至少需要遍历一遍所有文件，如果使用pandas的read_csv方法去循环遍历读取，速度慢的令人发指，比赛提交作品时甚至会超时（测试数据在线上，只能通过运行测试代码，线上处理数据后进行预测）。

数据分析

数据现况：

每个csv文件只有一列数据，数据长度不等；
文件名为样本名；
每个样本可能会有多个csv文件，每个csv文件都是该样本的特征；
数据除了列名外，只包含数字类型特征，存在空文件情况；
每个样本的csv特征文件存放在不同的文件夹路径中

目标数据格式：

一个文件包含一个样本的全部原始特征数据，文件名为样本名

解决思路

因为目标是将一个样本的所有特征数据存放到一个文件中，但每个特征数据长度不一致，而且并不存放在一起，所以不能直接一次读取一个样本的所有特征文件。并且最好也不要用numpy数组存放，因为数据长度不等，而numpy数组必须要补齐，个别文件长度异常的大，会导致最终的样本数据非常占空间。
所以想到可以用json保存数据，每个特征用字典存放，key是特征名称，value为具体的特征数据，一个json文件中包含一个字典，每个字典中用不同的key区分不同特征数据，读取速度也很快。
多进程方面，可以将所以的文件路径放到一个list中，利用多进程分别处理其中每个文件。

具体做法

获取所有文件的路径，去重；
按照文件名进行分组处理，一个样本的所有特征文件存放到一起；
open方法取代pandas的read_csv方法读取文件，按照样本先分组，每组内多个特征文件用多线程处理，同时每个样本组用多进程处理；

代码

单个样本完整路径为：
D:\train\OK\P1010\MeasuredData_I_LINE-VTS1_TCU_Raw\000dbac844493238aca63e33b28a0113.csv
其中 P1010@MeasuredData_I_LINE-VTS1_TCU_Raw 为该样本的特征

获取所有的数据路径

import json
import sys
import os
import re
import math
import datetime
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool

import numpy as np
import pandas as pd

def get_file(path):

    files_list = []

    for home, dirs, files in os.walk(path):
        for f in files:
            files_list.append(os.path.join(home, f))

    return list(set(files_list))

读取csv文件，第一行是特征名，这里直接从路径获取，所以文件中的可以扔掉。
如果是空文件，默认返回np.nan

def load_csv(path):

    col_name = "@".join([re.split(os.sep, path)[-3], re.split(os.sep, path)[-2]])

    with open(path, "r", encoding="UTF-8") as f:
        f.readline()  # 列名不需要
        data = f.readlines()

    data = [float(i.strip("\n")) for i in data]

    if len(data) > 0:
        return {
    
    col_name: data}
    else:
        return {
    
    col_name: [np.nan]}

按照样本分组，一个进程负责一个样本组内文件，每个样本组内多线程处理每个具体的文件

def map_process(file_list):

    save_path = r"C:\\csv2json"

    csv_name = file_list[0].split(os.sep)[-1]
    csv_df = {
    
    }

    pool = ThreadPool(6)
    df_list = pool.map(load_csv, file_list)
    pool.close()
    pool.join()

    for d in df_list:
        csv_df.update(d)

    if not os.path.exists(save_path):
        os.makedirs(save_path)

    with open(os.path.join(save_path, "{}".format(re.sub("csv", "json", csv_name))), "w") as f:
        json.dump(csv_df, f)


def main():

    path = 'C:\\train'

    files = get_file(path)

    data_info = pd.DataFrame({
    
    "file": files})
    data_info["sample"] = data_info["file"].apply(lambda x: x.split(os.sep)[-1])

    # 按照样本分组
    gr = data_info.groupby(by="sample")

    sample_list = []
    for csv, group in gr:
        sample_list.append(group["file"].values.tolist())

    # 多进程处理每个样本组
    t1 = datetime.datetime.now()
    print(t1)
    pool = Pool(6)
    pool.map(map_process, sample_list)
    pool.close()
    pool.join()
    print(datetime.datetime.now() - t1)


if __name__ == "__main__":
    main()

结果

不进行多线程处理时，一次运行时间大概需要几个小时，多线程处理后时间大概为25分钟，这还是在没有用全部核的情况下。多进程时不太建议使用全部核（pool = Pool(-1)），这样的话机器CPU资源会被全部占用，不能同时处理其他事情，出现卡顿。