Использование библиотек Polar вместо Pandas: подробный анализ производительности

В этой статье мы рассмотрим три реальные задачи на основе реальных данных, в которых Polars превосходит Pandas по всем показателям.

# Введение

В течение последнего десятилетия Pandas был основой для работы с данными в Python. Для наборов данных, помещающихся в память, он достаточно быстр и знаком, поэтому переход на другую библиотеку редко приходит в голову программисту.

Однако, как только вы начинаете работать с миллионами строк, начинают проявляться недостатки: операции группировки, занимающие несколько секунд, промежуточные копии, потребляющие оперативную память, и оконные функции, которые выполняются как циклы на уровне Python, а не как векторизованный код на C или Rust .

Polars — это библиотека DataFrame, написанная на Rust на основе Apache Arrow . Она разработана с учетом параллелизма и ленивой оценки как основных функций. Pandas выполняет каждую операцию заранее и последовательно, тогда как Polars может составить план запроса и оптимизировать его перед выполнением, при этом большинство операций автоматически выполняются параллельно на всех доступных ядрах ЦП.

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

# Использование функции rank() вместо with_row_count(): Рейтинг активности

В этом задании цель состоит в том, чтобы определить рейтинг активности электронной почты для каждого пользователя на основе общего количества отправленных писем. Пользователь с наибольшим количеством писем получает рейтинг 1. Результаты должны быть отсортированы по общему количеству писем в порядке убывания, используя алфавитный порядок в качестве критерия разрешения неоднозначностей, и каждый рейтинг должен быть уникальным, даже если у двух пользователей одинаковое количество писем.

// Представление данных

В таблице google_gmail_emails хранится одна строка для каждого отправленного электронного письма, содержащая идентификатор отправителя (from_user), идентификатор получателя (to_user) и дату отправки письма. Вот предварительный просмотр таблицы:

Grain (что означает одна строка выходных данных): один пользователь, с указанием общего количества его электронных писем и уникального рейтинга активности.

// Распространенная ошибка

В задании требуется присвоить уникальный ранг, даже если у двух пользователей одинаковое количество электронных адресов. Распространенная ошибка — использование метода `rank(method='dense')` в Pandas, который присваивает одинаковый ранг пользователям с одинаковым количеством адресов. Правильный метод — `first`, который разрешает совпадения по положению в отсортированном фрейме. Поскольку перед ранжированием мы сортируем по алфавиту по `user_id`, полученные ранги уникальны и детерминированы.

Оптимальное решение Polars полностью исключает функцию ранжирования. После сортировки по [«total_emails», «user_id»] в порядке убывания и возрастания соответственно, условие .with_row_count(«activity_rank», offset=1) присваивает последовательные целые числа, начиная с 1. Логика разрешения неоднозначностей не требуется, поскольку сортировка уже это обработала.

// Решения

1. Решение с использованием Pandas

Мы переименовываем from_user в user_id, группируем по пользователю, подсчитываем количество электронных писем, вычисляем первый ранг и сортируем по количеству электронных писем в порядке убывания, используя алфавитный метод для разрешения неоднозначностей.

import pandas as pd import numpy as np google_gmail_emails = google_gmail_emails.rename(columns={«from_user»: «user_id»}) result = google_gmail_emails.groupby( ['user_id']).size().to_frame('total_emails').reset_index() result['activity_rank'] = result['total_emails'].rank(method='first', ascending=False) result = result.sort_values(by=['total_emails', 'user_id'], ascending=[False, True])

2. Решение полярных растворов

Мы используем ленивую цепочку, которая переименовывает, группирует, сортирует и присваивает номера строкам за один проход. Вызов метода `.collect()` в конце материализует результат.

import polars as pl google_gmail_emails = google_gmail_emails.rename({«from_user»: «user_id»}) result = ( google_gmail_emails.lazy() .group_by(«user_id») .agg(total_emails = pl.count()) .sort( by=[«total_emails», «user_id»], descending=[True, False] ) .with_row_count(«activity_rank», offset=1) .select([ pl.col(«user_id»), «total_emails», «activity_rank» ]) .collect() )

// Сравнение производительности

Решение Pandas выполняет итерацию по данным дважды после группировки: один раз для вычисления размеров и один раз для присвоения рангов. Внутри функции rank(method='first') выделяется массив рангов, разрешаются совпадающие значения с помощью argsort и записываются обратно — что значительно дороже, чем кажется для одного столбца. Функция group_by в Polar распределяет рабочую нагрузку между всеми доступными ядрами ЦП, что приводит к значительно более быстрой агрегации для больших таблиц. А поскольку оператор .with_row_count() представляет собой один последовательный проход O(n) после сортировки, он заменяет функцию ранжирования самой дешевой возможной операцией. В таблице, содержащей миллионы записей электронной почты, использование параллельной агрегации без функции ранжирования может привести к 5–10-кратному сокращению времени выполнения по сравнению с подходом Pandas.

Вот предварительный просмотр результата выполнения кода:

# Использование cumcount() + pivot() вместо over(): поиск покупок пользователей

В этом задании нас просят определить активных пользователей, совершивших повторную покупку, — а именно, тех, кто совершил вторую покупку в течение 1 и 7 дней после первой. Покупки, совершенные в один и тот же день, не должны включаться. Результатом будет просто список соответствующих значений user_id.

// Представление данных

В таблице amazon_transactions содержится одна строка на каждую покупку, включающая user_id, item, created_at date и revenue.

Вот предварительный просмотр таблицы:

Grain (что означает одна строка выходных данных): один идентификатор пользователя, совершившего соответствующую возвратную покупку в течение 7 дней с момента первой.

// Крайний случай

Покупки, совершенные в один день, не учитываются, то есть промежуток между первой и второй покупкой должен превышать 0 дней и составлять не более 7 дней. Клиент, совершивший две покупки в один день, не имеет права на льготу.

// Решения

Оба решения находят самую раннюю дату покупки каждого пользователя, а затем фильтруют последующие покупки в течение 1–7 дней. Следует обратить внимание на один момент: если поле created_at содержит временные метки вместо обычных дат, необходимо обрезать значение до даты перед сравнением. В противном случае две покупки, совершенные в разное время в один и тот же день, будут некорректно проходить проверку на строгое неравенство.

1. Решение с использованием Pandas

В Pandas решение включает в себя выделение уникальных дат покупок для каждого пользователя, ранжирование их с помощью функции cumcount(), преобразование данных для получения первой и второй дат рядом и вычисление разницы в днях.

import pandas as pd amazon_transactions[«purchase_date»] = pd.to_datetime(amazon_transactions[«created_at»]).dt.date daily = amazon_transactions[[«user_id», «purchase_date»]].drop_duplicates() ranked = daily.sort_values([«user_id», «purchase_date»]) ranked[«rn»] = ranked.groupby(«user_id»).cumcount() + 1 first_two = (ranked[ranked[«rn»] = 1) & (first_two[«diff»] pl.col(«first_purchase_date»)) & (pl.col(«created_at»)

Источник: www.kdnuggets.com