如何从有序MongoDB集合中获取前N个类型的所有对应文档？

现有一个按date字段排序的MongoDB集合，单文档结构示例如下：

[{
    _id: new ObjectId(),
    type: 'foo',
    value: 123,
    date: '2022-06-30',
}, {
    _id: new ObjectId(),
    type: 'bar',
    value: 321,
    date: '2022-06-29',
}, {
    _id: new ObjectId(),
    type: 'foo',
    value: 456,
    date: '2022-06-28',
}, {
    _id: new ObjectId(),
    type: 'bar',
    value: 789,
    date: '2022-06-27',
}, {
    _id: new ObjectId(),
    type: 'baz',
    value: 234,
    date: '2022-06-26',
},
// 更多文档......
]

查询需要满足以下规则：

• 严格按照集合date排序顺序从头遍历文档
• 遍历过程中记录首次出现的type值，保留前2个首次出现的type对应的所有文档
• 遍历到第3个从未出现过的type时立刻终止
• 返回终止位置前，所有属于前2个首次出现type的文档

结果示例 #

初始集合场景 #

初始排序下最先出现的两个type为foo、bar，第三个首次出现的type为baz，预期返回：

// 预期返回结果
[{
    _id: new ObjectId(),
    type: 'foo',
    value: 123,
    date: '2022-06-30',
}, {
    _id: new ObjectId(),
    type: 'bar',
    value: 321,
    date: '2022-06-29',
}, {
    _id: new ObjectId(),
    type: 'foo',
    value: 456,
    date: '2022-06-28',
}, {
    _id: new ObjectId(),
    type: 'bar',
    value: 789,
    date: '2022-06-27',
}]

新增文档场景 #

插入如下date为2022-07-01的文档后：

{
    _id: new ObjectId(),
    type: 'baz',
    value: 567,
    date: '2022-07-01',
}

集合按date重排后，最先出现的两个type变为baz、foo，第三个首次出现的type是date为2022-06-29的bar，遍历到该文档即停止，预期返回：

[{
    _id: new ObjectId(),
    type: 'baz',
    value: 567,
    date: '2022-07-01',
}, {
    _id: new ObjectId(),
    type: 'foo',
    value: 123,
    date: '2022-06-30',
}]

注：该场景下2022-06-29的bar是有序遍历中第三个出现的新类型，其本身及后续所有文档均不纳入返回范围。

方案1：游标遍历（推荐，无版本限制、性能最优） #

通过排序游标逐条拉取文档，内存中记录已出现的type，遇到第三个新type直接关闭游标终止遍历，不需要扫描全表，适配所有MongoDB版本。
MongoDB Shell示例代码：

// 按业务排序规则创建游标，示例为date降序，升序将-1改为1即可
const cursor = db.collection.find().sort({ date: -1 })
const seenTypes = new Set()
const result = []

while (cursor.hasNext()) {
    const doc = cursor.next()
    if (!seenTypes.has(doc.type)) {
        // 已收集满2个type，当前为第三个新type，终止遍历
        if (seenTypes.size === 2) {
            cursor.close()
            break
        }
        seenTypes.add(doc.type)
    }
    // 前两个type对应的文档加入结果集
    result.push(doc)
}

// result即为符合要求的结果

方案2：聚合管道实现（数据库侧直接返回，需MongoDB 5.0+） #

基于窗口函数按排序顺序标记每个type的首次出现次序，截断第三个新type出现后的所有文档，适合需要在数据库层直接返回结果的场景。
聚合语句示例：

db.collection.aggregate([
    { $sort: { date: -1 } },
    // 标记每个type的首次出现次序
    {
        $setWindowFields: {
            sortBy: { date: -1 },
            output: {
                typeDenseRank: { $denseRank: { $partitionBy: "$type" } }
            }
        }
    },
    {
        $setWindowFields: {
            sortBy: { date: -1 },
            output: {
                typeFirstRank: { $min: "$typeDenseRank", $partitionBy: "$type" },
                docRank: { $documentNumber: {} }
            }
        }
    },
    // 找到第三个新type首次出现的位置
    {
        $setWindowFields: {
            sortBy: { date: -1 },
            output: {
                thirdTypeStart: { $min: { $cond: [{ $eq: ["$typeFirstRank", 3] }, "$docRank", null] } }
            }
        }
    },
    // 过滤符合要求的文档
    {
        $match: {
            typeFirstRank: { $lte: 2 },
            $or: [ { thirdTypeStart: null }, { docRank: { $lt: "$thirdTypeStart" } } ]
        }
    },
    // 清理计算产生的临时字段
    { $unset: ["typeDenseRank", "typeFirstRank", "docRank", "thirdTypeStart"] }
])

注意：如果使用的MongoDB版本低于5.0，不支持窗口函数，直接选择方案1即可，性能比低版本聚合拼接方案高很多。

内容的提问来源于stack exchange，提问作者Just a guy