Amazon DocumentDB聚合：$push+$slice提取各桶Top3数据

数据库文档 #

{uid: 1, score: 10}
{uid: 2, score: 11}
{uid: 3, score: 1}
{uid: 4, score: 6}
{uid: 5, score: 2}
{uid: 6, score: 3}
{uid: 7, score: 8}
{uid: 8, score: 10}

分桶规则 #

现有分桶聚合语句 #

已实现基础分桶需求的聚合代码：

db.scores.aggregate(
    [
        {
            $bucket: 
                {
                    groupBy: "$score",
                    boundaries: [0, 4, 7],
                    default: 7,
                    output:
                        {
                            "total": {$sum: 1},
                            "top_frustrated": 
                                {
                                    $push: {
                                        "uid": "$uid", "score": "$score"
                                    }
                                },
                        },
                }
        },
    ]
)

需求升级 #

希望每个桶仅返回Top3高分数据（例如7号桶仅返回uid2、1、8），同时保留该桶的总数total，输出示例如下：

{ "total" : 4, "top_scores" : 
    [
        {"uid" : 2, "score" : 11},
        {"uid" : 1, "score" : 10},
        {"uid" : 8, "score" : 10},
    ]
}

高效解决方案 #

针对Amazon DocumentDB（5.0及以上版本支持$sortArray），可以在现有$bucket阶段后添加两个聚合阶段，实现需求：

完整聚合语句 #

db.scores.aggregate([
    // 原分桶阶段
    {
        $bucket: {
            groupBy: "$score",
            boundaries: [0, 4, 7],
            default: 7,
            output: {
                "total": { $sum: 1 },
                "top_frustrated": {
                    $push: { "uid": "$uid", "score": "$score" }
                }
            }
        }
    },
    // 新增：对每个桶的数组按score降序排序，再取前3条
    {
        $addFields: {
            top_scores: {
                $slice: [
                    { $sortArray: { input: "$top_frustrated", sortBy: { score: -1 } } },
                    3
                ]
            }
        }
    },
    // 可选：移除原数组字段，只保留需要的输出字段
    {
        $project: {
            total: 1,
            top_scores: 1
        }
    }
])

方案说明 #

1. $sortArray：对每个桶内的top_frustrated数组按score降序排序，确保高分在前
2. $slice：截取排序后的数组前3条，得到Top3高分数据
3. 效率优势：仅在分桶后的小批量数组上做排序和切片操作，避免全局排序的性能损耗，比先$unwind再排序分组的方式更高效

兼容低版本DocumentDB（低于5.0）的备选方案 #

如果你的DocumentDB版本不支持$sortArray，可以用$unwind+$sort+$group的组合实现，但性能略低：

db.scores.aggregate([
    {
        $bucket: {
            groupBy: "$score",
            boundaries: [0, 4, 7],
            default: 7,
            output: {
                "total": { $sum: 1 },
                "docs": { $push: { "uid": "$uid", "score": "$score" } }
            }
        }
    },
    { $unwind: "$docs" },
    { $sort: { "_id": 1, "docs.score": -1 } },
    {
        $group: {
            _id: "$_id",
            total: { $first: "$total" },
            top_scores: { $push: "$docs" }
        }
    },
    {
        $project: {
            total: 1,
            top_scores: { $slice: ["$top_scores", 3] }
        }
    }
])

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