本文翻译自 Processing large JSON files in Python without running out of memory

在处理的大型 JSON 文件的时候很容易发生内存爆掉的问题。即便原始数据力量上能够为内存所容纳，但是由于 Python 在内存中的记录方式，其内存消耗会比原始数据的体积要更大。如果遇到了开启 SWAP 的计算机，即便内存不爆掉，如果程序的运行内存进入缓冲区，也会导致运行速度的急剧下降。解决这个问题的方法是流式解析 (Stream parsing)，也被称为 lazy parsing, iterative parsing，或者 chunked parsing。

1 问题：Python 的内存低效的 JSON 加载方式

考虑这样一个例子：一个大小是 24MB 的 JSON 文件，这个文件的内容代表了一系列的 Github 事件：

1
2
3
[{"id":"2489651045","type":"CreateEvent","actor":{"id":665991,"login":"petroav","gravatar_id":"","url":"https://api.github.com/users/petroav","avatar_url":"https://avatars.githubusercontent.com/u/665991?"},"repo":{"id":28688495,"name":"petroav/6.828","url":"https://api.github.com/repos/petroav/6.828"},"payload":{"ref":"master","ref_type":"branch","master_branch":"master","description":"Solution to homework and assignments from MIT's 6.828 (Operating Systems Engineering). Done in my spare time.","pusher_type":"user"},"public":true,"created_at":"2015-01-01T15:00:00Z"},
...
]

我们的目标是找到指定用交互过的仓库，下面这个简单的 Python 程序可以达成这一目标：

1
2
3
4
5
6
7
8
9
10
11
12
import json

with open("large-file.json", "r") as f:
    data = json.load(f)

user_to_repos = {}
for record in data:
    user = record["actor"]["login"]
    repo = record["repo"]["name"]
    if user not in user_to_repos:
        user_to_repos[user] = set()
    user_to_repos[user].add(repo)

程序运行结果是用户名和仓库名的映射字典。当我们使用 File memory profile 分析的时候，我们可以得到如下结果：

原文中这里是一个可交互的图表，建议在原文链接中查看

观察内存峰值，我们可以看到两处主要的内存分配行为：

  1. 读取文件；
  2. 将读取的内容转换成 Unicode 字符串。

我们来看 Python 的 json 模块的实现可以发现，这个标准库中的 json.load() 函数会先把整个文件读入内存。

1
2
3
4
5
6
7
def load(fp, *, cls=None, object_hook=None, parse_float=None,
        parse_int=None, parse_constant=None, object_pairs_hook=None, **kw):
    """Deserialize ``fp`` (a ``.read()``-supporting file-like object containing
    a JSON document) to a Python object.
    ...
    """
    return loads(fp.read(), ...)

注意上面记录到的是内存峰值的现象，所以后续创建字典对象时，其内存占用已经不是峰值处。整个程序执行过程中峰值是读取文件产生的。

有意思的是，尽管文件本身只有 24MB，但是读入内存之后其产生的内存峰值却远高于 24MB。为什么呢？

2 Python 的字符串内存表达方式

Python 的字符串表达经过优化可以使用较少的内存（这取决于字符串的内容）。首先，每个字符串都由于一个固有的开销 (overhead)。其次，如果字符串能够以 ASCII 编码表达，那么每个字符都只需要占用一个字节的内存。单如果有更多种类的字符需要表示，则每个字符占用的内存就上升到 4 个字节。我们来看下面的代码执行过程：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
>>> import sys
>>> s = "a" * 1000
>>> len(s)
1000
>>> sys.getsizeof(s)
1049

>>> s2 = "❄" + "a" * 999
>>> len(s2)
1000
>>> sys.getsizeof(s2)
2074

>>> s3 = "💵" + "a" * 999
>>> len(s3)
1000
>>> sys.getsizeof(s3)
4076

这三个 case 中每个字符串的长度都是 1000，但是他们使用的内存大小是不同的，这与其内容有关。

3 流式解决方案

显而易见将整个文件加载进入内存是一种内存的浪费。如果文件的体积非常大，那么我们甚至无法将整个文件读入内存。

如果 JSON 文件是一些对象组成的列表，那么理论上我们可以分片进行加载。个很多 Python 库支持这种行为，这里我们采用 ijson 库。

1
2
3
4
5
6
7
8
9
10
11
import ijson

user_to_repos = {}

with open("large-file.json", "rb") as f:
    for record in ijson.items(f, "item"):
        user = record["actor"]["login"]
        repo = record["repo"]["name"]
        if user not in user_to_repos:
            user_to_repos[user] = set()
        user_to_repos[user].add(repo)

在之前的标准库版本中，当数据被读入内存之后，文件就会被关闭。但是在现在这个场景下文件必须被保持打开。 这是因为文件的内容只有部分被读取，后续内容的读取取决后续的迭代过程。

items() 接口在此处的接受一个查询字符串用来指定加载的对象。在这个例子中 "item" 输入表示返回每个顶层对象。你可以参见 ijson 的文档来查询接口细节。

采用上面的代码我们可以发现程序的内存峰值降低到了 3.6MB。