Python使用yield读取(csv,txt)大文件(1g/4g), 解决遇到的性能问题

最近线上项目发现上传2GB左右的csv文件遇到不知名问题,然后通过我自己本地虚拟机还原bug发现在我的虚拟机的项目服务直接挂了, 后来虚拟机扩大到16运存, 会报错说是csv文件 int64类型太大之类的, 换成int 32 就是内存不够读取csv文件里的数据, 导致内存溢出的问题

首先看下原始代码:

他是通过pandas中的read_csv()一次性把csv文件的数据读取到内存中,如果只有几M的话可能没问题,大型文件就会出现内存溢出的问题

pandas.read_csv(filename)

使用4gb的csv文件测试(2003列, 100万行)

解决方案1:

通过把iterator参数把大型csv文件一个TextFileReader 对象，以便逐块处理文件, 然后使用get_chunk(int)来分块读取, 再统计总行数  用时:2分48秒左右, 而且占用内存较大

table_single_col = pd.read_csv(file_path, iterator=True)
row_count = 0
loop = True
while loop:
   try:
      chunk = table_single_col.get_chunk(50000)
      row_count += chunk.shape[0]
    except StopIteration:
       break

解决方案2:

使用csv模块来处理大文件, 用时35秒左右, 而且占用内存较小

import csv
import datetime
# 打开CSV文件，并指定编码和读取方式
s = datetime.datetime.now()
with open('/tmp/data9096/test_data1.csv', 'r', encoding='utf-8') as csvfile:
    # 创建 CSV 文件读取器
    reader = csv.reader(csvfile)
    # 遍历每一行数据，以列表形式返回
    num = 0
    for row in reader:
        num += 1
    print("1", num, datetime.datetime.now() - s)

最终解决方案: 使用yeild生成器处理大文件(推荐使用)

用时4.5秒左右, 占用的内存可以忽略不计

而且还达到了解耦作用, read_file()负责与“数据生成”相关的逻辑。这样 file_count() 里面的主循环就只需要负责计数即可

def read_file(f):
    while True:
        block = f.readline()
        if not block:
            break
        yield block


def file_count(file_path):
    row_count = 0
    with open(file_path, "r") as f:
        for _ in read_file(f):
            row_count += 1
    return row_count

print((file_count('/tmp/data9096/test_data1.csv')))