将数据从 MongoDB 并行加载到 python

Question

我在 MongoDB 中的集合中的所有文档都具有相同的字段。我的目标是将它们加载到 Python 中到 pandas.DataFrame 或 dask.DataFrame。

我想通过并行化来加快加载过程。我的计划是产生几个进程或线程。每个进程都会加载一个集合的一个块，然后这些块将被合并在一起。

如何正确使用 MongoDB？

我在 PostgreSQL 上尝试过类似的方法。我最初的想法是在 SQL 查询中使用SKIP 和LIMIT。它失败了，因为为每个特定查询打开的每个游标都从头开始读取数据表并且只是跳过了指定数量的行。所以我不得不创建额外的列，包含记录号，并在查询中指定这些数字的范围。

相反，MongoDB 为每个文档分配唯一的 ObjectID。但是，我发现不可能从另一个 ObjectID 中减去一个 ObjectID，它们只能与排序操作进行比较：小于、大于和等于。

另外，pymongo 返回光标对象，它支持索引操作并有一些方法，对我的任务似乎很有用，例如count、limit。

用于 Spark 的 MongoDB 连接器以某种方式完成了这项任务。不幸的是，我对 Scala 并不熟悉，因此，我很难知道他们是如何做到的。

那么，将数据从 Mongo 并行加载到 python 中的正确方法是什么？

到目前为止，我已经得出以下解决方案：

import pandas as pd
import dask.dataframe as dd
from dask.delayed import delayed

# import other modules.

collection = get_mongo_collection()
cursor = collection.find({ })

def process_document(in_doc):
    out_doc = # process doc keys and values
    return pd.DataFrame(out_doc)

df = dd.from_delayed( (delayed(process_document)(d) for d in cursor) )

但是，看起来dask.dataframe.from_delayed 在内部从传递的生成器创建了一个列表，有效地将所有集合加载到一个线程中。

更新。我发现in docs，pymongo.Cursor 的 skip 方法也从集合的开头开始，如 PostgreSQL。同一页面建议在应用程序中使用分页逻辑。到目前为止，我找到的解决方案为此使用 sorted _id。但是，它们还存储最后一次看到的_id，这意味着它们也可以在单个线程中工作。

更新2。我在官方的 MongoDb Spark 连接器中找到了 partitioner 的代码：https://github.com/mongodb/mongo-spark/blob/7c76ed1821f70ef2259f8822d812b9c53b6f2b98/src/main/scala/com/mongodb/spark/rdd/partitioner/MongoPaginationPartitioner.scala#L32

看起来，最初这个分区器从集合中的所有文档中读取关键字段并计算值的范围。

Update3：我的不完整解决方案。

不起作用，从 pymongo 获取异常，因为 dask 似乎错误地处理了 Collection 对象：

/home/user/.conda/envs/MBA/lib/python2.7/site-packages/dask/delayed.pyc in <genexpr>(***failed resolving arguments***)
     81         return expr, {}
     82     if isinstance(expr, (Iterator, list, tuple, set)):
---> 83         args, dasks = unzip((to_task_dask(e) for e in expr), 2)
     84         args = list(args)
     85         dsk = sharedict.merge(*dasks)

/home/user/.conda/envs/MBA/lib/python2.7/site-packages/pymongo/collection.pyc in __next__(self)
   2342 
   2343     def __next__(self):
-> 2344         raise TypeError("'Collection' object is not iterable")
   2345 
   2346     next = __next__

TypeError: 'Collection' object is not iterable

引发异常的原因：

def process_document(in_doc, other_arg):
    # custom processing of incoming records
    return out_doc

def compute_id_ranges(collection, query, partition_size=50):
    cur = collection.find(query, {'_id': 1}).sort('_id', pymongo.ASCENDING)
    id_ranges = [cur[0]['_id']]
    count = 1
    for r in cur: 
        count += 1
        if count > partition_size:
            id_ranges.append(r['_id'])
            count = 0
    id_ranges.append(r['_id'])
    return zip(id_ranges[:len(id_ranges)-1], id_ranges[1: ])    


def load_chunk(id_pair, collection, query={}, projection=None):
    q = query
    q.update( {"_id": {"$gte": id_pair[0], "$lt": id_pair[1]}} )
    cur = collection.find(q, projection)

    return pd.DataFrame([process_document(d, other_arg) for d in cur])


def parallel_load(*args, **kwargs):
    collection = kwargs['collection']
    query = kwargs.get('query', {})
    projection = kwargs.get('projection', None)

    id_ranges = compute_id_ranges(collection, query)

    dfs = [ delayed(load_chunk)(ir, collection, query, projection) for ir in id_ranges ]
    df = dd.from_delayed(dfs)
    return df

collection = connect_to_mongo_and_return_collection_object(credentials)

# df = parallel_load(collection=collection)

id_ranges = compute_id_ranges(collection)
dedf = delayed(load_chunk)(id_ranges[0], collection)

load_chunk 直接调用时完美运行。但是，调用 delayed(load_chunk)( blah-blah-blah ) 失败并出现异常，如上所述。

Answer 1

我正在研究 pymongo 并行化，这对我有用。我用了我不起眼的游戏笔记本电脑将近 100 分钟来处理我的 4000 万个文档的 mongodb。 CPU 使用率为 100%，我必须打开空调 :)

我使用跳过和限制函数来拆分数据库，然后将批次分配给进程。代码是为 Python 3 编写的：

import multiprocessing
from pymongo import MongoClient

def your_function(something):
    <...>
    return result

def process_cursor(skip_n,limit_n):
    print('Starting process',skip_n//limit_n,'...')
    collection = MongoClient().<db_name>.<collection_name>
    cursor = collection.find({}).skip(skip_n).limit(limit_n)
    for doc in cursor:        
        <do your magic> 
        # for example:
        result = your_function(doc['your_field'] # do some processing on each document
        # update that document by adding the result into a new field
        collection.update_one({'_id': doc['_id']}, {'$set': {'<new_field_eg>': result} })

    print('Completed process',skip_n//limit_n,'...')


if __name__ == '__main__':
    n_cores = 7                # number of splits (logical cores of the CPU-1)
    collection_size = 40126904 # your collection size
    batch_size = round(collection_size/n_cores+0.5)
    skips = range(0, n_cores*batch_size, batch_size)

    processes = [ multiprocessing.Process(target=process_cursor, args=(skip_n,batch_size)) for skip_n in skips]

    for process in processes:
        process.start()

    for process in processes:
        process.join()

最后一次拆分的限制将大于其余文档，但这不会引发错误

Answer 2

我认为dask-mongo 会在这里完成工作。您可以使用 pip 或 conda 安装它，在 repo 中您可以在 notebook 中找到一些示例。

dask-mongo 会将您在 MongoDB 中的数据作为 Dask 包读取，但随后您可以使用 df = b.to_dataframe() 从 Dask 包转到 Dask Dataframe，其中 b 是您使用 @987654326 从 mongo 读取的包@

Answer 3

“阅读手册，这是规则”:)

pymongo.Collection 具有返回游标列表的方法 parallel_scan。

更新。如果集合不经常更改并且查询始终相同（我的情况），则此功能可以完成工作。可以将查询结果存储在不同的集合中并运行并行扫描。