Python超大字典写入JSON内存溢出问题排查及存储方案咨询

需要基于20万篇文章的数据集构建词维度的倒排字典，数据集共覆盖12.5万个唯一词，原始数据格式如下：

data = [[['article_1', 'city', 0.43], ['article_1', 'big', 0.38], ['article_1', 'beautiful', 0.25]], 
        [['article_2', 'sun', 0.65], ['article_2', 'beautiful', 0.41], ['article_2', 'shining', 0.21]],
        [['article_3', 'big', 0.72], ['article_3', 'beautiful', 0.50], ['article_3', 'butterfly', 0.25]]]

每个二级列表对应1篇独立文章，单条元素格式为[文章编号, 词, 词权重(TF-IDF)]，目标输出格式为词对应(文章编号, 权重)列表的字典，示例如下：

{'city': [('article_1', 0.43)],
 'big': [('article_1', 0.38), ('article_3', 0.72)],
 'beautiful': [('article_3', 0.50), ('article_2', 0.41), ('article_1', 0.25)],
 'sun': [('article_2', 0.65)],
 'shining': [('article_2', 0.21)],
 'butterfly': [('article_3',0.25)]}

注：原示例中big对应的元组多写了重复的词字段，属于笔误，和其他条目的二元组格式保持一致即可。

原有实现在Colab运行时占满内存导致环境崩溃，原代码如下：

with open('drive/MyDrive/dictionary.json', 'w') as f:
    d = defaultdict(list) 
    [d[i[1]].append((i[0],i[2])) for j in data for i in j]
    dct = dict(d.items())
    dct = dict(sorted(d.items()))
    f.write(json.dumps(dct) + '\n')

原有代码的内存浪费点非常集中：

• 滥用列表推导式做循环逻辑：列表推导式会生成和遍历长度一致的结果列表，这里所有append操作的返回值都是None，相当于凭空生成了一个和总词数等长、全是None的临时列表，按20万篇文章平均每篇100个词算，这个临时列表会多占近百MB内存，完全无意义。
• 重复拷贝全量字典：defaultdict本身就是字典子类，不需要额外做dict(d.items())转换；后续sorted(d.items())会先生成一个包含所有键值对的排序临时列表，再转成新字典，相当于同一份倒排数据在内存里同时存了3份，直接把内存占用拉到原来的3倍。
• 全量序列化占内存：json.dumps(dct)会一次性把整个字典转成完整的大字符串存入内存，再写入文件，12.5万词的倒排表序列化后体积很容易达到GB级，直接撑爆Colab的内存限额。

基础内存优化（不改变存储格式） #

把无意义的冗余操作全部删掉，改用普通for循环遍历，流式写入JSON避免全量序列化，参考代码：

import json
from collections import defaultdict

word_index = defaultdict(list)
# 逐篇加载处理，不要一次性把全量data读进内存
for article in data:
    for art_id, word, tfidf in article:
        word_index[word].append( (art_id, tfidf) )
    # 处理完一篇文章就释放这篇的原始数据内存
    del article

# 逐键流式写入JSON，避免生成全量序列化字符串
with open('drive/MyDrive/dictionary.json', 'w') as f:
    f.write('{\n')
    is_first = True
    for word in sorted(word_index.keys()):
        if not is_first:
            f.write(',\n')
        is_first = False
        f.write(f'  {json.dumps(word)}: {json.dumps(word_index[word])}')
        # 内存极度紧张时，写完一个键就释放对应内存
        # del word_index[word]
    f.write('\n}\n')

更适合大规模倒排的存储方案 #

如果数据量继续增长，JSON格式读写效率和内存占用都不适合做倒排索引存储，可以换成以下方案：

• 用SQLite数据库存储：建表字段为word（建索引）、article_id、weight，逐批插入数据，查询时直接按word筛选匹配的文章列表，不需要把全量索引加载到内存，Colab环境自带SQLite依赖，不需要额外安装。
• 用二进制序列化格式替换JSON：比如用pickle或者msgpack做序列化，相同数据的存储体积比JSON小40%以上，序列化/反序列化速度快2-3倍，内存开销更低。
• 如果后续需要做检索相关计算，可以直接用专门的倒排索引库，避免自己维护字典结构。

内容的提问来源于stack exchange，提问作者Kate