将数据从 MongoDB 并行加载到 python

Parallelizing loading data from MongoDB into python

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

我想通过并行化来加快加载过程。我的计划是生成多个进程或线程。每个进程将加载一个集合的一个块,然后将这些块合并在一起。

如何正确使用 MongoDB?

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

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

此外,pymongo returns 游标对象,它支持索引操作并具有一些方法,对我的任务似乎很有用,例如 countlimit

MongoSpark 的数据库连接器以某种方式完成了这个任务。不幸的是,我不熟悉 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 docspymongo.Cursorskip 方法也从集合的开头开始,如 PostgreSQL。同一页面建议在应用程序中使用分页逻辑。到目前为止,我找到的解决方案为此使用 sorted _id 。但是,它们还存储最后一次看到的 _id,这意味着它们也在单个线程中工作。

更新 2。我在官方 MongoDb Spark 连接器中找到了分区程序的代码: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 ) 失败并出现异常,如上所述。

"Read the mans, thery're rulez":)

pymongo.Collection 有方法 parallel_scan returns 游标列表。

更新。这个函数可以完成这项工作,如果集合不会经常改变,并且查询总是相同的(我的情况)。可以将查询结果存储在不同的集合中并 运行 并行扫描。

我正在研究 pymongo 并行化,这对我有用。我的简陋游戏笔记本电脑花了将近 100 分钟来处理我的 mongodb 的 4000 万份文档。 CPU 已 100% 使用,我不得不打开空调 :)

我使用skip 和limit 函数拆分数据库,然后将批次分配给进程。代码是为 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()

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

我认为 dask-mongo 会完成这里的工作。您可以使用 pip 或 conda 安装它,并且在 repo 中您可以在 notebook.

中找到一些示例

dask-mongo 将读取您在 MongoDB 中的数据作为 Dask 包,然后您可以使用 df = b.to_dataframe() 从 Dask 包转到 Dask Dataframe,其中 b 是您使用 dask_mongo.read_mongo

从 mongo 读取的包