Параллелизм в pandas с помощью модуля multiprocessing

В pandas пока нет втроенной функции, которая позволяла бы производить вычисления, используя параллелизм, поэтому если требуется обработать большой объем данных, это может стать непростой задачей. Один из простейших способов ускорить обработку данных — разбить дейта-фрейм на отдельные части и затем, используя модуль multiprocessing, применить метод apply() и собрать дейта-фрейм обратно.

Как это можно сделать:

import pandas as pd
import numpy as np
import requests
import multiprocessing as mp


num_partitions = 10 # количество частей, на которые будет разбит дейта-фрейм
num_workers = 4 # количество процессов

# дейта-фрейм с тестовыми данными
df_source = pd.read_csv("https://ed-public-download.app.cloud.gov/downloads/Most-Recent-Cohorts-Scorecard-Elements.csv")

def parallelize_dataframe(df, func):
    df_split = np.array_split(df, num_partitions)
    pool = mp.Pool(num_workers)
    df = pd.concat(pool.map(func, df_split))
    pool.close()
    pool.join()
    return df

df и func — соответственно дейта-фрейм и функция, которая применяется к дейта-фрейму. В коде выше происходит разделение дейта-фрейма на набор фрагментов. Затем создается экземпляр Pool с 4 работниками (обычно количество соответствует числу ядер, доступных в машине). Функция pool.map() применяет func к каждому фрагменту дейта-фрейма в массиве df_split, а pd.concat() повторно объединяет все части данных в один дейта-фрейм.

Пробуем:

def combine_columns(df):
    df['INSTNM_CITY'] = df['INSTNM'] + ", " + df['CITY']
    return df

df_result = parallelize_dataframe(df_source, combine_columns)