使用PyMongo查询MongoDB数据写入S3的内存优化方案求助

我明白你现在的痛点——直接把MongoDB查询结果全拉到内存再转成JSON数组上传S3，内存占用直接飙到数据量的好几倍，这确实太头疼了。核心问题就是你一次性加载了所有文档，然后在内存里拼接完整的JSON数组，这必然会导致内存爆炸。咱们改成流式处理就能解决，边读MongoDB的文档，边往S3写数据，完全不用把所有数据存进内存。

1. 利用PyMongo游标懒加载特性 #

PyMongo的find()返回的是Cursor对象，它本身就是懒加载的——不会一次性把所有文档拉到本地内存，而是按需从MongoDB服务器获取数据。所以千万不要把游标转成列表（比如list(plans)），保持游标状态直接迭代就行。

2. 流式构建JSON数组并上传到S3 #

JSON数组的结构是[doc1, doc2, doc3,...]，我们可以分步骤流式输出：

• 先写入数组开头的[
• 迭代游标逐个处理文档：第一个文档直接序列化写入，后续文档先写逗号再序列化
• 最后写入数组结尾的]

下面是具体的代码实现，直接就能用：

import json
import boto3
from pymongo import MongoClient

def stream_mongo_to_s3(mongo_cursor, s3_bucket, s3_key):
    s3 = boto3.client('s3')
    
    # 自定义类文件对象，实现流式输出JSON数组
    class JSONStream:
        def __init__(self, cursor):
            self.cursor = cursor
            self.is_first_doc = True
        
        def read(self, size=-1):
            # 按需生成JSON片段，每次返回一段内容给S3
            if self.is_first_doc:
                self.is_first_doc = False
                try:
                    first_doc = next(self.cursor)
                    return f"[{json.dumps(first_doc)}"
                except StopIteration:
                    return "[]"
            else:
                buffer = []
                # 批量处理100个文档，提升效率（可根据需求调整数量）
                for _ in range(100):
                    try:
                        doc = next(self.cursor)
                        buffer.append(json.dumps(doc))
                    except StopIteration:
                        break
                if not buffer:
                    return "]"
                return "," + ",".join(buffer)
    
    # 流式上传到S3
    stream = JSONStream(mongo_cursor)
    s3.put_object(Bucket=s3_bucket, Key=s3_key, Body=stream)

# 你的业务代码部分
if __name__ == "__main__":
    # 初始化Mongo连接（根据你的实际配置调整）
    mongo_client = MongoClient("mongodb://your-host:27017/")
    mongo_db = mongo_client["your-db-name"]
    
    # 动态构建查询
    school_year_id = "2024"  # 示例值
    clientid = "client-123"  # 示例值
    mongo_query = {"$and": [{"ABC": school_year_id.upper()}, {"XYZ": clientid}]}
    
    # 获取游标（不要转成列表！）
    plans_cursor = mongo_db["your-collection"].find(mongo_query, no_cursor_timeout=True)
    
    try:
        # 流式上传到S3
        stream_mongo_to_s3(plans_cursor, "your-s3-bucket", "backup/plans.json")
    finally:
        # 手动关闭游标，避免MongoDB资源泄露
        plans_cursor.close()

3. 额外优化建议 #

• 禁用游标超时：如果查询时间很长，MongoDB的游标可能会超时，所以在find()时加上no_cursor_timeout=True，记得最后一定要手动关闭游标。
• 替换JSON序列化库：用ujson代替标准库的json，序列化速度更快，内存占用也更低，只需要把代码里的json.dumps换成ujson.dumps即可。
• 调整批量处理数量：代码里默认每次处理100个文档，你可以根据文档大小调整这个数字——文档越大，批量数可以越小，反之可以增大，平衡内存占用和IO效率。

为什么这个方案能降内存？ #

我们从来不会把所有MongoDB文档加载到内存里，每次只处理少量（1个或一批）文档，序列化后直接传给S3，S3会流式接收数据。内存里只会保留当前处理的几个文档和一小段JSON片段，内存占用会降到几MB级别，完全不用再担心1.5GB数据占12GB内存的问题。

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