如何为MongoDB JSON提取的嵌套数组添加post_id关联字段

现有Python代码可从本地MongoDB读取JSON文档，提取嵌套数组（tags、comments）并替换为ID，最终生成3份CSV文件，但缺少tags和comments与所属帖子的关联（即post_id列），不确定是在process_json处理阶段还是写入CSV阶段添加post_id，需要实现方案。

解决方案：在数据提取阶段（process_json）添加post_id #

推荐在提取嵌套数组时就为每个元素绑定对应的post_id，这样数据结构从一开始就保留关联关系，后续写入CSV时直接使用即可，逻辑更清晰且不易出错。

关键修改点： #

1. 修改process_json函数：处理tags数组时，将每个标签转为包含tag和post_id的字典；处理comments数组时，为每个评论字典新增post_id字段。
2. 调整CSV写入逻辑：适配修改后的数据结构，写入tags和comments时包含post_id列。

修改后的完整代码 #

import json
import csv
from pymongo import MongoClient
from bson import ObjectId, json_util


# Function to extract arrays and replace them with ID, add post_id to nested items
def process_json(json_data):
    result = {}
    arrays = {}

    # 获取并格式化当前帖子的ID作为post_id
    item_id = json_data["_id"]
    if isinstance(item_id, dict):
        item_id = item_id["$oid"]
    if isinstance(item_id, ObjectId):
        item_id = str(item_id)
    post_id = item_id

    result[post_id] = {}

    for key, value in json_data.items():
        if key != "_id":
            if isinstance(value, list):
                array_id = f"{key}_id"
                processed_items = []
                if key == "tags":
                    # 为每个标签添加post_id
                    processed_items = [{"tag": tag, "post_id": post_id} for tag in value]
                elif key == "comments":
                    # 为每个评论添加post_id
                    processed_items = [dict(comment, post_id=post_id) for comment in value]
                arrays[array_id] = processed_items
                result[post_id][array_id] = post_id
            else:
                result[post_id][key] = value

    return result, arrays

# Connecting to MongoDB
client = MongoClient("localhost", 27017)
db = client.mongotest
collection = db.collectiontest

# Selecting n JSON documents from collection
n = 2
documents = collection.find().limit(n)
json_documents = [json.loads(json_util.dumps(document, default=json_util.default)) for document in documents]

# Appending all JSONs into one
output = {}
extracted_arrays = {}

for doc in json_documents:
    processed_doc, arrays = process_json(doc)
    item_id = list(processed_doc.keys())[0]
    output[item_id] = processed_doc[item_id]

    for array_id, array_values in arrays.items():
        if array_id in extracted_arrays:
            extracted_arrays[array_id].extend(array_values)
        else:
            extracted_arrays[array_id] = array_values

# Creating main output CSV file
with open("output.csv", 'w', newline='') as csvfile:
    fieldnames_output = list(output[list(output.keys())[0]].keys())
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames_output)
    writer.writeheader()
    writer.writerows(output.values())

# Creating tags CSV file
with open("tags.csv", 'w', newline='') as csvfile:
    fieldnames_tags = ["tag", "post_id"]
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames_tags)
    writer.writeheader()
    writer.writerows(extracted_arrays['tags_id'])

# Creating comments CSV file
with open("comments.csv", 'w', newline='') as csvfile:
    # 取第一个评论的字段，加上post_id（确保字段完整）
    fieldnames_com = list(extracted_arrays['comments_id'][0].keys())
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames_com)
    writer.writeheader()
    writer.writerows(extracted_arrays['comments_id'])

# Printing output and extracted arrays for easier visual control
print("Extracted arrays:")
for array_id, array_values in extracted_arrays.items():
    print(array_id, ":", array_values)

print("\nOutput:")
print(json.dumps(output, indent=2))

改动说明： #

• process_json函数：针对tags数组，将每个字符串标签转为{"tag": "标签内容", "post_id": "帖子ID"}的字典；针对comments数组，用dict(comment, post_id=post_id)为每个评论字典新增post_id字段。
• tags.csv写入：改用DictWriter，指定tag和post_id字段，直接写入处理后的字典列表。
• comments.csv写入：因为处理后的评论已经包含post_id，直接用新的字段列表写入即可。

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