读取大型压缩文件

Reading large compressed files

这可能是一个简单的问题,但我似乎无法找到这个问题的答案或者为什么它在这个特定案例中不起作用。

我想读取大文件,可以压缩也可以不压缩。我用 contextlib 写了一个 contextmanager 函数来处理这个。然后使用 with 语句读取主脚本中的文件。

我的问题是脚本使用了大量内存然后被终止(使用压缩文件进行测试)。我究竟做错了什么?我应该采用不同的方法吗?

def process_vcf(location):
    logging.info('Processing vcf')
    logging.debug(location)
    with read_compressed_or_not(location) as vcf:
        for line in vcf.readlines():
            if line.startswith('#'):
                logging.debug(line)

@contextmanager
def read_compressed_or_not(location):
    if location.endswith('.gz'):
        try: 
            file = gzip.open(location)
            yield file
        finally:
            file.close()
    else:
        try: 
            file = open(location, 'r')
            yield file
        finally:
            file.close()

除了使用 for line in vcf.readlines(),您还可以:

line = vcf.readline()
while line:
    # Do stuff
    line = vcf.readline()

这一次只会将一行加载到内存中

影响最小的解决方案就是跳过 readlines 函数的使用。 readlines returns 包含文件中每一行的列表,因此它确实在内存中包含了整个文件。使用文件名本身使用生成器一次读取一行,因此它不必将整个文件都放在内存中。

    with read_compressed_or_not(location) as vcf:
        for line in vcf:
            if line.startswith('#'):
                logging.debug(line)

文件打开功能是读取gzip文件和非gzip文件的主要区别。因此可以动态分配开启器,然后读取文件。那么就不需要自定义上下文管理器了。

import gzip

open_fn = gzip.open if location.endswith(".gz") else open
with open_fn(location, mode="rt") as vcf:
    for line in vcf:
        ...