[Перевод] 6 способов значительно ускорить pandas с помощью пары строк кода. Часть 1

Стас Молотов25.05.2020

0 2 Время чтения: 4 мин.

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

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

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

Часть 1:

Numba
Multiprocessing
Pandarallel

Часть 2:

Swifter
Modin
Dask

Numba

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

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

import numpy as npimport numba # создадим таблицу в 100 000 строк и 4 колонки, заполненную случайными числами от 0 до 100df = pd.DataFrame(np.random.randint(0,100,size=(100000, 4)),columns=['a', 'b', 'c', 'd']) # функция для создания новой колонкиdef multiply(x): return x * 5 # оптимизированная с помощью numba версия@numba.vectorizedef 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) # встроенная имплементация PandasIn [2]: %timeit df['new_col'] = df['a'] * 5545 µ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.njitdef 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 = 4pool = 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-компиляции и распараллеливание задачи. В следующей части я расскажу о более интересных и сложных инструментах, ну а пока я предлагаю вам потестировать инструменты самим, дабы убедиться в их эффективности.