python使用协程解决内存瓶颈问题

前言

本来程序跑得好好的，突然快到结尾时程序被强制关闭，报错 zsh: killed  xxx/bin/python

后面经查找发现是因为内存耗尽，程序强制被关闭，追其溯源，发现使用了类似这样的代码：

从一个文件夹下读取全部文件数据作为list返回，由于list太大导致内存不足

能不能通过分块或者惰性加载的思想来迭代nc_datas从而解决内存瓶颈问题？

答案是可以的

协程模型

python中的协程采用了generator实现，而generator就类似惰性加载的思想，你调用它一次它才计算一次。比如 x * x 就是generator，它代表了一个理论上无限长的 list 集合，但实际上占用的内存只有当前的值，这也与Java中的Stream是类似的

由于前面内存耗尽、程序强制被关闭时，读数据占了大概10G的内存，因此我想分3～4次读取，这样每次只会消耗2～3G的内存，不会出现内存瓶颈问题

下面给出一个协程简化模型

# 判断当前下标c在第几个段
def check(c,red_lines):
    for i,red_v in enumerate(red_lines):
        if c < red_v:
            return i
    return len(red_lines)

# 比如3个段就会cut出2条分界线
def cut(x,chunck_count):
    per_num = len(x) / chunck_count
    red_lines = []
    # 比如分3段处理就有3-1=2条线(切2刀)
    for i in range(chunck_count - 1):
        red_lines.append((i + 1) * per_num)
    return red_lines

# 读数据
def read_nc_files():
    
    # 假设一共有10个文件要读取
    filepaths = ["1.nc","2.nc","3.nc","4.nc","5.nc","6.nc","7.nc","8.nc","9.nc","10.nc"]
    
    # 切成3次读取
    red_lines = cut(filepaths,3)
    ans = []
    old = 0
    for i,v in enumerate(filepaths):
        idx = check(i, red_lines)
        if idx != old:
            old = idx
            # 协程返回，下一步转到代码 "for data in datas:" 处进行下一次遍历
            yield ans
            del ans
            ans = []
        ans.append(v)

    yield ans


def test():
    
    # 由于python检测到read_nc_files函数里有yield关键字
    # 因此将其作为生成器generator返回，即这里datas不是list，而是一个generator
    datas = read_nc_files()

    # 对于遍历generator，遍历一次计算一次（返回一次）
    # 而不是一次性把所有datas返回
    # 遍历后的下一步将跳转到代码 "yield ans" 处并继续往下执行
    for data in datas:
        print(data)
        
test()

实验将10个文件切成3次读取，打印结果：

对于11、12、13个文件测试，同样打印结果：

可以看到采用协程的内存瓶颈解决方案成功