具有大型数据集的Python和psycopg2多处理性能

时间:2013-12-19 18:01:16

标签: python memory psycopg2

我有一个非常标准的多处理脚本,可以处理数据库表中的200万条记录。在我完成将工作放入worker_queue之前,内存使用量超过12GB且崩溃。有没有更好的方法来设计它?

import math
import psycopg2
from psycopg2.extras import DictCursor
from multiprocessing import Process, Manager

from config import DB

connection = psycopg2.connect(DB)
cursor = connection.cursor(cursor_factory=DictCursor)

def worker(worker_queue, progress):
    for row in iter(worker_queue.get, None):
        # Do work
        progress.put(1)


if __name__ == "__main__":
    total, done = 0, 0

    cursor.execute("SELECT * from table")

    manager = Manager()
    worker_queue = manager.Queue()
    progress = manager.Queue()

    for row in cursor:
        worker_queue.put(row)
        total += 1

    workers = [Process(target=worker, args=(worker_queue, progress)) for i in range(50)]

    for each in workers:
        each.start()

    for i in iter(progress.get, None):
        done += 1

        remaining = total - done

        if remaining == 0:
            print 'Done'
        elif ((remaining % (10 ** int(math.log10(remaining)))) == 0):
            print str(remaining) + ' remaining'

1 个答案:

答案 0 :(得分:2)

值得注意的两件事

1)不要使用select *。这有两个原因:首先,您可能需要加载更多数据。其次是你无法控制数据的顺序(一旦我们进入第2点就很重要))

2)不要使用DictCursor。它将每一行变成一个dict,它吃掉 lot 的内存(因为你有效地复制了每个dict中的列名)。请改用默认的cursor_factory。现在,为了了解元组中返回的字段的顺序,您必须在select查询中指定该顺序。

那应该照顾好你的问题。如果没有,那么你必须在较小的数据集上完成工作。