В этой статье я расскажу о шести инструментах, способных значительно ускорить ваш pandas код. Инструменты я собрал по одному принципу — простота интеграции в существующую кодовую базу. Для большинства инструментов вам достаточно установить модуль и добавить пару строк кода.

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

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

Часть 1:

• Numba
• Multiprocessing
• Pandarallel

Часть 2:

• Swifter
• Modin
• Dask

Numba

Этот инструмент ускоряет непосредственно сам Python. Numba — это JIT компилятор, который очень любит циклы, математические операции и работу с Numpy, поверх которого как раз и построен Pandas. Проверим на практике, какие преимущества он дает.

Смоделируем типичную ситуацию — нужно добавить новую колонку, применив какую-то функцию к уже существующей с помощью метода apply.

import numpy as np
import numba

# создадим таблицу в 100 000 строк и 4 колонки, заполненную случайными числами от 0 до 100
df = pd.DataFrame(np.random.randint(0,100,size=(100000, 4)),columns=['a', 'b', 'c', 'd'])

# функция для создания новой колонки
def multiply(x):
    return x * 5
    
# оптимизированная с помощью numba версия
@numba.vectorize
def multiply_numba(x):
    return x * 5

Как видите, вам не нужно ничего менять в своем коде. Достаточно всего лишь добавить декоратор. А теперь посмотрите на время работы.

# наша функция
In [1]: %timeit df['new_col'] = df['a'].apply(multiply)
23.9 ms ± 1.93 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# встроенная имплементация Pandas
In [2]: %timeit df['new_col'] = df['a'] * 5
545 µs ± 21.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# наша функция с numba
# мы отдаем весь вектор значений, чтобы numba сам провел оптимизацию цикла
In [3]: %timeit df['new_col'] = multiply_numba(df['a'].to_numpy())
329 µs ± 2.37 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Оптимизированная версия быстрее в ~70 раз! Впрочем, в абсолютных величинах, версия от Pandas отстала совсем несильно, поэтому возьмем более сложный кейс. Определим новые функции:

# возводим значения строки в квадрат и берем их среднее 
def square_mean(row):
    row = np.power(row, 2)
    return np.mean(row)
# применение:
# df['new_col'] = df.apply(square_mean, axis=1)

# numba не умеет работать с примитивами pandas (Dataframe, Series и тд.)
# поэтому мы даем ей двумерный массив numpy
@numba.njit
def square_mean_numba(arr):
    res = np.empty(arr.shape[0])
    arr = np.power(arr, 2)
    for i in range(arr.shape[0]):
        res[i] = np.mean(arr[i])
    return res
# применение:
# df['new_col'] = square_mean_numba(df.to_numpy())

Построим график зависимости времени вычислений от количества строк в датафрейме:

Итоги

• Возможно добиться ускорения в тысячи раз
• Иногда нужно переписать код, чтобы использовать Numba
• Можно использовать далеко не везде, в основном для оптимизации математических операций
• Стоит учесть, что Numba поддерживает не все возможности python и numpy

Multiprocessing

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

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

df = pd.read_csv('abcnews-date-text.csv', header=0)
# увеличиваю датасет в 10 раз, добавляя копии в конец
df = pd.concat([df] * 10)
df.head()

# считаем среднюю длину слова в заголовке
def mean_word_len(line):
    # этот цикл просто усложняет задачу
    for i in range(6):
        words = [len(i) for i in line.split()]
        res = sum(words) / len(words)
    return res

def compute_avg_word(df):
    return df['headline_text'].apply(mean_word_len)

Параллельность будет обеспечиваться следующим кодом:

from multiprocessing import Pool

# у меня 4 физических ядра
n_cores = 4
pool = Pool(n_cores)

def apply_parallel(df, func):
    # делим датафрейм на части
    df_split = np.array_split(df, n_cores)
    # считаем метрики для каждого и соединяем обратно
    df = pd.concat(pool.map(func, df_split))
    return df
# df['new_col'] = apply_parallel(df, compute_avg_word)

Сравним скорость работы:

Итоги

• Работает на стандартной библиотеке питона
• Мы получили ускорение в x2-3 раза
• Использовать распараллеливание на маленьких данных — плохая идея, т.к накладные расходы на межпроцессорное взаимодействие превышают выигрыш по времени

Pandarallel

Pandarallel — это небольшая библиотека для pandas, добавляющая в него возможность работать с несколькими ядрами. Под капотом работает на стандратном мультипроцессинге, так что увеличения скорости по сравнению с предыдущим подходом ждать не стоит, зато все из коробки + немного сахара в виде красивого progress bar ;) К слову, недавно видел неплохую статью на хабре по pandarallel.

Приступим к тестированию. Сейчас и далее я буду использовать те же данные и функции для их обработки, что и в предыдущей части. Сначала настроим pandarallel — это очень просто:

from pandarallel import pandarallel
# pandarallel сам определит сколько у вас ядер, но можно указать и самому
pandarallel.initialize()

Осталось только написать оптмизированную версию нашего обработчика, что тоже очень просто — достаточно заменить apply на parallel_aply:

df['headline_text'].parallel_apply(mean_word_len)

Сравним скорость работы:

Итоги

• В глаза сразу бросается довольно большой overhead в примерно 0.5 сек. При каждом применении parallel_apply под капотом сначала создается, а потом закрывается пул воркеров. В самописном варианте выше я создавал пул 1 раз, а потом его переиспользовал, поэтому накладные расходы были сильно ниже.
• Если не учитывать описанные выше издержки, то ускорение такое же, как и в предыдущем варианте, примерно в 2-3 раза.
• Pandarallel также умеет в parallel_apply над группированными данными (groupby), что довольно удобно. Полный список функционала и примеры смотрите здесь

В целом, я бы предпочел этот вариант самописному, т.к при средних/больших объемах данных разницы в скорости почти нет, при этом мы получаем крайне простой API и progress bar.

To be continued

В этой части мы рассмотрели 2 довольно простых подхода к оптимизации pandas — использование jit-компиляции и распараллеливание задачи. В следующей части я расскажу о более интересных и сложных инструментах, ну а пока я предлагаю вам потестировать инструменты самим, дабы убедиться в их эффективности.

> Часть 2

P.S.: Trust, but verify — весь код, использованный в статье (бенчмарки и отрисовка графиков), я выложил на github