Python 多进程计算示例

Python 多进程计算示例

import logging
from multiprocessing import Pool
import time
import pandas as pd

# 计算注册买单时间 (时间较长，推荐分割数据后用多进程进行计算)
bins = [0,1,15,30,45,60,75,90,105,120,135,150,165,180,360]
start = time.time()

def df_mp(df, row_start, row_end):
    """多进程函数"""
    # 分拆数据
    df_sub = df[row_start : row_end]
    df_sub = df_sub[df_sub['JoinTime'] < '2021-01-01']
    # 计算注册买单时间    
    df_sub['tranTime'] = df_sub.apply(lambda x: x['FirstPaymentTime'] - x['JoinTime'], axis=1)    
    del df_sub['FirstPaymentTime']
    df_sub['tranTime'] = df_sub['tranTime'].apply(lambda x: x.days)
    df_sub.fillna(181.0, inplace=True)
    df_sub['tranTime'] = df_sub['tranTime'] + 1
    df_sub['tranTime_bin'] = pd.cut(df_sub['tranTime'], bins)    
    return df_sub

if __name__ == '__main__':
    pool = Pool()
    batch_uids  = 500000  # 每个数据自己的数量
    batchs = int(len(allData2) / batch_uids) + 1
    res_l = []
    for i in range(batchs):
        m = i * batch_uids # 切片始
        n = (i + 1) * batch_uids # 切片终
        res = pool.apply_async(df_mp, args=(allData2, m, n,)) # 此处不能用get方法，会阻塞进程池
        res_l.append(res)
    print("==============================>")
    pool.close()
    pool.join()   #调用join之前，先调用close函数，否则会出错。执行完close后不会有新的进程加入到pool,join函数等待所有子进程>结束
    df = res_l[0].get()
    for res in res_l[1:]:
        df = df.append(res.get())

logger.info("all done. cost: %3f 秒" % (time.time() - start))