使用PyMongo从MongoDB读取数据到Pandas时OOM问题求助

你遇到的问题其实很典型——MongoDB返回的BSON文档转成Python字典后，内存开销比CSV转Pandas的原生数据类型大得多，而且你的两种写法本质上都是一次性把所有数据加载到内存中，这才是内存爆掉的核心原因。下面给你拆解问题并给出具体的优化方案：

1. 根本原因：内存里的数据形态差异 #

CSV读取时，Pandas会直接将数据解析为高效的原生数据类型（比如int32、float64、category），而从MongoDB读取的每个文档是Python字典，每个字段都带有Python对象的额外开销（比如每个字符串是Python的str对象，比Pandas的object或category类型占更多内存）。更关键的是，list(cursor)会把所有90万条文档一次性加载到内存的列表里，这一步的内存占用会远超最终DataFrame的大小。

2. 优化方案：分批读取+直接生成DataFrame #

不要一次性把所有文档转成列表，而是用迭代器分批处理，直接生成小DataFrame再合并，这样能大幅降低峰值内存占用。

分批读取的单进程实现 #

import pandas as pd
from pymongo import MongoClient

client = MongoClient(host, port)
collection = client[db_name][collection_name]

# 每次读取10000条数据，可根据内存情况调整批次大小
batch_size = 10000
total_docs = collection.count_documents({})
dfs = []

for skip in range(0, total_docs, batch_size):
    cursor = collection.find().skip(skip).limit(batch_size)
    # 直接从cursor生成DataFrame，避免中转列表占用额外内存
    batch_df = pd.DataFrame.from_records(cursor)
    dfs.append(batch_df)
    print(f"已读取 {min(skip + batch_size, total_docs)}/{total_docs} 条记录")

# 合并所有批次
final_df = pd.concat(dfs, ignore_index=True)

3. 进一步优化：指定数据类型减少内存 #

提前给Pandas指定更高效的数据类型，比如把不需要精确计算的数值设为更小的类型，重复率高的字符串列设为category：

# 定义字段的目标类型，根据你的实际字段调整
dtype_spec = {
    "numeric_col1": "int32",
    "numeric_col2": "float32",
    "high_repeat_string_col": "category",
    # 如果不需要_id字段，后续可以直接删除或者在查询时排除
    "_id": "object"
}

# 在生成batch_df时指定dtype
batch_df = pd.DataFrame.from_records(cursor, dtype=dtype_spec)

另外，如果你不需要_id字段，直接在find()里排除，减少数据量：

cursor = collection.find({}, {"_id": 0}).skip(skip).limit(batch_size)

4. 改进多进程读取：避免全局列表收集数据 #

你的多进程写法最后还是把所有数据放到all_lists里再转DataFrame，这和单进程一次性加载没区别。应该让每个进程直接返回小DataFrame，然后合并：

import multiprocessing
import concurrent.futures
import pandas as pd
from pymongo import MongoClient

def read_mongo_batch(params):
    skip, batch_size, db_name, coll_name, host, port = params
    client = MongoClient(host, port)
    collection = client[db_name][coll_name]
    cursor = collection.find({}, {"_id": 0}).skip(skip).limit(batch_size)
    return pd.DataFrame.from_records(cursor)

if __name__ == "__main__":
    host = "your_host"
    port = 27017
    db_name = "your_db"
    coll_name = "your_coll"
    batch_size = 10000
    total_docs = MongoClient(host, port)[db_name][coll_name].count_documents({})
    
    # 生成所有批次的参数
    params_list = [
        (skip, batch_size, db_name, coll_name, host, port)
        for skip in range(0, total_docs, batch_size)
    ]
    
    dfs = []
    with concurrent.futures.ProcessPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
        for batch_df in executor.map(read_mongo_batch, params_list):
            dfs.append(batch_df)
            print(f"已完成一个批次，当前累计行数：{sum(len(df) for df in dfs)}")
    
    final_df = pd.concat(dfs, ignore_index=True)

5. 终极优化：MongoDB端预处理减少数据量 #

如果你的数据里有很多不需要的行或字段，先在MongoDB端用聚合过滤，减少返回的数据量：

# 示例：只取2023年及以后的数据，并且只返回需要的字段
pipeline = [
    {"$match": {"create_time": {"$gte": pd.Timestamp("2023-01-01")}}},
    {"$project": {"_id": 0, "col1": 1, "col2": 1, "col3": 1}}
]
cursor = collection.aggregate(pipeline)
# 然后再分批读取这个聚合结果

这样处理后，你应该能把内存控制在可接受的范围内，毕竟CSV只需要1.5GB，优化后的MongoDB读取应该能接近这个水平。

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