基于MongoDB双集合构建OrchestrationUser到UserGroup映射的高效方案问询

需求说明 #

使用Python的MongoEngine操作MongoDB，构建OrchestrationUser.user_id到UserGroup.id列表的映射：通过OrchestrationUser.group_ids匹配UserGroup.orchestrations_groups中的对应分组。

不可修改的MongoEngine模型 #

class OrchestrationUser(Document):
    user_id = StringField(required=True, primary_key=True)
    group_ids = ListField(StringField())

class OrchestrationGroup(EmbeddedDocument):
    group_id = StringField()
    group_name = StringField()

class OrchestrationGroups(EmbeddedDocument):
    orchestration_id = StringField()
    groups = ListField(EmbeddedDocumentField(OrchestrationGroup))

class UserGroup(Document):
    id = StringField(required=True, primary_key=True, default=lambda: str(uuid.uuid4()))
    orchestrations_groups = ListField(EmbeddedDocumentField(OrchestrationGroups))

约束条件 #

• OrchestrationUser实例最多160000个
• UserGroup实例最多250个
• 每个group_ids和orchestrations_groups.groups列表包含1-1000个ID
• 目标映射格式：{orch_user_id: [user_group_ids]}

当前实现方案 #

def map_users_to_usergroups():
    import datetime
    from itertools import chain
    from collections import defaultdict
    from ... import OrchestrationUser, UserGroup

    start_time = datetime.datetime.utcnow()
    group_to_usergroup_mapping = {}
    pipeline = [
        {'$unwind': '$orchestrations_groups'},
        {'$unwind': '$orchestrations_groups.groups'},
        {
            '$group': {
                '_id': '$orchestrations_groups.groups.group_id',
                'user_group_ids': {'$addToSet': '$_id'}
            }
        },
        {
            '$project': {
                '_id': 0,
                'group_id': '$_id',
                'user_group_ids': 1
            }
        }
    ]

    print("Starting aggregation of user groups...")
    for entry in UserGroup.objects.aggregate(*pipeline):
        group_to_usergroup_mapping[entry['group_id']] = entry['user_group_ids']
    print(f"Aggregation complete. Total {len(group_to_usergroup_mapping)}")

    user_to_usergroups_mapping = defaultdict(set)
    last_user_id = None

    print("Mapping users to user groups...")
    while True:
        query = OrchestrationUser.objects.order_by('user_id')
        if last_user_id is not None:
            query = query.filter(user_id__gt=last_user_id)
        chunk = query[:5000]

        if not chunk:
            break

        for orchestration_user in chunk:
            user_to_usergroups_mapping[orchestration_user.user_id].update(
                chain.from_iterable(
                    group_to_usergroup_mapping.get(group_id, [])
                    for group_id in orchestration_user.group_ids
                    if group_id in group_to_usergroup_mapping
                )
            )
            last_user_id = orchestration_user.user_id

        print(f"Processed chunk up to user_id: {last_user_id}")

    total_time = (datetime.datetime.utcnow() - start_time).total_seconds()
    print(f"Took {total_time} seconds.")
    
    return user_to_usergroups_mapping

user_to_usergroups_mapping = map_users_to_usergroups()

示例数据 #

OrchestrationUsers #

[
  {"user_id": "orch_user_1", "group_ids": ["group_1", "group_2"]},
  {"user_id": "orch_user_2", "group_ids": ["group_2", "group_3"]}
]

UserGroups #

[
  {
    "id": "user_group_1",
    "orchestrations_groups": [
      {
        "orchestration_id": "orch_1",
        "groups": [
          {"group_id": "group_1", "group_name": "Group One"},
          {"group_id": "group_2", "group_name": "Group Two"}
        ]
      }
    ]
  },
  {
    "id": "user_group_2",
    "orchestrations_groups": [
      {
        "orchestration_id": "orch_2",
        "groups": [
          {"group_id": "group_3", "group_name": "Group Three"}
        ]
      }
    ]
  }
]

预期结果 #

{
    "orch_user_1": ["user_group_1"],
    "orch_user_2": ["user_group_1", "user_group_2"]
}

高效优化方案 #

由于UserGroup实例仅250个，完全可以在内存中完成映射构建，避免MongoDB聚合的额外开销，同时优化用户数据的查询和处理逻辑：

def map_users_to_usergroups_optimized():
    import datetime
    from collections import defaultdict
    from ... import OrchestrationUser, UserGroup

    start_time = datetime.datetime.utcnow()
    group_to_usergroup = defaultdict(set)

    print("Loading and processing UserGroups...")
    # 仅加载必要字段，减少数据传输
    for user_group in UserGroup.objects.only('id', 'orchestrations_groups'):
        user_group_id = user_group.id
        # 遍历嵌套结构，构建group_id到user_group_id的映射
        for orch_group in user_group.orchestrations_groups:
            for group in orch_group.groups:
                group_to_usergroup[group.group_id].add(user_group_id)
    # 转换为列表格式，匹配预期输出
    group_to_usergroup = {k: list(v) for k, v in group_to_usergroup.items()}
    print(f"Processed {len(group_to_usergroup)} unique group IDs from {UserGroup.objects.count()} UserGroups")

    user_to_usergroups = defaultdict(set)
    batch_size = 10000  # 增大批次，减少查询次数

    print("Mapping OrchestrationUsers to UserGroups...")
    # 批量加载用户，仅获取必要字段
    for batch in OrchestrationUser.objects.only('user_id', 'group_ids').batch_size(batch_size):
        user_id = batch.user_id
        # 快速匹配关联的用户组
        for group_id in batch.group_ids:
            if usergroups := group_to_usergroup.get(group_id):
                user_to_usergroups[user_id].update(usergroups)
    
    # 转换为最终的列表格式
    user_to_usergroups = {k: list(v) for k, v in user_to_usergroups.items()}
    
    total_time = (datetime.datetime.utcnow() - start_time).total_seconds()
    print(f"Completed in {total_time:.2f} seconds.")
    
    return user_to_usergroups

user_to_usergroups_mapping = map_users_to_usergroups_optimized()

优化点说明 #

1. 内存映射构建：直接遍历少量的UserGroup实例，在内存中构建映射，避免MongoDB聚合的IO和计算开销。
2. 按需加载字段：使用only()方法只获取必要字段，大幅减少网络传输的数据量。
3. 增大批次大小：将用户查询批次提升至10000，减少查询MongoDB的次数，降低IO开销。
4. 简化匹配逻辑：去掉冗余的链式迭代操作，直接通过字典查找完成映射，逻辑更简洁高效。

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