Polars: быстрая библиотека для работы с DataFrame

Polars — это высокопроизводительная библиотека для работы с DataFrame, спроектированная для эффективной и быстрой обработки данных. Вдохновлённая популярной библиотекой pandas, она предлагает еще более мощные возможности для работы с большими наборами данных, которые могут не поместиться в оперативную память.

Polars объединяет гибкость и удобство Python с быстродействием и масштабируемостью Rust, что делает его привлекательным выбором для широкого спектра задач по работе с данными.

Особенности Polars:

1. Скорость и производительность. Ядро Polars написано на Rust — низкоуровневом языке, который работает без внешних зависимостей. Rust эффективно использует память и обеспечивает производительность, сравнимую с C или C++. Также Polars использует Apache Arrow для выполнения векторизованных запросов. Apache Arrow - платформой разработки, которая спроектирована для быстрой обработки данных в памяти и использует преимущества столбцового хранения. Это обеспечивает Polars сильный прирост производительности.
2. Поддержка многопоточности. Polars позволяет использовать все доступные ядра CPU параллельно. Это особенно важно при работе с большими наборами данных.
3. Простая в использовании. Если вы уже знакомы с pandas или другими подобными библиотеками, то вы быстро освоитесь и с Polars. Она предоставляет свой уникальный, но все же знакомый интерфейс, так что переход к Polars будет лёгким.

Начало работа с DataFrame в Polars

Для начала еобходимо установить Polars. Она поддерживает версии Python 3.7 и выше.

• Установка через командную строку/терминал: pip install polars
• Установка в Google Colab: !pip install polars
• Установка через терминал в Anaconda: conda install -c conda-forge polars

Как и в большинстве других библиотек для обработки данных, основная структура данных, используемая в Polars, - это DataFrame. DataFrame - это двумерная структура данных, состоящая из строк и столбцов. Столбцы DataFrame состоят из Series, которые являются одномерными массивами с метками.

Для примера создадим датафрейм в Polars из словаря на основе случайно сгенерированных данными с информацией о домах:

import numpy as np
import polars as pl

num_rows = 5000
rng = np.random.default_rng(seed=7)

buildings_data = {
     "sqft": rng.exponential(scale=1000, size=num_rows),
     "year": rng.integers(low=1995, high=2023, size=num_rows),
     "building_type": rng.choice(["A", "B", "C"], size=num_rows),
 }
buildings = pl.DataFrame(buildings_data)
buildings

Тут сначала импортируются библиотеки numpy и polars под псевдонимами np и pl соответственно. Затем определяется переменная num_rows, которая определяет количество строк в случайно сгенерированных данных. Для генерации случайных значений вызывается функция default_rng() из модуля random библиотеки NumPy. Эта функция возвращает генератор, который может создавать случайные числа по различным законам распределения.

Затем определяется словарь с ключами sqft, year и building_type, которые представляют собой случайно сгенерированные массивы длиной num_rows. Массив sqft содержит числа с плавающей точкой, year содержит целые числа, а массив building_type содержит строки. Эти данные становятся тремя столбцами в датафрейме.

А конструктор класса DataFrame в Polars принимает двумерные данные различных форм, в данном случае словарь. При отображении датафрейма в консоли Polars сначала выводит размерность данных в виде кортежа, где первый элемент показывает количество строк, а второй - количество столбцов. Затем выводится предварительный просмотр данных в виде таблицы с именами столбцов и их типы данных. Например, “year” имеет тип float64, а “building_type“ - str. Polars поддерживает и другие разные типы данных, в основном основанные на реализации из Arrow.

У датафреймов в Polars есть много полезных методов и атрибутов для исследования данных. При этом в pandas и в Polars используется в основном те же их названия.

Обращение buildings.schema возвращает схему DataFrame в виде словаря с типом данных каждого столбца. Есть и привычные по pandas методы head и describe.

Контексты и выражения в Polars

Контексты и выражения - основные элементы синтаксиса преобразования данных в Polars. Выражения - это по сути вычисления или преобразованиями над столбцами данных. В т.ч. это математические операции, агрегации, сравнения, манипуляции строками и многое другое.

Контекст - это основное действие в рамках которого исполняется выражение. В Polars есть три основных контекста:

• Выборка (selection) столбцов из датафрейма.
• Фильтрация (filtering) - уменьшение размера DataFrame путем извлечения строк, которые соответствуют определенным условиям.
• Группировка/агрегация (groupby/aggregation) - вычисление сводных статистик внутри подгрупп данных.

Рассмотрим пример выборки:

# Выбор только столбца "sqft" из DataFrame buildings
buildings.select("sqft")  

 # То же самое, но с использованием метода col() для явного указания столбца
buildings.select(pl.col("sqft")) 

Используя выражение pl.col() в контексте .select(), можно выполнять дополнительные манипуляции со столбцами:

# Сортируем значения столбца "sqft" по возрастанию и затем делим их на 1000
buildings.select(pl.col("sqft").sort() / 1000)

Ещё один контекст, который часто используется совместно с .select(), это .filter(). Как следует из названия, эта операция фильтрует данные на основе заданного выражения. Например, если мы хотим оставить только все дома, построенные после 2015 года, то это можно сделать так:

 # Фильтруем данные, оставляя только дома, построенные после 2015 года
after_2015 = buildings.filter(pl.col("year") > 2015)

# Количество строк, количество столбцов DataFrame
after_2015.shape  
# Результат: (1230, 3)

# Выбираем минимальное значение из столбца "year" DataFrame
after_2015.select(pl.col("year").min())  

Еще один часто используемый контекст - это groupby, он же агрегация. Агрегация используется для вычисления сводных статистик внутри подгрупп набора данных. В примере с данными о зданиях предположим, что мы хотим узнать среднюю площадь, медианный год постройки и количество зданий для каждого типа здания. Следующий запрос выполняет эту задачу:

# Вернём DataFrame, сгруппированный по типу здания и содержащий 
# среднюю площадь, медианный год и количество зданий для каждого типа

buildings.groupby("building_type").agg(
     [
         pl.mean("sqft").alias("mean_sqft"),  # Средняя площадь
         pl.median("year").alias("median_year"),  # Медианный год
         pl.count(),  # Количество зданий
     ]
)

В этом примере сначала вызываем buildings.groupby("building_type"), чтобы создать объект GroupBy. У GroupBy есть метод .agg(), который принимает список выражений, вычисляемых для каждой группы. Например, pl.mean("sqft") вычисляет среднюю площадь для каждого типа здания, а pl.count() возвращает количество зданий каждого типа. А .alias() используется, чтобы дать имя новым агрегированным столбцам.

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

Работа с “ленивым” API

API с отложенной загрузкой (Lazy API) или просто “ленивый” API в Polars является одной из самых мощных особенностей библиотеки. Он позволяет указать последовательность операций без их немедленного выполнения: эти операции сохраняются как вычислительный граф и выполняются только при необходимости. Это позволяет Polars оптимизировать запросы перед выполнением, обнаруживать ошибки до обработки данных и выполнять запросы, эффективно используя память при обработки больших наборах данных, которые не помещаются в память.

Основным объектом в “ленивом” API является LazyFrame. LazyFrame можно создать следующим способом:

import numpy as np
import polars as pl

num_rows = 5000
rng = np.random.default_rng(seed=7)

buildings = {
    "sqft": rng.exponential(scale=1000, size=num_rows),
    "price": rng.exponential(scale=100_000, size=num_rows), #Добавим новую колонку "price"
    "year": rng.integers(low=1995, high=2023, size=num_rows),
    "building_type": rng.choice(["A", "B", "C"], size=num_rows),
}

# Создание объекта LazyFrame
buildings_lazy = pl.LazyFrame(buildings)
buildings_lazy

# Вывод LazyFrame
<polars.LazyFrame object at 0x106D19950>

Посмотрим, как работает “ленивый” API:

# Создание отложенного запроса для DataFrame buildings_lazy

lazy_query = (
    buildings_lazy
    .with_columns(  # Добавление нового столбца с ценой за квадратный фут
        (pl.col("price") / pl.col("sqft")).alias("price_per_sqft")  # Вычисление цены за квадратный фут
    )
    .filter(pl.col("price_per_sqft") > 100)  # Фильтрация по цене за квадратный фут
    .filter(pl.col("year") < 2010)  # Фильтрация по году
)
lazy_query  

Перед выполнением запроса можно посмотреть его  план с помощью метода .show_graph(). План запроса выводит последовательность шагов, которые вызовет запрос.

lazy_query.show_graph()

План запроса считывается снизу вверх, и каждый прямоугольник соответствует этапу в запросе. Символы Sigma (σ) и Pi (π) - это символы из реляционной алгебры, которые сообщают операцию, которая выполняется с данными.

В этом примере π */4 означает, что работа осуществляется над всеми четырьмя столбцами DataFrame; а σ(col(“year”)) < 2010 означает, что обрабатываются только строки, в которых инофрмация до 2010 года. Можно интерпретировать полный план запроса так:

• 1. Используй четыре столбца из buildings_lazy и отфильтруй buildings_lazy, оставив строки, в которых хранится инофрмация до 2010 года;.
• 2. Создай столбец price_per_sqft;
• 3. Отфильтруй buildings_lazy, оставив только строки, в которых price_per_sqft больше 100.

Важно отметить, что Polars фильтрует buildings_lazy по году перед выполнением любой другой части запроса, несмотря на то, что это последняя операция, указанная в коде. Такого рода оптимизация в Polars называется "выталкивание предикатов" (predicate pushdown), она повышает эффективность использования памяти за счет применения фильтров как можно раньше и последующего уменьшения размера данных перед дальнейшей обработкой.

Иногда граф плана запроса  может обрезать важные детали, которые не помещаются в прямоугольник. Это особенно часто случается, когда запросы становятся более сложными. Если необходимо увидеть полное представление плана запроса, то можно использовать метод .explain(). Как и с графическим планом запроса, строковый план запроса читается снизу вверх, и каждый этап находится на своей собственной строке.

print(lazy_query.explain())

# Вывод:
FILTER [(col("price_per_sqft")) > (100.0)] 
  FROM WITH_COLUMNS: [[(col("price")) / (col("sqft"))].alias("price_per_sqft")] 
  DF ["sqft", "price", "year", "building_type"];
⮑ PROJECT */4 COLUMNS; SELECTION: "[(col(\"year\")) < (2010)]"

Далее для фактического выполнения запроса используется метод .collect():

# Создание “ленивого” запроса, который вычисляет цену и фильтрует данные
lazy_query = buildings_lazy.with_columns(
    (pl.col("price") / pl.col("sqft")).alias("price_per_sqft")).filter(pl.col("price_per_sqft") > 100) .filter(pl.col("year") < 2010
)

# Выполнение “ленивого” запроса и выбор определенных столбцов
result = lazy_query.collect().select(pl.col(["price_per_sqft", "year"]))

Чтобы дополнительно проверить, что запрос правильно отфильтровал данные, можно посмотреть на сводную статистику:

lazy_query.collect().select(pl.col(["price_per_sqft", "year"])).describe()

Сканирование данных с помощью LazyFrames

Если весь набор данных уже хранится в памяти, то зачем нужны отложенные запросы?

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

При работе с файлами, такими как CSV, обычно все данные загружаются в память и далее с ними происходит работа. С помощью “ленивого” API можно минимизировать объём считываемых в память данных, обрабатывая необходимые. Это позволяет Polars оптимизировать использование памяти и время вычислений.

В следующем примере поработаем с данными об электромобилях из датасета Data.gov. Датасет содержит 140 000 строк данных о зарегистрированных в Департаменте лицензирования штата Вашингтон электрических и гибридных автомобилях.

import requests 
import pathlib 

# Функция для загрузки данных
def download_file(file_url: str, local_file_path: pathlib.Path) -> None:
    """Загружаем данные и сохраниям их в указаном файле."""
    response = requests.get(file_url)  
        local_file_path.write_bytes(response.content)  
        print(f"Данные успешно загружены и сохранены в: {local_file_path}")
    else:  
        raise requests.exceptions.RequestException(
            f"Не удалось загрузить данные. Статус: {response.status_code}"
        )


from downloads import download_file 

# URL для загрузки данных
url = "https://data.wa.gov/api/views/f6w7-q2d2/rows.csv?accessType=DOWNLOAD"  

# Путь к локальному файлу
local_file_path = pathlib.Path("electric_cars.csv")  

# Вызов функции для загрузки файла
download_file(url, local_file_path)  

Ключевой момент для эффективной работы с файлами через ленивый API заключается в использовании функциональности сканирования Polars. Вместо чтения всего файла в память Polars создает LazyFrame, который ссылается на данные файла. Обработка данных не происходит до тех пор, пока вы явно не выполните запрос. Для сканирования данных используется pl.scan_csv():

# Сканирование CSV файла и создание LazyFrame
lazy_car_data = pl.scan_csv(local_file_path)
lazy_car_data

# Вывод:
<polars.LazyFrame object at 0x10292EC50>


# Cхема данных LazyFrame
lazy_car_data.schema

# Вывод:
{'VIN (1-10)': Utf8, 'County': Utf8, 'City': Utf8, 'State': Utf8,
'Postal Code': Int64, 'Model Year': Int64, 'Make': Utf8, 'Model': Utf8,
'Electric Vehicle Type': Utf8, 'Clean Alternative Fuel Vehicle (CAFV) Eligibility': Utf8,
'Electric Range': Int64, 'Base MSRP': Int64, 'Legislative District': Int64,
'DOL Vehicle ID': Int64, 'Vehicle Location': Utf8, 'Electric Utility': Utf8,
'2020 Census Tract': Int64}

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

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

# Создание отложенного запроса для фильтрации и агрегации данных 
# по автомобилям на основе lazy_car_data
lazy_car_query = (
    lazy_car_data
    .filter((pl.col("Model Year") >= 2018))  # Фильтрация по году выпуска
    .filter(
        pl.col("Electric Vehicle Type") == "Battery Electric Vehicle (BEV)"  # Фильтрация по типу автомобиля
    )
    .groupby(["State", "Make"])  # Группировка по штату и производителю
    .agg(
        pl.mean("Electric Range").alias("Average Electric Range"),  # Вычисление среднего запаса хода электромобиля
        pl.min("Model Year").alias("Oldest Model Year"),  # Вычисление самого старого года выпуска
        pl.count().alias("Number of Cars"),  # Подсчет количества автомобилей
    )
    .filter(pl.col("Average Electric Range") > 0)  # Фильтрация по среднему запасу хода электромобиля
    .filter(pl.col("Number of Cars") > 5)  # Фильтрация по количеству автомобилей
    .sort(pl.col("Number of Cars"), descending=True)  # Сортировка по количеству автомобилей
)

# Выполнение "ленивого" запроса и вывод результата
lazy_car_query.collect()

В данном запросе фильтруются данные по всем автомобилям, у которых год выпуска 2018 или позднее, и тип транспорта - Battery Electric Vehicle (BEV). Затем вычисляется средний запас хода электромобиля, самый старый год выпуска и количество автомобилей для каждого штата и марки. Наконец, дополнительно данные фильтруются по среднему запасу хода и количеству автомобилей.

Поскольку это "ленивый" запрос, никаких вычислений не выполняется до тех пор, пока вы не вызовете lazy_car_query.collect(). После выполнения запроса хранится и возвращается только запрошенные данные - ничего более.

В каждой строке в DataFrame, возвращенном из lazy_car_query.collect(), хранится информация о среднем запасе хода электромобиля, самом старом годе выпуска и количестве автомобилей по каждому штату и производителю. Например, в первой строке указано, что в штате Вашингтон есть 55 690 автомобилей Tesla с годом выпуска 2018 или позднее, и их средний запас хода составляет около 89,11 миль.

Безшовная интеграция

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

Интеграция с внешними источниками данных

Ранее мы рассмотрели пример, как Polars выполняет отложенные запросы при работе с CSV файлами при помощи метода .scan_csv(). Polars также может работать с источниками данных, такими как JSON, Parquet, Avro, Excel и различными базами данных.

Для сохранения данных в разные форматы файлов используются методы .write_csv(),  .write_ndjson(),  .write_parquet():

# Пример сохранения данных в разные форматы
import polars as pl

# Создание DataFrame
data = pl.DataFrame({
    "A": [1, 2, 3, 4, 5],
    "B": [6, 7, 8, 9, 10],
})

# Запись данных в формат CSV
data.write_csv("data.csv")

# Запись данных в формат JSON
data.write_ndjson("data.json")

# Запись данных в формат Parquet
data.write_parquet("data.parquet")

Для чтения данных из разных форматов файлов используются методы .read_csv(),  .read_ndjson(),  .read_parquet():

# Чтение данных из CSV и создание отложенного запроса
data_csv = pl.read_csv("data.csv")
data_csv_lazy = pl.scan_csv("data.csv")

# Отображение схемы данных CSV
data_csv_lazy.schema
{'A': Int64, 'B': Int64}

# Чтение данных из JSON и создание отложенного запроса
data_json = pl.read_ndjson("data.json")
data_json_lazy = pl.scan_ndjson("data.json")

# Отображение схемы данных JSON
data_json_lazy.schema
{'A': Int64, 'B': Int64}

# Чтение данных из Parquet и создание отложенного запроса
data_parquet = pl.read_parquet("data.parquet")
data_parquet_lazy = pl.scan_parquet("data.parquet")

# Отображение схемы данных Parquet
data_parquet_lazy.schema
# {'A': Int64, 'B': Int64}

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

Интеграция с экосистемой Python

Polars без проблем интегрируется с существующими библиотеками Python, что позволяет внедрять Polars в существующий код, не изменяя зависимостей или сильно перепроекти́руя код. В следующем примере рассмотрим, как датафреймы Polars без проблем преобразуются между массивами NumPy и датафреймами pandas:

import numpy as np
import pandas as pd
import polars as pl

# Создание датафрейма с помощью Polars
polars_data = pl.DataFrame({
    "A": [1, 2, 3, 4, 5],
    "B": [6, 7, 8, 9, 10]
})

# Создание датафрейма с помощью Pandas
pandas_data = pd.DataFrame({
    "A": [1, 2, 3, 4, 5],
    "B": [6, 7, 8, 9, 10]
})

# Создание массива NumPy
numpy_data = np.array([
    [1, 6],
    [2, 7],
    [3, 8],
    [4, 9],
    [5, 10]
])

# Транспонирование массива NumPy
numpy_data = numpy_data.T

pandas-датафрейм и массива NumPy можно преобразовать в Polars-датафрейм с помощью pl.from_pandas(pandas_data), pl.from_numpy(numpy_data, schema={"A": pl.Int64, "B": pl.Int64}):

pl.from_pandas(pandas_data)

pl.from_numpy(numpy_data, schema={"A": pl.Int64, "B": pl.Int64})

Для преобразования DataFrame из Polars обратно в pandas или NumPy используются методы  .to_pandas(),  .to_numpy():

# Преобразование объекта DataFrame из Polars в DataFrame из pandas
polars_data.to_pandas()

# Вывод:
   A   B
0  1   6
1  2   7
2  3   8
3  4   9
4  5  10

# Преобразование объекта DataFrame из Polars в массив NumPy
polars_data.to_numpy()

# Вывод:
array([[ 1,  6],
       [ 2,  7],
       [ 3,  8],
       [ 4,  9],
       [ 5, 10]])

Источник: Real Python

Перевод и адаптация: Дмитрий Каленов