python - 减少 for 循环和列表的内存占用

Question

我有一个正在尝试减少内存占用的函数。我可以使用的最大内存量只有 500MB。看起来使用 .split('\t') 和 for 循环确实使用了大量内存。有什么方法可以减少内存使用吗？

Line #    Mem usage  Increment   Line Contents
==============================================
10     35.4 MiB      0.0 MiB   @profile
11                             def function(username):
12     35.4 MiB      0.0 MiB       key = s3_bucket.get_key(username)
13     85.7 MiB     50.2 MiB       file_data = key.get_contents_as_string()
14    159.3 MiB     73.6 MiB       g = [x for x in file_data.splitlines() if not x.startswith('#')]
15    144.8 MiB    -14.5 MiB       del file_data
16    451.8 MiB    307.1 MiB       data = [x.split('\t') for x in g]
17    384.0 MiB    -67.8 MiB       del g
18
19    384.0 MiB      0.0 MiB       d = []
20    661.7 MiB    277.7 MiB       for row in data:
21    661.7 MiB      0.0 MiB           d.append({'key': row[0], 'value':row[3]})
22    583.7 MiB    -78.0 MiB       del data
25    700.8 MiB    117.1 MiB       database[username].insert_many(d)
26    700.8 MiB      0.0 MiB       return

更新1

根据 @Jean-FrançoisFabre 和 @Torxed 的建议，这是一个改进，但生成器似乎仍然占用大量内存。

@martineau 我更喜欢使用 MongoDB .insert_many() 因为迭代键并执行 .insert() 是很多速度较慢。

20     35.3 MiB      0.0 MiB   @profile
21                             def function(username):
22     85.4 MiB     50.1 MiB       file_data = s3_bucket.get_key(username).get_contents_as_string()
23    610.5 MiB    525.2 MiB       data = (x.split('\t') for x in isplitlines(file_data) if not x.startswith('#'))
24    610.5 MiB      0.0 MiB       d = ({'key': row[0], 'value':row[3]} for row in data)
25    123.3 MiB   -487.2 MiB       database[username].insert_many(d)
26    123.3 MiB      0.0 MiB       return

UDPATE2

我已经确定了内存使用的来源，如本配置文件所示:

21     41.6 MiB      0.0 MiB   @profile
22                             def insert_genotypes_into_mongodb(username):
23     91.1 MiB     49.4 MiB       file_data = s3_bucket.get_key(username).get_contents_as_string()
24     91.1 MiB      0.0 MiB       genotypes = (x for x in isplitlines(file_data) if not x.startswith('#'))
25     91.1 MiB      0.0 MiB       d = ({'rsID': row.split('\t')[0], 'genotype':row.split('\t')[3]} for row in genotypes)
26                                 # snps_database[username].insert_many(d)
27     91.1 MiB      0.0 MiB       return

insert_many() 函数清楚地解析了前面的行，导致整个列表加载到内存中并混淆了分析器。

解决方案是将键分块插入 MongoDB 中:

22     41.5 MiB      0.0 MiB   @profile
23                             def insert_genotypes_into_mongodb(username):
24     91.7 MiB     50.2 MiB       file_data = s3_bucket.get_key(username).get_contents_as_string()
25    180.2 MiB     88.6 MiB       genotypes = (x for x in isplitlines(file_data) if not x.startswith('#'))
26    180.2 MiB      0.0 MiB       d = ({'rsID': row.split('\t')[0], 'genotype':row.split('\t')[3]} for row in genotypes)
27     91.7 MiB    -88.6 MiB       chunk_step = 100000
28
29     91.7 MiB      0.0 MiB       has_keys = True
30    127.4 MiB     35.7 MiB       keys = list(itertools.islice(d,chunk_step))
31    152.5 MiB     25.1 MiB       while has_keys:
32    153.3 MiB      0.9 MiB           snps_database[username].insert_many(keys)
33    152.5 MiB     -0.9 MiB           keys = list(itertools.islice(d,chunk_step))
34    152.5 MiB      0.0 MiB           if len(keys) == 0:
35    104.9 MiB    -47.6 MiB               has_keys = False
36                                 # snps_database[username].insert_many(d[i*chunk_step:(i+1)*chunk_step])
37    104.9 MiB      0.0 MiB       return

感谢您的帮助。

Answer 1

首先，不要使用splitlines()，因为它会创建列表，您需要一个迭代器。所以你可以使用 Iterate over the lines of a string获取 splitlines() 迭代器版本的示例:

def isplitlines(foo):
    retval = ''
    for char in foo:
        retval += char if not char == '\n' else ''
        if char == '\n':
            yield retval
            retval = ''
    if retval:
        yield retval

个人注意:由于字符串连接，这不是很有效。我已经使用列表和 str.join 重写了它。我的版本:

def isplitlines(buffer):
    retval = []
    for char in buffer:
        if not char == '\n':
            retval.append(char)
        else:
            yield "".join(retval)
            retval = []
    if retval:
        yield "".join(retval)

然后，避免使用 del，因为不使用中间列表(分割行的列表除外)。只需“压缩”您的代码，跳过 g 部分并创建 d 作为生成器理解而不是列表理解:

def function(username):
   key = s3_bucket.get_key(username)
   file_data = key.get_contents_as_string()
   data = (x.split('\t') for x in isplitlines(file_data) if not x.startswith('#'))
   d = ({'key': row[0], 'value':row[3]} for row in data)
   database[username].insert_many(d)

这可能有点“单行”，但很难理解。并且当前的代码是可以的。将其视为链式生成器理解/表达式仅与一大块源内存一起工作:file_data