使用大数据集 Python 优化循环

Question

这是我第一次在 Python 上大展拳脚，所以我需要一些帮助。

我有一个具有以下结构的 mongodb（或 python dict）：

{
  "_id": { "$oid" : "521b1fabc36b440cbe3a6009" },
  "country": "Brazil",
  "id": "96371952",
  "latitude": -23.815124482000001649,
  "longitude": -45.532670811999999216,
  "name": "coffee",
  "users": [
    {
      "id": 277659258,
      "photos": [
        {
          "created_time": 1376857433,
          "photo_id": "525440696606428630_277659258",
        },
        {
          "created_time": 1377483144,
          "photo_id": "530689541585769912_10733844",
        }
      ],
      "username": "foo"
    },
    {
      "id": 232745390,
      "photos": [
        {
          "created_time": 1369422344,
          "photo_id": "463070647967686017_232745390",
        }
      ],
      "username": "bar"
    }
  ]
}

现在，我想创建两个文件，一个包含摘要，另一个包含每个连接的权重。我的适用于小型数据集的循环如下：

#a is the dataset
data = db.collection.find()
a =[i for i in data]

#here go the connections between the locations
edges = csv.writer(open("edges.csv", "wb"))
#and here the location data
nodes = csv.writer(open("nodes.csv", "wb"))

for i in a:

    #find the users that match
    for q in a:
        if i['_id'] <> q['_id'] and q.get('users') :
            weight = 0
            for user_i in i['users']:
                for user_q in q['users']:
                    if user_i['id'] == user_q['id']:
                        weight +=1
            if weight>0:
                edges.writerow([ i['id'], q['id'], weight])


    #find the number of photos
    photos_number =0
    for p in i['users']:
        photos_number += len(p['photos'])


    nodes.writerow([ i['id'],
                    i['name'],
                    i['latitude'],
                    i['longitude'],
                    len(i['users']),
                    photos_number
                ])

缩放问题：我有 20000 个位置，每个位置可能有多达 2000 个用户，每个用户可能有大约 10 张照片。

有没有更有效的方法来创建上述循环？也许是多线程、JIT、更多索引？ 因为如果我在单个线程中运行上述内容，最多可以得到 20000^2 *2000 *10 结果...

那么我怎样才能更有效地处理上述问题呢？ 谢谢

Answer 1

@YuchenXie 和@PaulMcGuire 建议的微优化可能不是您的主要问题，即您要循环超过 20,000 x 20,000 = 400,000,000 对条目，然后有一个包含 2,000 x 2,000 用户对的内部循环。这会很慢。

幸运的是，通过在 i['users'] 中预缓存用户 ID 的 sets 并用简单的集合交集替换您的内循环，可以使内循环变得更快。这会将 Python 解释器中发生的 O(num_users^2) 操作更改为 C 中发生的 O(num_users) 操作，这应该会有所帮助。 （我只是用大小为 2,000 的整数列表对其进行计时；在我的计算机上，它从你这样做的方式的 156 毫秒变为这种方式的 41 微秒，加速了 4,000 倍。）

您还可以通过注意关系是对称的，在成对的位置上切断主循环的一半工作，因此同时执行i = a[1]、q = a[2] 和i = a[2]、q = a[1] 是没有意义的。

考虑到这些和@PaulMcGuire 的建议，以及其他一些风格上的变化，你的代码变成了（警告：前面的代码未经测试）：

from itertools import combinations, izip

data = db.collection.find()
a = list(data)

user_ids = [{user['id'] for user in i['users']} if 'users' in i else set()
            for i in a]

with open("edges.csv", "wb") as f:
    edges = csv.writer(f)
    for (i, i_ids), (q, q_ids) in combinations(izip(a, user_ids), 2):
        weight = len(i_ids & q_ids)
        if weight > 0:
            edges.writerow([i['id'], q['id'], weight])
            edges.writerow([q['id'], i['id'], weight])

with open("nodes.csv", "wb") as f:
    nodes = csv.writer(f)
    for i in a:
        nodes.writerow([
            i['id'],
            i['name'],
            i['latitude'],
            i['longitude'],
            len(i['users']),
            sum(len(p['photos']) for p in i['users']), # total number of photos
        ])

希望这应该足以加快速度。如果没有，@YuchenXie 的建议可能会有所帮助，尽管我对此表示怀疑，因为 stdlib/OS 非常擅长缓冲这类事情。 （您可以使用文件对象的缓冲设置。）

否则，它可能归结为尝试从 Python 中获取核心循环（在 Cython 或手写 C 中），或者尝试使用 PyPy。不过，我怀疑这是否会给你带来任何巨大的加速。

您还可以将硬重量计算推送到 Mongo 中，这可能更聪明；我从来没有真正使用过它，所以我不知道。

Answer 2

瓶颈是磁盘 I/O。

当您合并结果并使用一个或多个writerows 调用而不是多个writerow 时，它应该会快得多。

Answer 3

是否折叠此循环：

photos_number =0
for p in i['users']:
    photos_number += len(p['photos'])

到：

photos_number = sum(len(p['photos']) for p in i['users'])

有帮助吗？

你的体重计算：

        weight = 0
        for user_i in i['users']:
            for user_q in q['users']:
                if user_i['id'] == user_q['id']:
                    weight +=1

也应该可以折叠成：

        weight = sum(user_i['id'] == user_q['id'] 
                        for user_i,user_q in product([i['users'],q['users']))

由于 True 等于 1，因此对所有布尔条件求和与计算所有为 True 的值相同。