MongoDB中按条件对相邻文档进行分组的实现方法

现有按order字段排序的MongoDB文档集合，需要按state字段对相邻的文档进行分组，输入输出示例如下：

输入示例 #

[
  {order: 1, state: 'one'},
  {order: 2, state: 'one'},
  {order: 3, state: 'one'},
  {order: 4, state: 'two'},
  {order: 5, state: 'two'},
  {order: 6, state: 'one'},
  {order: 7, state: 'two'},
  {order: 8, state: 'three'},
  {order: 9, state: 'three'}
]

期望输出 #

[
  [
    {order: 1, state: 'one'},
    {order: 2, state: 'one'},
    {order: 3, state: 'one'}
  ],
  [
    {order: 4, state: 'two'},
    {order: 5, state: 'two'}
  ],
  [
    {order: 6, state: 'one'}
  ],
  [
    {order: 7, state: 'two'}
  ],
  [
    {order: 8, state: 'three'},
    {order: 9, state: 'three'}
  ]
]

实现方案 #

核心思路是用窗口函数标记相邻同state的文档为同一分组，再通过分组聚合得到结果，完整聚合管道如下：

db.collection.aggregate([
  // 1. 确保文档按order升序排列（集合已有序可省略）
  {$sort: {order: 1}},
  // 2. 生成分组key：相邻state不同则key递增
  {$setWindowFields: {
    partitionBy: null,
    sortBy: {order: 1},
    output: {
      group_key: {
        $sum: {
          $cond: [
            {$ne: ["$state", {$lag: "$state", default: null}]},
            1,
            0
          ]
        },
        window: {documents: ["unbounded", "current"]}
      }
    }
  }},
  // 3. 按group_key分组，收集同组文档
  {$group: {
    _id: "$group_key",
    docs: {$push: "$$ROOT"}
  }},
  // 4. 清理文档，只保留需要的字段（可根据需求调整）
  {$project: {
    _id: 0,
    docs: {$map: {
      input: "$docs",
      as: "doc",
      in: {order: "$$doc.order", state: "$$doc.state"}
    }}
  }},
  // 5. 按分组顺序排序，保证输出顺序和原文档一致
  {$sort: {group_key: 1}},
  // 6. 将所有分组数组汇总成最终结果
  {$group: {
    _id: null,
    result: {$push: "$docs"}
  }},
  // 7. 只输出结果数组
  {$project: {_id: 0, result: 1}}
])

步骤说明 #

1. $sort：确保文档按order排序，这是相邻分组的前提。
2. $setWindowFields：通过$lag获取前一个文档的state，和当前文档对比，不同则给group_key加1，相同则加0；再用累加窗口计算每个文档的group_key，这样相邻同state的文档会拥有相同的group_key。
3. $group：按group_key将文档分组，用$push把同组文档收集到数组中。
4. $project：清理文档结构，移除不需要的group_key字段，只保留order和state（如果原始文档有其他字段，可修改$map中的内容保留）。
5. $sort：按group_key排序，保证分组顺序和原文档的顺序一致。
6. 最后一次$group：把所有分组的数组汇总成一个大数组，得到期望的输出格式。

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