MongoDB中SQL OVER PARTITION子句的等价实现方案问询

初始数据 #

原始工程师数据如下：

对应的MongoDB插入代码：

db.engineers.insertMany([
    { FirstName: 'John', LastName: 'Doe', BirthdateOn: ISODate('1990-01-01'), CupsOfCoffee: 5, HoursOfSleep: 8 },
    { FirstName: 'James', LastName: 'Bond', BirthdateOn: ISODate('1990-01-01'), CupsOfCoffee: 1, HoursOfSleep: 6 },
    { FirstName: 'Leeroy', LastName: 'Jenkins', BirthdateOn: ISODate('2000-06-20'), CupsOfCoffee: 16, HoursOfSleep: 10 },
    { FirstName: 'Jane', LastName: 'Doe', BirthdateOn: ISODate('2000-06-20'), CupsOfCoffee: 8, HoursOfSleep: 2 },
    { FirstName: 'Lorem', LastName: 'Ipsum', BirthdateOn: ISODate('2010-12-25'), CupsOfCoffee: 4, HoursOfSleep: 5 }
])

期望输出 #

需要得到包含以下信息的结果：

• 工程师个人信息（姓名、出生日期、咖啡杯数）
• 按出生日期分组后，咖啡杯数降序的行号
• 同出生日期的工程师总数
• 同出生日期的咖啡总杯数
• 同出生日期的平均睡眠时间

对应的SQL查询参考：

SELECT
    FirstName,
    LastName,
    BirthdateOn,
    CupsOfCoffee,
    ROW_NUMBER() OVER (PARTITION BY BirthdateOn ORDER BY CupsOfCoffee DESC) AS 'Row Number',
    COUNT(EngineerId) OVER (PARTITION BY BirthdateOn) AS TotalEngineers,
    SUM(CupsOfCoffee) OVER (PARTITION BY BirthdateOn) AS TotalCupsOfCoffee,
    AVG(HoursOfSleep) OVER (PARTITION BY BirthdateOn) AS AvgHoursOfSleep
FROM Engineers

期望的查询结果格式：

当前问题 #

现有的聚合查询仅能统计分组数据，但无法保留个人信息和行号：

db.engineers.aggregate([
    {
        $group: {
            _id: '$BirthdateOn',
            TotalEngineers: {
                $count: {  }
            },
            TotalCupsOfCoffee: {
                $sum: '$CupsOfCoffee'
            },
            AvgHoursOfSleep: {
                $avg: '$HoursOfSleep'
            }
        }
    }
])

希望不修改原集合，将个人数据与分组统计数据关联，同时生成行号。

解决方案 #

通过MongoDB聚合管道的分组、排序、添加行号、展开等步骤实现：

db.engineers.aggregate([
    // 按BirthdateOn分组，计算统计字段，同时保留分组内的所有文档
    {
        $group: {
            _id: "$BirthdateOn",
            TotalEngineers: { $count: {} },
            TotalCupsOfCoffee: { $sum: "$CupsOfCoffee" },
            AvgHoursOfSleep: { $avg: "$HoursOfSleep" },
            engineers: { $push: "$$ROOT" } // 将分组内的所有文档存入数组
        }
    },
    // 对每个分组内的工程师数组，按CupsOfCoffee降序排序
    {
        $set: {
            engineers: {
                $sortArray: {
                    input: "$engineers",
                    sortBy: { CupsOfCoffee: -1 }
                }
            }
        }
    },
    // 为排序后的数组添加行号
    {
        $set: {
            engineers: {
                $map: {
                    input: "$engineers",
                    as: "eng",
                    in: {
                        $mergeObjects: [
                            "$$eng",
                            { "Row Number": { $add: [{ $indexOfArray: ["$engineers", "$$eng"] }, 1] } }
                        ]
                    }
                }
            }
        }
    },
    // 展开工程师数组，将统计字段合并到每个文档中
    { $unwind: "$engineers" },
    // 投影出需要的字段，调整格式
    {
        $project: {
            _id: 0,
            FirstName: "$engineers.FirstName",
            LastName: "$engineers.LastName",
            BirthdateOn: "$_id",
            "Row Number": "$engineers.Row Number",
            CupsOfCoffee: "$engineers.CupsOfCoffee",
            TotalEngineers: 1,
            TotalCupsOfCoffee: 1,
            AvgHoursOfSleep: { $round: ["$AvgHoursOfSleep", 0] } // 取整，和SQL结果一致
        }
    },
    // 按BirthdateOn和行号排序，和示例结果顺序一致
    {
        $sort: {
            BirthdateOn: 1,
            "Row Number": 1
        }
    }
])

结果说明 #

执行上述聚合后，将得到与SQL查询几乎一致的结果：

• 每个工程师文档保留个人信息
• 生成按出生日期分组、咖啡杯数降序的行号
• 包含分组统计的TotalEngineers、TotalCupsOfCoffee、AvgHoursOfSleep字段
• 无需修改原集合，完全通过聚合管道计算得到

内容的提问来源于stack exchange，提问作者Andrew K.