5

在 python 操作大文件时节省内存

 8 months ago
source link: https://muyuuuu.github.io/2023/12/26/python-memory-optimization/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

在 python 操作大文件时节省内存

2023-12-26

35 2.8k 3 分钟

没想到有一天写 python 的时候也会想着如何去节省内存。平时写 python 的时候根本不会关注这些，变量什么的直接创建和使用就完了，也不用考虑内存的释放，反正有垃圾回收机制。只不过这次数据量过大，debug 的时候发现内存一直在申请，导致系统彻底的卡死。

可能也是从事算法的优化工作养成了职业病，每次写代码的时候都会想，这些代码消耗的时间怎么样，占用的空间怎么样，数据结构是否可以继续优化，这些逻辑有没有更优雅的写法。

注：本文程序中使用 psutil 库来监测进程使用的内存大小，需要 pip install psutil一下。

需要解析一个很大的日志文件，日志文件中含有一些无用的信息，像下面这样：

有用信息1
无用信息1
有用信息2
有用信息3
无用信息2
...
有用信息N

解析文件的时候，需要从文件中解析并提取出有用的信息，存入一个对象中，完成后续的处理。
但是呢，对于某些特殊的任务和需求，发现文件只解析一次是不行的，也就是需要对文件进行二次解析。

所以为了避免重复的解析文件，在第一次文件解析完毕后，直接把有用的核心信息序列化出去，这样二次解析的话就不用重新读取源文件在解析，直接读取序列化后的核心数据就好了。

序列化导出

最开始的方案是使用一个 list 持续追加解析得到的核心数据，文件解析完毕后把这个很大的 list 序列化出去。监测到进程占用的内存大小为：700MB。

import random
import pickle
import time
import psutil
import os

data = []

for i in range(10000000):
    data.append(str(random.randint(10000, 109070987)))

with open("data.pkl", "wb") as f:
    pickle.dump(data, f)

# 获取当前 Python 进程占用的内存
memory_info = process.memory_info()

# 打印占用的内存大小，rss 单位为字节
print(memory_info.rss / 1024 / 1024, "MB")

而如果使用序列化追加的方式，仅用 15MB，耗时增加 2s，毕竟每次序列化的时候都需要打开文件并在末尾追加内容：

with open("data.pkl", "ab") as f:
    for i in range(10000000):
        pickle.dump(str(random.randint(10000, 109070987)), f)

这里可以设置一个 buffer 进行优化，buffer 达到一定大小后在统一序列化出去。

class SeriesModel:
    def __init__(self) -> None:
        self._buf = []

def series(self, stack, finish=False):
        self._buf.append(stack)
        if 100 < len(self._buf) or finish is True:
            with open(config.SERIES_PATH, "ab") as f:
                for item in self._buf:
                    pickle.dump(item, f)
                self._buf = []

序列化读入

在二次解析的时候，需要把序列化的数据 load 进来。如果加载序列化的文件并且直接处理数据，同样需要使用 700MB 的内存。这种一次性创建所有元素的行为是没有必要的。

with open("data.pkl", "rb") as f:
    data = pickle.load(f)

for i in data:
        i += " "

可以使用惰性计算来解决这一问题，只有在真正需要这个变量的时候才去创建，而不是一开始就创建所有的变量。考虑到生成器表达式的局限性，我们直接使用 yield 关键字创建一个生成器函数。

yield 语句类似 return 会返回一个值，但它会记住这个返回的位置，下次迭代的时候就从这个位置继续执行，返回下一个元素。这样就消耗内存 15MB。

def read(file):
    with open(file, "rb") as f:
        data = pickle.load(f)
        for i in data:
            yield i

# data 是生成器
data = read("data.pkl")
for i in data:
    i += " "

任何一个生成器都会定义一个名为 __next__ 的方法，这个方法要在最后一个元素之后需抛出 StopIteration 异常。next() 函数的本质就是调用对象的 __next__()。这个方法要么返回迭代的下一项，要么引起结束迭代的异常 StopIteration，下面的示例揭示了生成器的本质。

class FibGenerator():
    def __init__(self, n):
        self.__n = n

self.__s0 = 0
        self.__s1 = 1
        self.__count = 0

def __next__(self):  # 用于内建函数 next()
        if self.__count < self.__n:
            ret = self.__s0
            self.__s0, self.__s1 = self.__s1, (self.__s0 + self.__s1)
            self.__count += 1
            return ret
        else:
            raise StopIteration

def __iter__(self):  # 用于 for 循环语句
        return self

fg = FibGenerator(5)
print(type(fg))
print(isinstance(fg, Iterable))
for i in fg:
    print(i, end=' ')

>>>
<class '__main__.FibGenerator'>
True
0 1 1 2 3

示例中如果没有定义 __iter__() 方法则只能使用 next() 函数进行迭代，当它定义后，就可以使用 for 和 in 语句访问了，同时定义了这两种方法的对象称为迭代器。生成器表达式和生成器函数产生生成器时，会自动生成名为 __iter__ 和 __next__ 的方法，所以生成器也是一种迭代器。

https://`python`howto.readthedocs.io/zh-cn/latest/iterator.html


report

Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK