我正在尝试将大量数据集~146m行写入CSV。我试过这个:
def paramlist():
for row in nodes.itertuples():
l = []
for row2 in ref_stops.itertuples():
l.append((row[1], row[2], row[3], row2[1],
row2[2], row2[3], row2[4], haversine(row[3], row[2], row2[3], row2[2])))
yield l
pool = multiprocessing.Pool()
pool.map(func, paramlist())
def func(params):
with open(r'big_file.csv', 'a') as f:
writer = csv.writer(f)
for row in params:
writer.writerow(row)
此代码有效,但它会占用我的所有RAM和中止。
如何优化它?
答案 0 :(得分:3)
pool.map
将使用整个iterable。这就是为什么你会遇到内存问题。
您应该使用pool.imap
来避免这种情况。有关详细说明,请参阅this post。
话虽如此,我真诚地怀疑多处理会以您编写的方式加速您的程序,因为瓶颈是磁盘I / O.一遍又一遍地打开,追加和关闭文件几乎不比一次顺序写入快。并行写入单个文件是不可能的。
假设l
的生成需要一些时间,如果你编写这样的程序可能会加速:
from contextlib import closing
import multiprocessing
import csv
import pandas as pd
import numpy as np
# Just for testing
ref_stops = pd.DataFrame(np.arange(100).reshape((-1, 5)))
nodes = pd.DataFrame(np.arange(400).reshape((-1, 4)))
def haversine(a, b, c, d):
return a*b*c*d
# This function will be executed by the workers
def join_rows(row):
row_list = []
# join row with all rows from `ref_stops` and compute haversine
for row2 in ref_stops.itertuples():
row_list.append((row[1], row[2], row[3],
row2[1], row2[2], row2[3], row2[4],
haversine(row[3], row[2], row2[3], row2[2])))
return row_list
def main():
with closing(multiprocessing.Pool()) as pool:
# joined_rows will contain lists of joined rows in arbitrary order.
# use name=None so we get proper tuples, pandas named tuples cannot be pickled, see https://github.com/pandas-dev/pandas/issues/11791
joined_rows = pool.imap_unordered(join_rows, nodes.itertuples(name=None))
# open file and write out all rows from incoming lists of rows
with open(r'big_file.csv', 'w') as f:
writer = csv.writer(f)
for row_list in joined_rows:
writer.writerows(row_list)
if __name__ == '__main__':
main()
我假设你不关心订单,否则你不会首先选择多处理,对吧? 这样,它不是生成行列表但是工作进程的主要进程。只要一个工作进程完成了一个列表,它就会将它返回给主进程,然后主进程将其条目附加到文件中。然后,worker将获取一个新行并开始构建另一个列表。
一般来说,在程序中使用更多的pandas功能也可能更好(我假设你因为itertuples
而使用pandas数据帧)。例如,您可以创建一个新的Dataframe而不是行列表,并使haversine
与pandas.Series
对象兼容,这样您就不必在每个条目上调用它。
答案 1 :(得分:0)
尝试以块的形式写入数据。 读取您的数据框(假设您从数据框写入)部分,即某些块。 写一次每个块,这表现得相当快。