使用 joblib 对 Pandas 数据进行并行处理

使用 joblib 对 Pandas 数据进行并行处理

如果需要对一个很大的数据集进行操作，而基于一列数据生成新的一列数据可能都需要耗费很长时间。

于是可以使用 joblib 进行并行处理。

假设我们有一个 dataframe 变量 data，要基于它的 source 列生成新的一列 double，其实就是把原来的 source 列做了个平方运算。感觉就这个简单的运算，应该有更简单的方法，在这里只是举个例子，我们使用 apply 方法并行实现。

如果直接使用 apply 那么直接如下实现

import pandas as pd

def double_func(data):
    return pow(data,2)

data["double"] = data["source"].apply(double_func)

使用并行实现如下

import pandas as pd
from joblib import Parallel, delayed

def double_func(data):
    return pow(data,2)

def key_func(subset):
    subset["double"] = subset["source"].apply(double_func)

data_grouped = data.groupby(data.index)
results = Parallel(n_jobs=8)(delayed(key_func)(group) for name, group in data_grouped)
data = pd.concat(results)

基本原理就是把整个 dataframe 根据 index，每行生成了一个子数据集，而把每个子数据集作为子任务使用多进程运行，最终生成 results 是多进程运行生成的结果的 list，使用 concat 重新组合就是我们最终想要的结果了。

n_jobs 参数就是需要使用几个进程池来运行程序。貌似一般 CPU 是几核的用几个进程会比较好？

其实速度并不是成倍减少的，具体原因我也……不太好讲清，但是还是可以很大幅度提升运行速度的。

---

顺便一提，如果数据集很大，程序一跑起来，根本不知道它跑得怎么样了，还是说卡死了。

注意到，我们生成的 data_grouped 是一个可迭代的对象，那么就可以使用 tqdm 来可视化进度条。

如果在 jupyter 里面使用的话，代码可以是下面这样

import pandas as pd
from joblib import Parallel, delayed
from tqdm import tqdm, tqdm_notebook

tqdm_notebook().pandas()

def double_func(data):
    return pow(data,2)

def key_func(subset):
    subset["double"] = subset["source"].apply(double_func)

data_grouped = data.groupby(data.index)
results = Parallel(n_jobs=8)(delayed(key_func)(group) for name, group in tqdm(data_grouped))
data = pd.concat(results)

友情提示，在我自己使用的时候遇到 bug ，提示无法从 Pandas 导入 PanelGroupby 的错误。查了许久才发现，是新版 Pandas 删除了PanelGroupby 这个模块。解决办法其实就是……升级 tqdm，在最新版已经修复了这个 bug 了。