如何合并两个pymongo查询的输出结果 避免使用高开销lookup操作

实现思路 #

• 先将第一个查询的结果转换为以员工姓名为键的字典映射，匹配时间复杂度为O(1)，内存开销极低
• 遍历第二个查询的游标，直接匹配对应员工的持股数量合并生成目标文档，无需一次性加载全量数据到内存

完整实现代码 #

from bson.son import SON
import pprint

# 第一步：执行查询1，构建员工姓名->持股数量的映射字典
pipeline = [
    {"$unwind":"$EmpName"},
    {"$group": {"_id": "$EmpName", "count": {"$sum": 1}}},
    {"$sort": SON([("count", -1), ("_id", -1)])}
]
emp_count_list = list(qatClient.get_database(dbName).get_collection(EmpCol).aggregate(pipeline))
# 构建映射字典，key为员工姓名，value为持股数
emp_count_map = {item["_id"]: item["count"] for item in emp_count_list}

# 第二步：执行查询2，遍历游标逐行合并数据
docCur = qatClient.get_database(dbName).get_collection(sapCol).find(
    {'currentTo' : currentToDate},
    {"_id" : 0,"EmpName":1, "Approvedby" : 1, "Expdate" : 1, "marked" : 1}
)

# 遍历合并，得到最终结果
final_result = []
for doc in docCur:
    emp_name = doc["EmpName"]
    # 匹配持股数量，不存在则默认填0，可根据需求调整为跳过该条记录
    count = emp_count_map.get(emp_name, 0)
    # 合并生成目标文档
    merged_doc = {
        "_id": emp_name,
        "count": count,
        **doc
    }
    final_result.append(merged_doc)
    # 可直接打印查看
    pprint.pprint(merged_doc)

可选优化说明 #

• 如果数据量过千万级别，可跳过把查询2结果存入final_result的步骤，直接逐行处理写入目标位置即可，避免占用过多内存
• 如果需要保留查询1中存在但查询2中不存在的员工记录，可额外遍历emp_count_map补全对应的记录，缺少的字段可根据需求填充默认值

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