В этой статье я расскажу о шести инструментах, способных значительно ускорить ваш 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()

publish_date headline_text
0 20030219 aba decides against community broadcasting lic...
1 20030219 act fire witnesses must be aware of defamation
2 20030219 a g calls for infrastructure protection summit
3 20030219 air nz staff in aust strike for pay rise
4 20030219 air nz strike to affect australian travellers

# считаем среднюю длину слова в заголовке
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


image


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