Вид на Polars с высоты птичьего полета

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

Общий обзор

Итак, что же такое Polars? Коротко можно описать так: “движок запросов с фронтендом DataFrame”. Это слишком высокий уровень даже для обзора с высоты птичьего полета. Поэтому давайте углубимся в эти два элемента – DataFrame и механизм запросов – и рассмотрим, как выполняется запрос. Пройдя пошаговый путь выполнения запроса, мы сможем увидеть каждый компонент в действии и понять его роль и назначение.

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

Запрос

Когда вы взаимодействуете с Polars, вы используете наш API DataFrame. Этот API специально разработан для параллельного выполнения и с учетом производительности. Написание запроса в Polars в этом смысле – это написание небольшой программы (или в данном случае запроса) на языке, специфичном для данной области (DSL), разработанном Polars. Этот DSL имеет свой собственный набор правил, определяющих, какие запросы являются корректными, а какие – нет.

Для этого поста давайте воспользуемся известным набором данных по такси на Нью-Йоркские праздники, содержащим данные о поездках на такси1. В примере ниже мы вычисляем среднюю стоимость минуты поездки стоимостью более 25 долларов по зонам. Этот пример достаточно прост для понимания и в то же время достаточно глубок, чтобы продемонстрировать назначение механизма запросов.

import polars as pl

query = (
    pl.scan_parquet("yellow_tripdata_2023-01.parquet")
    .join(pl.scan_csv("taxi_zones.csv"), left_on="PULocationID", right_on="LocationID")
    .filter(pl.col("total_amount") > 25)
    .group_by("Zone")
    .agg(
        (pl.col("total_amount") /
        (pl.col("tpep_dropoff_datetime") - pl.col("tpep_pickup_datetime")).dt.total_minutes()
        ).mean().alias("cost_per_minute")
    ).sort("cost_per_minute",descending=True)
)

Приведенный выше запрос имеет тип LazyFrame. Он возвращается мгновенно, в то время как набор данных о поездках на такси в Нью-Йорке превышает 3 миллиона строк, так что же произошло? Оператор определяет запрос, но еще не выполняет его. Эта концепция известна как ленивая оценка и является одним из ключевых преимуществ Polars. Если вы посмотрите на структуру данных на стороне Rust, то увидите, что она содержит два элемента: логический_план и флаги конфигурации для оптимизатора opt_state.

pub struct LazyFrame {
    pub logical_plan: LogicalPlan,
    pub(crate) opt_state: OptState,
}

Логический план представляет собой дерево, в котором источники данных являются листьями дерева, а преобразования – узлами. План описывает структуру запроса и содержащиеся в нем выражения.

pub enum LogicalPlan {
    /// Filter on a boolean mask
    Selection {
        input: Box<LogicalPlan>,
        predicate: Expr,
    },
    /// Column selection
    Projection {
        expr: Vec<Expr>,
        input: Box<LogicalPlan>,
        schema: SchemaRef,
        options: ProjectionOptions,
    },
    /// Join operation
    Join {
        input_left: Box<LogicalPlan>,
        input_right: Box<LogicalPlan>,
        schema: SchemaRef,
        left_on: Vec<Expr>,
        right_on: Vec<Expr>,
        options: Arc<JoinOptions>,
    },
    ...
}

Одним из важных шагов при преобразовании запроса в логический план является проверка. Polars заранее знает схему данных и может проверить правильность преобразований. Это гарантирует, что вы не столкнетесь с ошибками на полпути к выполнению запроса. Например, определение запроса, в котором вы выбираете несуществующий столбец, возвращает ошибку перед выполнением

pl.LazyFrame([]).select(pl.col("does_not_exist"))

polars.exceptions.ColumnNotFoundError: column_does_not_exist

Error originated just after this operation:
DF []; PROJECT */0 COLUMNS; SELECTION: "None"

Мы можем просмотреть логический план, вызвав show_graph на LazyFrame:

query.show_graph(optimized=False)

Оптимизация запросов

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

query.show_graph()

На первый взгляд может показаться, что оба плана (оптимизированный и неоптимизированный) одинаковы. Однако произошли две важные оптимизации: Projection pushdown и Predicate pushdown.

Компания Polars проанализировала запрос и отметила, что используется лишь небольшой набор столбцов. Для данных о поездках есть четыре столбца. Для данных о зонах – два столбца. Чтение всего набора данных было бы расточительством, поскольку другие столбцы не нужны. Поэтому, проанализировав ваш запрос, Projection Pushdown значительно ускорит чтение данных. Оптимизацию можно увидеть в узлах листа под π� 4/19 и π� 2/4.

С помощью Predicate pushdown Polars фильтрует данные как можно ближе к источнику. Это позволяет избежать чтения данных, которые на более позднем этапе запроса будут отброшены. Узел фильтра был перемещен в паркетный ридер под σ�, что означает, что наш ридер будет немедленно удалять строки, которые не соответствуют нашему фильтру. Следующая операция объединения будет выполняться гораздо быстрее, так как данных поступает меньше.

Polars поддерживает ряд оптимизаций, с которыми можно ознакомиться здесь.

Выполнение запросов

После того как логический план оптимизирован, наступает время его выполнения. Логический план – это план того, что пользователь хочет выполнить, а не того, как это сделать. Именно здесь в игру вступает физический план. Наивным решением было бы иметь один алгоритм объединения и один алгоритм сортировки; таким образом, можно было бы выполнять логический план напрямую. Однако это сопряжено с огромными затратами на производительность, поскольку знание характеристик ваших данных и среды, в которой вы работаете, позволяет Polars выбирать более специализированные алгоритмы. Таким образом, существует не один алгоритм join, а несколько, каждый со своим уникальным стилем и производительностью. Планировщик запросов преобразует LogicalPlan в PhysicalPlan и выбирает лучшие алгоритмы для запроса. Затем наш вычислительный механизм выполняет операции. В этом посте мы не будем подробно рассказывать о модели выполнения наших движков и о том, как им удается работать так быстро. Это мы оставим на другой раз.

Если посмотреть на разницу в производительности обоих планов (оптимизированный и неоптимизированный), то можно увидеть 4-кратное улучшение. В этом и заключается сила ленивого выполнения и использования движка запросов вместо того, чтобы с нетерпением оценивать каждое выражение по порядку. Это позволяет механизму оптимизировать и избежать ненужной работы. Все это улучшение происходит без особых затрат для пользователя, поскольку все, что ему нужно сделать, – это написать запрос. Вся сложность скрыта внутри механизма запросов.

%%time
query.collect(no_optimization=True);

CPU times: user 2.45 s, sys: 1.18 s, total: 3.62 s
Wall time: 544 ms

%%time
query.collect();

CPU times: user 616 ms, sys: 54.2 ms, total: 670 ms
Wall time: 135 ms

Заключение

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