Запустите функцию с позиционными и необязательными аргументами параллельно в python

Question

Я пытаюсь вычислить различные метрики в Pandas DataFrame, используя метод apply. Поскольку DataFrame, с которым я работаю, довольно большой (1 миллион строк x 20 столбцов), я решил распараллелить процесс вычислений.

Чтобы воспроизвести проблему, с которой я столкнулся, я собираюсь использовать набор данных радужной оболочки. Вот шаги:

# Step 1: Import all required modules + load iris dataset to Pandas DataFrame

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

iris = pd.DataFrame(sns.load_dataset('iris'))

# Step 2: Define function that adds some metric to initial iris DataFrame

def add_metrics(data):
    data['x_1'] = data['species'].apply(lambda x: len(x))
    return data

# Step 3: Define parallelization function

num_partitions = 10 # number of partitions to split dataframe
num_cores = 4 

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

# Step 4: Add metrics to initial iris DataFrame using parallelization function

iris = parallelize_dataframe(iris, add_metrics)

Вышеописанный процесс работает отлично, так как НО Я хочу иметь возможность иметь дополнительные позиционные и / или необязательные аргументы в моей функции add_metrics , Например, моя функция add_metrics может выглядеть следующим образом:

def add_metrics(data, num, keep = False):
    data['x_1'] = data['species'].apply(lambda x: len(x))
    data['x_2'] = data['sepal_length'].apply(lambda x: x * num)
    if keep == True:
        data['x_3'] = data['sepal_width'].apply(lambda x: x * num)
    return data

Теперь, независимо от того, как я пытаюсь вызвать функцию parallelize_dataframe, я получаю сообщение об ошибке. Например:

iris = parallelize_dataframe(iris, add_metrics(iris, 2, keep = True)) бросает TypeError: 'DataFrame' object is not callable.

Я довольно новичок в Python, поэтому я не знаю, что здесь не так и как решить мою проблему. Я знаю, что выбранный мной пример не требует параллельной обработки, поскольку набор данных iris содержит только 150 наблюдений. Я использовал это, чтобы легко воспроизвести мою проблему.

Буду признателен за любую помощь.

Answer 1

Вы можете использовать functools.partial для установки переменных в вашей функции перед переходом на карту.

def add(x,y):
    return(x+y)

a = [1, 2, 3]
import functools
map(functools.partial(add, y=2), a) # map object [3, 4, 5]