Собесов

Сценарий: переписать pandas-пайплайн на polars lazy

PythonOptimisation и big dataСредняяMiddle

Условие

Есть pandas-скрипт: фильтр → группировка → join со справочником → сортировка. На 10 ГБ работает 8 минут. Переписать на polars lazy.

Решение

Pandas

import pandas as pd
df = pd.read_parquet('events.parquet')
df = df[df['country'] == 'RU']
g = df.groupby('user_id', as_index=False).agg(amount=('amount', 'sum'))
res = g.merge(users, on='user_id', how='left').sort_values('amount', ascending=False)

Polars lazy

import polars as pl
 
q = (
    pl.scan_parquet('events.parquet')
      .filter(pl.col('country') == 'RU')
      .group_by('user_id').agg(pl.col('amount').sum())
      .join(pl.scan_parquet('users.parquet'), on='user_id', how='left')
      .sort('amount', descending=True)
)
 
# План смотрим — что polars оптимизировал
print(q.explain())
 
res = q.collect()   # выполнение

Что polars делает сам

  • Projection pushdown: читает только нужные колонки из parquet.
  • Predicate pushdown: фильтр country = 'RU' уходит в чтение, мимо проходит мало данных.
  • Common subexpression elimination: повторные выражения считаются один раз.
  • Parallel execution: все ядра задействованы.

Streaming для очень больших данных

res = q.collect(streaming=True)   # обрабатывает потоком, не держит всё в RAM

Подводные камни

  1. Polars API не совпадает с pandas: group_by (не groupby), pl.col(...) обязателен в expressions, нет inplace.
  2. На lazy-плане ошибка может всплыть только при collect() — отлаживайте на .head(1000).collect().
  3. mergejoin, sort_valuessort. Привычные имена не работают.
  4. pd.to_datetimepl.col('ts').str.to_datetime(). Парсинг строгий.
  5. Если в проде exchange-формат — pandas, не забыть .to_pandas() в конце пайплайна.

Эталонный ответ

polars lazy с scan_parquet + filter + group_by + join + sort, всё за один .collect(). Predicate/projection pushdown даёт обычно 5–20× ускорение и в разы меньше RAM.

Хочешь увидеть разбор?

Зарегистрируйся бесплатно — откроется развёрнутое решение этой задачи и ещё 4 на выбор.

Зарегистрироваться и увидеть разбор
Уже есть аккаунт? Войти