Почему не всегда Pandas — лучший выбор (и когда стоит попробовать Polars)

Меня зовут Данила Ляпин, я Senior Data Scientist в Яндексе и автор курса «Специалист по Data Science» ^[1] в Яндекс Практикуме.

В современном мире анализа данных пользу библиотеки Pandas трудно переоценить — она используется везде экспертами любого уровня: от стажёров до техлидов, а последние годы это де-факто стандарт в аналитике.

У Pandas есть огромная экосистема с большой базой знаний и интеграциями в различные библиотеки (тут и встроенная визуализация от matplotlib, и переход к данным в numpy формате, и много чего ещё). Практически каждый разведывательный анализ данных начинается с таких слов, как: read_csv, describe, head, isna().sum(). 

Но объём данных растёт ежегодно, память ^[2] дорожает, а пайплайны усложняются. Чтобы загрузить локально большой датасет и обработать его, приходится не один раз словить “kernel crashed”:

И вот на этом моменте самое время поговорить про Polars, потому что там с этим как раз всё в порядке (и не только с этим :-)

Что такое Polars и чем он хорош

Polars — это dataframe-библиотека (есть реализация на Python и Rust), ориентированная на скорость и эффективность памяти. Основана на Apache Arrow — совместимом колоночном представлении. Может работать в двух режимах:

1. Pandas-like: «сделал операцию — получил результат».

2. Ленивые вычисления: строим план вычислений, затем оптимизируем его и выполняем.

Глядя на логотип, можно увидеть отсылку к колоночной природе Polars :-)

Если посмотреть на официальную документацию (и отталкиваться от практики реального использования), у Polars довольно много плюсов:

• работает быстро, потому что написан с нуля на Rust и без внешних зависимостей

• поддерживает все распространённые уровни хранения данных: локальные, облачные хранилища и базы данных

• очень понятный API: пишем запросы так, как удобно, а Polars сам определит наиболее эффективный способ их выполнения (с помощью оптимизатора запросов)

• потоковый API позволяет обрабатывать результаты без того, чтобы хранить все данные в памяти одновременно

• параллельная обработка всего использует «железо» по максимуму, распределяя нагрузку между доступными ядрами процессора без дополнительной настройки

• векторизованный механизм выполнения запросов

• при необходимости можно выполнять запросы на графических процессорах NVIDIA для достижения максимальной производительности (ну или там, где ещё есть CUDA)

• Polars может принимать и генерировать данные в формате Apache Arrow, часто с использованием операций без копирования.

❗️ Обратите внимание ^[3], что Polars опирается на Arrow-совместимое колоночное представление данных, но не использует PyArrow как основной движок вычислений. Вместо этого вычисления и работа с памятью реализованы в самом Polars.

Установка Polars

Тут всё максимально просто: достаточно написать pip install polars и уже можно пользоваться.  Библиотека не имеет зависимостей (вообще никаких!), поэтому можно не переживать, что потом для запуска внезапно потребуется что-то ещё, а этого нет.

Импортируется чаще всего как import polars as pl, но никто не может помешать вам импортировать её любым другим способом (хоть так: import polars as pandas, но делать так мы точно не рекомендуем:-)

Основные типы данных

Если вы знакомы с пандасовскими объектами Series и DataFrame, то здесь мало что меняется: в Polars также реализованы свои типы этих данных. На практике очень удобно переходить к датафреймам Polars от Pandas (и наоборот):

polars_df = pl.from_pandas(pandas_df)  # Polars DataFrame
pandas_df = polars_df.to_pandas()  # Pandas DataFrame

Но кроме этих объектов в Polars есть несколько специфичных типов данных:

LazyFrame (ленивый датафрейм). Можем перейти к нему от обычного датафрейма (и наоборот):

lazy_polars_df = polars_df.lazy()  # LazyFrame
polars_df = lazy_polars_df.collect()  # DataFrame

Schema (названия столбцов и их типы данных). Можно проверить схему датафрейма, вызвав:

print(polars_df.schema)  
# Получим что-то вида:
# Schema({'name': String, 'birthdate': Date, 'weight': Float64, 'height': Float64})

Lazy API

Lazy API позволяет работать с LazyFrame-объектами. При использовании lazy API Polars не выполняет каждый запрос построчно, а обрабатывает весь запрос целиком, причём он выполняется только после команды collect(). Отсрочка выполнения до последнего момента может значительно повысить производительность, поэтому в большинстве случаев предпочтение отдаётся Lazy API. Давайте рассмотрим это на примере:

q = (
    pl.scan_csv("../data/iris.csv")
    .filter(pl.col("sepal_length") > 5)
    .group_by("species")
    .agg(pl.col("sepal_width").mean())
)
df = q.collect()

При использовании Lazy API вы можете использовать функцию объяснения, чтобы попросить Polars создать описание плана запроса, который будет выполнен после сбора результатов. Это может быть полезно, если вы хотите увидеть, какие типы оптимизации выполняет Polars по вашим запросам. Мы можем попросить Polars объяснить запрос, который мы определили выше: q.explain()

print(q.explain())

Использование lazy API позволяет выполнить несколько оптимизаций запроса. Некоторые из них выполняются заранее, другие — по мере поступления материализованных данных. Вот пример таких оптимизаций из жизни:

• применить фильтры как можно раньше / на уровне сканирования файла

• выбрать только те столбцы, которые нужны на уровне сканирования

• загрузить только необходимый срез с уровня сканирования

• определить, какие объединения следует выполнять в первую очередь, чтобы снизить нагрузку на память

• привести типы данных к чему-то одному для успешного выполнения операций с минимальным потреблением памяти

• оценить количество элементов для определения оптимальной стратегии группировки

Если нужно больше оптимизаций (а их мало не бывает), примеры с разборами есть в официальных доках ^[4].

Коротко про потоковый режим

Ещё одним преимуществом ленивого API является то, что он позволяет выполнять запросы в потоковом режиме. Вместо обработки всех данных одновременно Polars может выполнять запрос пакетно, что позволяет обрабатывать наборы данных, которые не помещаются в памяти. 

Чтобы сообщить Polars, что мы хотим выполнить запрос в потоковом режиме, мы передаём аргумент engine=”streaming” для сбора:

q1 = (
    pl.scan_csv("../data/iris.csv")
    .filter(pl.col("sepal_length") > 5)
    .group_by("species")
    .agg(pl.col("sepal_width").mean())
)
df = q1.collect(engine="streaming")

Поддержка SQL-режима

Начинающим специалистам по данным SQL куда ближе, чем Python, поэтому данный режим может быть вполне полезен. 

❗️ Важно: в Polars отсутствует отдельный SQL-движок, поскольку он преобразует SQL-запросы в выражения, которые затем выполняются с помощью собственного обработчика. Такой подход гарантирует, что Polars сохраняет свои преимущества в производительности и масштабируемости как нативная библиотека DataFrame, одновременно предоставляя пользователям возможность работать с SQL.

Polars использует объект SQLContext для управления SQL-запросами. Контекст содержит сопоставление имен идентификаторов DataFrame и LazyFrame с соответствующими наборами данных. Вот как он создаётся:

ctx = pl.SQLContext()

Есть несколько способов регистрации объектов DataFrame во время инициализации SQLContext:

• зарегистрировать все объекты LazyFrame и DataFrame в глобальном пространстве имён

• зарегистрировать явно с помощью сопоставления словарей или ключевых аргументов (kwargs)

Посмотрим на оба варианта в коде:

df = pl.DataFrame({"a": [1, 2, 3]})
lf = pl.LazyFrame({"b": [4, 5, 6]})

# Зарегистрировать все датафреймы в глобальном пространстве имен: df как "df", lf как "lf"
ctx = pl.SQLContext(register_globals=True)

# Зарегистрировать явное сопоставление имени идентификатора с датафреймом
ctx = pl.SQLContext(frames={"table_one": df, "table_two": lf})

# Зарегистрировать фреймы с помощью kwargs: dataframe df как "df", а  lazyframe lf как "lf"
ctx = pl.SQLContext(df=df, lf=lf)

Ещё мы можем зарегистрировать DataFrames из библиотеки Pandas, предварительно преобразовав их в Polar-объекты (подробнее, как это работает: https://docs.pola.rs/user-guide/sql/intro ^[5]):

import pandas as pd

df_pandas = pd.DataFrame({"c": [7, 8, 9]})
ctx = pl.SQLContext(df_pandas=pl.from_pandas(df_pandas))

После можем запускать наши SQL-запросы:

iris = pl.read_csv(
    "../data/iris.csv"
)
with pl.SQLContext(register_globals=True, eager=True) as ctx:
    df_small = ctx.execute("SELECT * from iris LIMIT 5")
    print(df_small)

SQL-запросы при этом можно выполнять из нескольких источников:

with pl.SQLContext(
    data=pl.scan_csv(DATA_FILE_2),
    districts=pl.scan_csv(DATA_FILE_2_DISTRICTS),
    groups=pl.scan_csv(DATA_FILE_2_GROUPS),
    eager=True,
) as ctx:
    query = """
    SELECT
        user_id,
        sex,
        districts.district,
        groups.group
    FROM
        data
    LEFT JOIN districts USING (user_id)
    LEFT JOIN groups USING (user_id)
    """
    print(ctx.execute(query))

На всякий случай напомню и повторюсь, что Polars не поддерживает всю спецификацию SQL, но поддерживает подмножество наиболее распространённых типов операторов. Уже доступны: SELECT, WHERE,ORDER,LIMIT,GROUP BY,UNION, JOIN, но оконных функций пока нет. Обещают добавить, но пока вот так.

Примеры кода из серии «одна задача — два решения»

Теперь посмотрим, насколько различается синтаксис библиотек Pandas и Polars при выполнении одних и тех же задач (спойлер: местами практически идентичен!).

1. Чтение большого файла + фильтр + агрегация

Цель: сравнить производительность при работе с большим parquet-файлом (5.5 млн строк, 65 колонок) при выполнении типичной аналитической задачи — фильтрации и агрегации данных.

(Файл DATA_FILE_1.parquet) ^[6]

Классический код на Pandas:

data = pd.read_parquet(DATA_FILE_1)
data = data[data['target'] == 1]
data_aggr = data.groupby('user_id')['item_price'].sum()

Код на Polars:

data = pl.scan_parquet(DATA_FILE_1)
data = data.filter(pl.col('target') == 1)
data_aggr = data.group_by('user_id').agg(pl.col('item_price').sum()).collect()

Как видим, код получился практически одинаковый (с учётом синтаксиса, конечно).

2. Join нескольких таблиц + расчёт метрик

Цель: проверить эффективность объединения нескольких таблиц и последующего расчёта бизнес-метрик.

(Файлы DATA_FILE_2.csv, DATA_FILE_2_DISTRICTS.csv и DATA_FILE_2_GROUPS.csv) ^[6]

Сначала напишем, используя Pandas:

data = pd.read_csv(DATA_FILE_2)
districts = pd.read_csv(DATA_FILE_2_DISTRICTS)
data = data.merge(districts, on='user_id', how='left')
groups = pd.read_csv(DATA_FILE_2_GROUPS)
data = data.merge(groups, on='user_id', how='left')

data_aggr = data.groupby('group', as_index=False).agg({'payment_page': 'sum', 'payment': 'sum'})
data_aggr['conversion'] = data_aggr['payment'] / data_aggr['payment_page']

А потом — с помощью Polars:

data = pl.scan_csv(DATA_FILE_2)
districts = pl.scan_csv(DATA_FILE_2_DISTRICTS)
data = data.join(districts, on='user_id', how='left')
groups = pl.scan_csv(DATA_FILE_2_GROUPS)
data = data.join(groups, on='user_id', how='left')

data_aggr = data.group_by('group').agg(pl.col('payment_page').sum(), pl.col('payment').sum())
data_aggr = data_aggr.with_columns(
    (pl.col('payment') / pl.col('payment_page')).alias('conversion')
).collect()

Тоже довольно близко друг к другу по коду (но не по производительности, можете проверить, хотя к ней мы тоже сейчас перейдём отдельно).

3. Оконные функции (ranking, rolling)

Цель: сравнить производительность оконных функций.

(Файл DATA_FILE_3.csv) ^[6]

Сначала — на Pandas:

data = pd.read_parquet(DATA_FILE_1)
data["rank"] = data.groupby("user_id")["ials_score"].rank(method="dense", ascending=False)
data['rolling_mean'] = data.groupby('user_id')['item_price'].rolling(10).mean().values

И на Polars:

data = pl.scan_parquet(DATA_FILE_1)
data = data.with_columns([
          pl.col("ials_score")
            .rank(method="dense", descending=True)
            .over("user_id")
            .alias("rank"),

          pl.col("item_price")
            .rolling_mean(window_size=10)
            .over("user_id")
            .alias("rolling_mean"),
      ]).collect()

Выводы по скорости и эффективности будут ниже, для затравки скажу, что когда первый раз видишь такой отрыв — происходит довольно сильная переоценка ценностей :-)

4. Строковые операции (regex, split)

Цель: проверить, как обрабатываются текстовые данные с использованием регулярных выражений и разделения строк.

(Файл DATA_FILE_4.csv) ^[6]

Код на Pandas:

data = pd.read_csv(DATA_FILE_4)
data['race/ethnicity'] = data['race/ethnicity'].str.split(' ').str[1]
data['parental level of education'] = data['parental level of education'].str.replace(r"[^ws]", "", regex=True)

И код на Polars:

data = pl.scan_csv(DATA_FILE_4)
data = data.with_columns(
    [
        pl.col("race/ethnicity")
          .str.split(" ")
          .list.get(1)
          .alias("race/ethnicity"),
        pl.col("parental level of education")
          .str.replace_all(r"[^ws]", "")
          .alias("parental level of education"),
    ]
).collect()

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

5. Работа со «сложными» типами (JSON)

Цель: сравнить производительность работы с вложенными структурами данных (json/словари). Например, у нас есть csv-файл с метаданными, а нам нужно извлечь параметры из вложенной структуры:

Сделаем это на Pandas:

data = pd.read_parquet(DATA_FILE_5)
for param in ['accompaniment', 'description', 'general_mood', 'genre_tags', 'lead_instrument', 'production_quality', 'tempo_and_rhythm', 'vocal_presence']:
    data[param] = data['enriched_metadata'].str[param]
data['main_genre'] = data['genre_tags'].str[0]

А потом — на Polars (и посмотрим на разницу в синтаксисе, тут она более явная):

data = pl.scan_parquet(DATA_FILE_5)

for param in ['accompaniment', 'description', 'general_mood', 'genre_tags', 'lead_instrument', 'production_quality', 'tempo_and_rhythm', 'vocal_presence']:
    data = data.with_columns(
        (pl.col('enriched_metadata').struct.field(param)).alias(param)
    )

data = data.with_columns(
    (pl.col('genre_tags').list.get(0, null_on_oob=True)).alias('main_genre')
).collect()

Профилирование кода и результаты сравнения

Для каждого эксперимента выше я провёл профилирование по времени и по памяти. Для замера времени выполнения использовалась магическая команда %%timeit Python, а для замера памяти — memory_profiler ^[7].

На какой конфигурации производился замер:

• Python version: 3.13.0

• Pandas version: 2.2.3

• Polars version: 1.34.0

• железо: MacBook M3 Pro

Результаты профилирования:

Ключевые выводы

Polars — точно не универсальный инструмент для всех задач, которые можно делать с помощью Pandas, но во многих ситуациях он оказывается гораздо эффективнее. 

Если коротко по выводам, то получилось такое:

• Polars показывает ускорение от 2x до 120x в зависимости от задачи.

• Наибольший выигрыш достигается при работе с большими файлами и фильтрации данных.

• Polars значительно эффективнее использует память благодаря lazy evaluation.

• Даже на малых датасетах Polars не уступает по скорости.

Если у вас появились вопросы по Polars (или хотите разобрать кейсы из своей практики) — пишите в комментариях, постараюсь ответить и помочь.