引入

工作时,使用 mongo 查询全部 collections 的数据,在运行一段时间后,因为内存的问题导致程序被 kill,现使用迭代的查询方式解决

想法

使用 ObjectId 进行排序,从最小的 id 开始,到最大的 id 结束,并逐条返回内容,最终当完成后抛出 StopIteration 结束迭代

实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2017/7/21 下午3:45
# @Author  : Hou Rong
# @Site    : 
# @File    : MongoUtils.py
# @Software: PyCharm
each_line_num = 2000


def mongo_find_iter(collections):
    # 获取最小,最大 id 以便遍历全部
    min_id = list(collections.find().sort('_id', 1).limit(1))[0]['_id']
    max_id = list(collections.find().sort('_id', -1).limit(1))[0]['_id']

    now_id = min_id

    # 返回第一项
    __line = collections.find_one({'_id': {'$gte': now_id}})
    now_id = __line['_id']
    yield __line
    if now_id == max_id:
        raise StopIteration()

    # 返回剩余部分
    while True:
        for __line in collections.find({'_id': {'$gt': now_id}}).sort('_id').limit(each_line_num):
            now_id = __line['_id']
            yield __line
            if now_id == max_id:
                # 采用异常跳出两级循环
                raise StopIteration()


if __name__ == '__main__':
    import pymongo

    client = pymongo.MongoClient(host='your mongo host')

    src_collections = client['db name']['collection name']

    count = 0
    for line in mongo_find_iter(src_collections):
        count += 1
        if count % 10000 == 0:
            print(count)
    print(count)