如何在MongoDB聚合管道中按Python映射数值排序color字段？

需求说明 #

你有如下Python颜色权重映射：

color_map = {"Red": 1,
             "Yellow": 2,
             "Green": 3}

需要在MongoDB聚合管道中，依据该映射对集合中的color字段排序，集合示例文档：

{_id: 1,
 color: "Red",
 name: "Pony"},
{_id: 2,
 color: "Green",
 name: "Unicorn"},
{_id: 3,
 color: "Yellow",
 name: "Mole"}

期望排序后的聚合结果：

[{_id: 1, name: "Pony"},
 {_id: 3, name: "Mole"},
 {_id: 2, name: "Unicorn"}]

实现思路 #

核心通过聚合管道三步完成：

1. 添加权重字段：用$addFields结合$switch，根据color_map为每个文档添加对应排序权重值
2. 按权重排序：使用$sort基于新增的权重字段升序排列
3. 清理字段：用$project移除不需要的color和权重字段

MongoEngine代码实现 #

假设你的MongoEngine文档模型定义如下：

from mongoengine import Document, StringField

class Animal(Document):
    color = StringField(required=True)
    name = StringField(required=True)

利用Python的color_map动态生成聚合管道，代码如下：

color_map = {"Red": 1, "Yellow": 2, "Green": 3}

# 生成$switch分支规则，避免硬编码
switch_branches = [
    {"case": {"$eq": ["$color", color]}, "then": weight}
    for color, weight in color_map.items()
]

# 构建聚合管道
pipeline = [
    # 新增排序权重字段
    {
        "$addFields": {
            "sort_weight": {
                "$switch": {
                    "branches": switch_branches,
                    "default": 0  # 未在映射中的颜色默认排最后
                }
            }
        }
    },
    # 按权重升序排序
    {"$sort": {"sort_weight": 1}},
    # 移除不需要的字段，仅保留_id和name
    {"$project": {"color": 0, "sort_weight": 0}}
]

# 执行聚合查询并获取结果
sorted_results = list(Animal.objects.aggregate(*pipeline))
print(sorted_results)

补充说明 #

• 若color字段是MongoDB枚举类型或其他非字符串类型，只需调整$eq匹配的类型即可
• 动态生成switch_branches可避免硬编码，当color_map更新时无需手动修改聚合规则，维护更便捷

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