pandas 多进程处理数据框

我假设看到这里的朋友都已经掌握了 apply。

核心函数:

def parallelize_dataframe(df, func, **kwargs):
    CPUs = multiprocessing.cpu_count()
    num_partitions = CPUs  # number of partitions to split dataframe
    num_cores = CPUs  # number of cores on your machine

    df_split = np.array_split(df, num_partitions)
    pool = Pool(num_cores)
    func = partial(func, **kwargs)
    df = pd.concat(pool.map(func, df_split))
    pool.close()
    pool.join()
    return df

一个具体的例子(仔细看注释):

import seaborn as sns
import numpy as np
import pandas as pd
from functools import partial
from multiprocessing import Pool
import multiprocessing

# 这个函数就是用来做多线程的，输入是你要处理的数据框和相应的函数
def parallelize_dataframe(df, func, **kwargs):
    CPUs = multiprocessing.cpu_count()
    num_partitions = CPUs # number of partitions to split dataframe
    num_cores = CPUs  # number of cores on your machine

    df_split = np.array_split(df, num_partitions)
    pool = Pool(num_cores)
    func = partial(func, **kwargs)
    df = pd.concat(pool.map(func, df_split))
    pool.close()
    pool.join()
    return df

# 导入 iris 数据集
iris = pd.DataFrame(sns.load_dataset('iris'))

# 首先你需要定义好你要对每一行处理的函数
def count(x):
    return len(x)

# 然后需要定义一个使用上述count函数的parall_func函数
def parall_func(df):
    df['length_of_word'] = df['species'].apply(count)
    return df

# 然后把上一步定义好的这个 parall_func 和 数据集，一起放入 parallelize_dataframe
iris = parallelize_dataframe(iris, parall_func)