Привет! Если после заголовка вы решили, что это очередная статья в стиле «Топ-10 способов ускорить Pandas», то не торопитесь с выводами. Вместо топов и подборок предлагаю взглянуть на бенчмарки скорости и потребления памяти в зависимости от характеристик датафрейма и убедиться, что часть советов из статей по ускорению могут оказаться даже вредными. Разберём, какой из способов ускорения нужно пробовать в разных ситуациях, как это зависит от размера датафрейма и как ведёт себя в реальном проекте.
Эта статья будет вам полезна, если вы новичок, который по каким-то причинам обрабатываете большие объёмы данных при помощи Pandas, и вам нужно его ускорить с минимальным количеством правок в коде.
Так вот, предыстория. Летом 2023 прямо перед уходом в отпуск мне выпало рефакторить проект с довольно сложным дата-пайплайном. Ключевой проблемой была скорость работы всей этой махины: переобучить модель было практически невозможно, так как подгрузка ретроданных для этого занимала целый год (!), а до отпуска оставался всего месяц. Отчаянные попытки дотащить этот проект и успеть переобучить модель до ухода в отпуск и привели к появлению этой статьи. Так что кидайте этот набор лайфхаков товарищу, у которого код с Pandas-датафреймами работает дольше, чем может позволить бренная человеческая жизнь.
Завязка
Итак, вводные:
Исходный пайплайн подтягивал данные со скоростью один день данных за один день. Данные нужны были за год, а значит и их подгрузка заняла бы целый год.
Управиться надо было за месяц — до начала отпуска.
Нужно было что-то сделать с памятью, потому что код постоянно вылетал по Out of memory error (OOM).
Итого, весь пайплайн нужно ускорить как минимум раз в 20, а за одно научиться жрать меньше памяти.
Дисклеймер, который предвосхищает комментарии с советами переходить на Spark:
На тот момент мы решили не менять Pandas на что-то другое, потому что код ещё дорабатывался в блокнотах разработчиками, которые были незнакомы со Spark.
Итак, начинаем оптимизировать. Первое, на что стоит обратить внимание когда код работает дольше, чем хотелось — это, конечно, циклы. С них и начнём.
Итерируемся
В документации Pandas сказано, что есть два основных метода итерации по датафрейму — iterrows и itertuples. С iterrows можно обратиться к конкретному значению по названию колонки. Itertuples перебирает таплы, поэтому работать с ним чуть сложнее — нужно помнить, какой индекс соответствует каждой колонке.
При этом, согласно документации, с itertuples производительность должна быть выше. Бенчмарк ниже это подтверждает, особенно разница в скорости становится видна при росте количества строк в датафрейме.
Казалось бы, лайфхак проверен, можно менять все iterrows на itertuples в своем проекте и лутать профит в скорости. Собственно, я заменила все iterrows на itertuples в своём проекте, но, внезапно, он стал работать ещё медленнее и вылетать по памяти.
Оказалось, что при большом количестве колонок itertuples начинает работать хуже.
Из графика видно, что в датафреймах с большим количеством колонок itertuples отрабатывает заметно медленнее.
Теперь про память. Как видно из профилировщика ниже, iterrows примерно в пять раз менее прожорливый.
Итого, если у вас в датафрейме больше 1000 колонок, стоит идти вразрез с официальной документацией и использовать iterrows вместо itertuples: сэкономите и память и время.
Конкретно в моём проекте до рефакторинга было 21 (!) iterrows и ноль itertuples. В итоге там, где было мало колонок, я поменяла метод итерации на itertuples, что позволило ускориться в два раза. Результат отличный, но недостаточный.
Применяем функции
Продолжаем читать документацию Pandas. Там говорят, что вообще-то итерироваться не стоит, потому что это слишком медленно. Вместо этого нужно применять apply, а в идеале — просто умножать на вектор или матрицу.
Я сравнила все варианты: for, apply, NumPy и apply с аргументом raw=True (с этим аргументом одна строка обрабатывается как numpy.ndarray, а не Pandas.Series) — это классная штука, которую почему-то мало кто использует.
Судя по графику, for работает очень медленно, а все остальные примерно на одном уровне. Но мы помним, что было в прошлом эксперименте, поэтому попробуем увеличить количество колонок. И тут ситуация становится уже иной: синий график for оказался прямо под apply, а apply с аргументом raw=True и NumPy — примерно на одном уровне.
По памяти картина приблизительно такая же: лучший результат у apply с аргументом raw=True. Для меня это оказался оптимальный вариант. Во-первых, не нужно вносить много изменений в код, во-вторых — выигрываем по времени и памяти.
В нашем проекте было 58 apply и 0 из них с аргументом raw=True. Замена всех апплаев на raw-вариант помогла ускориться в 1,8 раз. Однако, это всё ещё слишком медленно, поэтому продолжаем.
Мёрджимся
Merge или loc
Pandas позиционирует себя как библиотека, которая умеет делать много всего с табличными данными. В том числе, там есть функции, с помощью которых можно сделать абсолютно любой join (как в SQL). В проекте, подвергнутому рефакторингу, их хоть отбавляй. Например, в фрагменте ниже мы извлекаем из большого датафрейма те строки, которые имеют определённые индексы — и делаем это почему-то через merge. Выбор merge для этой операции сходу напрягает. Попробуем заменить merge на loc — ожидаемо, это действительно работает быстрее.
По памяти тоже есть небольшой выигрыш. Но не такой существенный, как когда мы игрались с итерациями.
Вроде бы все отлично, но внимательный читатель наверняка заметит, что зелёная строчка делает не то же самое, что красная. Если у нас есть key (пара индексов), которого нет в нашем df, то зелёная строчка сломается, потому что мы не можем выбрать индекс, которого нет в датафрейме. Красная строчка отработает корректно, выполнив inner join.
Конкретно для этого куска кода это нормально: нужные индексы гарантировано существуют в датафрейме исходя из логики кода, так что можно безбоязненно менять merge на loc и тем самым чуть ускоряться.
Но что делать, если нам нужен честный (и, главное, быстрый) left join?
Покажу кастомную функцию loc_with_fill. Глобально её суть сводится к последней строчке: мы берём наш датафрейм и используем set_index(), как делали в случае, когда переписывали на loc, а потом достаём данные, обращаясь напрямую к индексу, но делаем это при помощи reindex. По сути, reindex — это то же самое, что и loc, просто он подставляет NaN во всех колонках для строчек с «несуществующим» индексом.
Если посмотрим на скорость, то увидим, что с loc_with_fill всё работает в два раза быстрее.
По памяти всё не так феерично, но небольшое улучшение тоже есть.
Merge или join
Возможно, у вас возник логичный вопрос: если мы так много переписывали merge на что-нибудь ещё, и оно работало быстрее и лучше, то зачем нам вообще merge? Я тоже так подумала и решила заменить его на стандартный join.
До этого мы брали только индексы из одного датафрейма, и не использовали другие колонки. Но чаще встречается ситуация, когда нам нужно из обоих датафреймов тянуть либо все, либо несколько колонок.
Я попыталась переписать это на join, но не очень успешно. Видно, что голый join отрабатывает хуже, чем merge, что в целом не удивительно, так как и сам Pandas пишет в доке, что merge классно работает для общего случая.
Ещё merge требует в два раза меньше памяти по сравнению с join. Так что этот эксперимент считаем провальным и делаем вывод, что при виде merge в коде не стоит судорожно менять его на join. Нужно сначала разобраться с какими данными мы работаем и какую операцию делаем.
Indicator и его оптимизация
Последний кейс слегка странный — я нашла место в коде, где использовался merge, чтобы понять, какие ключи есть в кэше.
При передаче аргумента indicator=True в итоговый датафрейм добавляется колонка, в которой указано, откуда пришел тот или иной индекс. Там могут быть значения типа: both, left only, right only и т.д. По крайней мере, так задумал автор кода.
Мне показалось это странным и неэффективным. Поэтому я переписала логику — взяла множество ключей в одном датафрейме, множество в другом, и посмотрела на их пересечения.
Работало это чуть-чуть быстрее, но не так феерично, как все прошлые штуки.
Вывод к моему удивлению такой: можно и нужно использовать merge с аргументом indicator=True. Скорость и память проседают несильно, но писать код будет значительно проще.
В проекте, который я рефакторила, merge использовался 56 раз. Больше половины из них я переписала на loc и reindex. Это дало ускорение в два раза, но мы всё ещё не успеваем до отпуска :)
Типы данных
Int, uint, float
Теперь давайте применять стандартные подходы к ускорению, которые работают не только с Pandas, но и с произвольным кодом — попробуем поменять тип данных. Начнём с тех, которые есть примерно везде. Pandas поддерживает и int8, и int64, но по дефолту всегда создаёт int64 и float64.
На бенчмарке видно, что int64 отрабатывает дольше остальных, а самый быстрый - int8 (что, в целом, неудивительно). При этом int64 и uint64, а также int32 и uint32 — совпали. Но есть какой-то аномальный int16. Тут причина, вероятно, в том, что на разных машинах конкретная разрядность иногда отрабатывает быстрее, чем все остальные.
По памяти int8 и uint8 — самые экономные. Int64 просит подозрительно мало, скорее всего потому, что я конвертирую данные из NumPy в Pandas. А NumPy по дефолту тоже использует int64. И поэтому в ячейках дополнительно замеряется память конвертации, а при int64 эта конвертация не происходит.
Похожая ситуация с float: 64 отрабатывает дольше, чем 32, а памяти требует меньше. Но здесь делаем скидку на то, что я конвертирую под капотом.
Categorical и sparse
Ещё у Pandas есть свои специфичные типы данных — categorical и sparse.
Categorical классно работает, когда у вас много похожих значений: Pandas превращает их в категории и благодаря этому работает с ними чуть более оптимально. А sparse подходит для датафреймов с большим количеством отсутствующих значений. Он не запоминает ячейки, равные NaN, а запоминает только те строки, где есть значение.
В моём коде есть строка, которая в теории должна выбирать — categorical или sparse. Я построила бенчмарк и оказалось, что в общем случае на операции group&count categorical всегда работает лучше, чем string, вне зависимости от того, сколько у нас уникальных значений.
У sparse похожая картина: здесь я смотрела на операцию суммы, и получилось, что в целом он тоже работает быстрее, чем string. И это не зависит от процента нулей.
Конечно, было бы логичнее тестировать categorical и sparse на одной операции. Но оказалось, что sparse-колонки нельзя группировать. Поэтому этот тип данных не всегда получится безболезненно внедрить в код.
По памяти categorical и sparse тоже отрабатывают лучше, чем string.
В своём проекте я поменяла string на sparse и categorical, int64 на int8, и получила ускорение в 2,5 раза. Пока это самая большая цифра, но мы всё ещё не успеваем закончить до отпуска. Хотя уже перестали вылетать по памяти.
Параллелизация
Чтобы параллелизовать код на Pandas, можно использовать:
pandarallel.
multiprocessing.
В моём сетапе использовалось 4 воркера. По бенчмарку видно, что pandarallel работает быстрее, чем pandas apply. При этом pandas apply не имеет смысла применять, если у вас данные маленького размера — в этом случае он будет быстрее параллелизма, так как не будет оверхеда.
При этом, чтобы добавить pandarallel, вам практически не нужно менять код.
Ииии… Нам это дало ускорение в 3 раза. Наш Гарольд наконец доволен, потому что время до конца прогрузки данных меньше, чем до начала отпуска. И у меня даже остался один день на то, чтобы переобучить модель.
Выводы
Кратко сформулировала главные инсайты:
iterrows не имеет смысла переписывать на itertuples, если у вас больше 1000 колонок.
Apply прикольно использовать с аргументом raw=True — можно безболезненно получить буст скорости и памяти.
Merge стоит переписать на loc или reindex, если от одного из двух датафреймов вам нужны только индексы.
Имеет смысл потратить немного времени и переписать int64, float64 на int8 и float32.
String можно заменить на sparse и categorical, но проверяйте, чтобы код после этого не сломался.
Добавляйте параллелизацию, лучше именно pandarallel, а не самописную (но это стоит делать только если у вас всё хорошо с памятью).
Из всего этого списка проще всего сделать параллелизацию и поменять типы данных, а потом уже заморачиваться с переписываем apply или iterrows.
Ну и не забывайте вовремя ходить в отпуск!