从 JSON 加载 Pydantic 模型而不耗尽内存
Python⇒Speed ─ About ─ Consulting ─ Contact Articles ─ Docker packaging ─ Faster data science ─ Climate crisis Products ─ Docker packaging ─ Faster data science
从 JSON 加载 Pydantic 模型而不耗尽内存
作者：Itamar Turner-Trauring，最后更新于 2025 年 5 月 22 日，最初创建于 2025 年 5 月 22 日
你有一个很大的 JSON 文件，并且你想将数据加载到 Pydantic 中。不幸的是，这会占用大量内存，以至于大型 JSON 文件很难读取。该怎么办？
假设你必须使用 JSON，在本文中，我们将介绍：
- 使用 Pydantic 的默认 JSON 加载方式导致的高内存使用。
- 如何通过切换到另一个 JSON 库来减少内存使用。
- 更进一步，切换到带有 slots 的 dataclasses。
问题：20 倍的内存倍增
我们将从一个 100MB 的 JSON 文件开始，并将其加载到 Pydantic (v2.11.4) 中。我们的模型如下所示：
from pydantic import BaseModel, RootModel
class Name(BaseModel):
first: str | None
last: str | None
class Customer(BaseModel):
id: str
name: Name
notes: str
# Map id to corresponding Customer:
CustomerDirectory = RootModel[dict[str, Customer]]
我们加载的 JSON 大致如下所示：
{
"123": {
"id": "123",
"name": {
"first": "Itamar",
"last": "Turner-Trauring"
},
"notes": "关于 Itamar 的一些笔记"
},
# ... etc ...
}
Pydantic 内置了对加载 JSON 的支持，但遗憾的是，它不支持从文件读取。因此，我们将文件加载到字符串中，然后解析它：
with open("customers.json", "rb") as f:
raw_json = f.read()
directory = CustomerDirectory.model_validate_json(
raw_json
)
这非常简单直接。
但存在一个问题。如果我们测量峰值内存使用量，它会占用_大量_内存：
$/usr/bin/time -v python v1.py
...
Maximum resident set size (kbytes): 2071620
...
大约 2000MB 的内存，是 JSON 文件大小的 20 倍。如果我们的 JSON 文件是 10GB，那么内存使用量将是 200GB，我们可能会耗尽内存。我们能做得更好吗？
减少内存使用
解析 JSON 时，峰值内存使用有两个根本来源：
- 解析过程中使用的内存；许多 JSON 解析器不注意内存使用，并且使用了超出必要的内存。
- 最终表示形式（我们正在创建的对象）使用的内存。
我们将尝试减少每种情况下的内存使用。
1. 内存高效的 JSON 解析
我们将使用 ijson，一个增量 JSON 解析器，允许我们_流式传输_我们正在解析的 JSON 文档。我们不会将整个文档加载到内存中，而是每次加载一个键/值对。结果是，大部分内存使用将来自结果对象的内存表示形式，而不是解析：
import ijson
with open("customers.json", "rb") as f:
# 我们将自己创建根字典：
data = {}
# 空字符串是 ijson 查询语言的一部分，
# 在这种情况下，它表示“遍历顶层”，并且
# 由于我们正在使用 kvitems()，这意味着顶层
# 根 JSON 对象/字典中的键值对：
for cid, cust_dict in ijson.kvitems(f, ""):
# 为值字典创建一个 Customer：
customer = Customer.model_validate(cust_dict)
# 使用键将其存储在根字典中：
data[cid] = customer
# 现在创建根对象：
directory = CustomerDirectory.model_validate(data)
虽然以这种方式解析速度明显较慢（5 倍），但它可以显着减少内存使用，仅为 1200MB。
它还要求我们做更多解析 JSON 的工作，但是低于顶层 JSON 对象或列表的任何内容都可以由 Pydantic 完成。
2. 内存高效的表示
我们正在创建大量 Python 对象，并且节省 Python 对象内存的一种方法是使用“slots”。本质上，slots 是 Python 对象更有效的内存表示形式，其中可能的属性列表是固定的。这节省了内存，但代价是不允许向对象添加额外的属性，这在实践中并不常见，因此通常是一个很好的权衡。
不幸的是，pydantic.BaseModel 目前似乎不支持这一点，因此我切换到 Pydantic 的 dataclass 支持，它支持这一点。这是我们的新模型：
from pydantic import RootModel
from pydantic.dataclasses import dataclass
# 使用 slots 创建一个类；这意味着你不能
# 添加额外的属性，但它会使用更少的内存：
@dataclass(slots=True)
class Name:
first: str | None
last: str | None
@dataclass(slots=True)
class Customer:
id: str
name: Name
notes: str
# Map id to corresponding Customer:
CustomerDirectory = RootModel[dict[str, Customer]]
我们还需要稍微调整我们的解析代码：
import ijson
with open("customers.json", "rb") as f:
data = {}
for cust_id, cust_dict in ijson.kvitems(f, ""):
customer = Customer(**cust_dict)
data[cust_id] = customer
directory = CustomerDirectory.model_validate(data)
使用此版本的代码，内存使用量已降至 450MB。
最后的想法
以下是使用我们介绍的三种技术解析 100MB JSON 文件时的峰值内存使用情况摘要：
| 实现 | 峰值内存使用量 (MB) |
| ------------------------------------- | ------------------- |
| Model.model_validate_json() | 2000 |
| ijson | 1200 |
| ijson + @dataclass(slots=True) | 450 |
这种加载大量对象的特殊用例可能不是 Pydantic 开发者关心或有时间优先考虑的事情。但是 Pydantic 当然可以在内部更像 ijson 一样工作，并且可以添加使用 __slots__ 到 BaseModel 的选项。最终结果将使用更少的内存，同时仍然受益于 Pydantic 更快的 JSON 解析器。
在那之前，你可以自己实现这些选项。