API Python

Глобальный SQL

Для выполнения SQL-запросов с помощью движка выполнения Polars к данным Polars типа DataFrame, LazyFrame и Series можно использовать как SQLContext, так и функцию polars.sql(). Также поддерживаются DataFrame и Series из Pandas, а также объекты Table и RecordBatch из PyArrow.

При использовании в SQL-запросе объекты, отличные от Polars, неявно преобразуются в DataFrame. Для PyArrow и данных Pandas с типами данных PyArrow это преобразование выполняется без копирования, если базовые данные корректно соответствуют встроенному типу данных.

Пример:

import polars as pl
import pandas as pd

polars_df = pl.DataFrame({"a": [1, 2, 3, 4], "b": [4, 5, 6, 7]})
pandas_df = pd.DataFrame({"a": [3, 4, 5, 6], "b": [6, 7, 8, 9]})
polars_series = (polars_df["a"] * 2).rename("c")
pyarrow_table = polars_df.to_arrow()

pl.sql(
    """
    SELECT a, b, SUM(c) AS c_total FROM (
      SELECT * FROM polars_df                  -- polars frame
        UNION ALL SELECT * FROM pandas_df      -- pandas frame
        UNION ALL SELECT * FROM pyarrow_table  -- pyarrow table
    ) all_data
    INNER JOIN polars_series
      ON polars_series.c = all_data.b          -- polars series
    GROUP BY "a", "b"
    ORDER BY "a", "b"
    """
).collect()

# shape: (3, 3)
# ┌─────┬─────┬─────────┐
# │ a   ┆ b   ┆ c_total │
# │ --- ┆ --- ┆ ---     │
# │ i64 ┆ i64 ┆ i64     │
# ╞═════╪═════╪═════════╡
# │ 1   ┆ 4   ┆ 8       │
# │ 3   ┆ 6   ┆ 18      │
# │ 5   ┆ 8   ┆ 8       │
# └─────┴─────┴─────────┘

SQL для фреймов

Выполняет SQL-запрос непосредственно к указанному eager/lazy-фрейму, ссылаясь на него как на «self»; возвращает новый фрейм того же типа, содержащий результат запроса.

Пример:

import polars as pl

df = pl.DataFrame({
    "a": [1, 2, 3],
    "b": [4, 5, 6],
})
df.sql("""
  SELECT a::uint4, (b * b) AS bb
  FROM self WHERE a != 2
""")

# shape: (2, 2)
# ┌─────┬─────┐
# │ a   ┆ bb  │
# │ --- ┆ --- │
# │ u32 ┆ i64 │
# ╞═════╪═════╡
# │ 1   ┆ 16  │
# │ 3   ┆ 36  │
# └─────┴─────┘

SQL для выражений

Функцию polars.sql_expr() можно использовать для создания встроенных выражений Polars из фрагментов SQL.

Пример:

import polars as pl

df = pl.DataFrame({
    "a": [1, 2, 3],
    "b": [4, 5, 6],
})
df.with_columns(
    pl.sql_expr("(a * a) + (b::float / 2) AS expr1"),
    pl.sql_expr("CONCAT_WS(':',a,b) AS expr2")
)

# shape: (3, 4)
# ┌─────┬─────┬───────┬───────┐
# │ a   ┆ b   ┆ expr1 ┆ expr2 │
# │ --- ┆ --- ┆ ---   ┆ ---   │
# │ i64 ┆ i64 ┆ f64   ┆ str   │
# ╞═════╪═════╪═══════╪═══════╡
# │ 1   ┆ 4   ┆ 3.0   ┆ 1:4   │
# │ 2   ┆ 5   ┆ 6.5   ┆ 2:5   │
# │ 3   ┆ 6   ┆ 12.0  ┆ 3:6   │
# └─────┴─────┴───────┴───────┘

SQLContext

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

class polars.SQLContext

Выполняет SQL-запросы к данным DataFrame/LazyFrame.

__init__(
    frames: Mapping[str,
    CompatibleFrameType | None] | None = None,
    *,
    register_globals: bool | int = False,
    eager: bool = False,
    **named_frames: CompatibleFrameType | None,
) → None

Инициализирует новый SQLContext.

Изменено в версии 0.20.31: Параметр eager_execution переименован в eager.

Параметры:

frames

Сопоставление {name:frame, ...}, которое может содержать фреймы Polars, а также DataFrame и Series из pandas и объекты Table и RecordBatch из pyarrow.

register_globals

Регистрирует совместимые объекты (DataFrame, LazyFrame и Series из polars), найденные в глобальной области видимости, автоматически сопоставляя имена переменных с именами таблиц. Чтобы зарегистрировать другие объекты (данные pandas/pyarrow), передайте их явно или вызовите метод класса execute_global. Если задано целое число, будут зарегистрированы только последние «n» найденных объектов.

eager

Если значение True, возвращает результаты выполнения в виде DataFrame вместо LazyFrame. (Обратите внимание: сам запрос всегда выполняется в ленивом режиме; этот параметр определяет, вернёт ли execute() eager- или lazy-фрейм с результатом.)

**named_frames

Именованные eager/lazy-фреймы, переданные как именованные аргументы.

Примеры

>>> lf = pl.LazyFrame({"a": [1, 2, 3], "b": ["x", None, "z"]})
>>> res = pl.SQLContext(frame=lf).execute(
...     "SELECT b, a*2 AS two_a FROM frame WHERE b IS NOT NULL"
... )
>>> res.collect()
shape: (2, 2)
┌─────┬───────┐
│ b   ┆ two_a │
│ --- ┆ ---   │
│ str ┆ i64   │
╞═════╪═══════╡
│ x   ┆ 2     │
│ z   ┆ 6     │
└─────┴───────┘

Примечание: также можно использовать как менеджер контекста.

__enter__() → SQLContext[FrameType]

Отслеживает зарегистрированные таблицы при входе в область видимости; поддерживает вложенные области.

__exit__(
    exc_type: type[BaseException] | None,
    exc_val: BaseException | None,
    exc_tb: TracebackType | None,
) → None

При выходе из контекста отменяет регистрацию всех таблиц, созданных в заданной области видимости.

См. также

unregister

import polars as pl

df1 = pl.DataFrame({"id": [1, 2, 3], "value": [0.1, 0.2, 0.3]})
df2 = pl.DataFrame({"id": [3, 2, 1], "value": [25.6, 53.4, 12.7]})

with pl.SQLContext(df_a=df1, df_b=df2, eager=True) as ctx:
    df = ctx.execute("""
      SELECT
        a.id,
        a.value AS value_a,
        b.value AS value_b
      FROM df_a AS a INNER JOIN df_b AS b USING (id)
      ORDER BY id
    """)

    # shape: (3, 3)
    # ┌─────┬─────────┬─────────┐
    # │ id  ┆ value_a ┆ value_b │
    # │ --- ┆ ---     ┆ ---     │
    # │ i64 ┆ f64     ┆ f64     │
    # ╞═════╪═════════╪═════════╡
    # │ 1   ┆ 0.1     ┆ 25.6    │
    # │ 2   ┆ 0.2     ┆ 53.4    │
    # │ 3   ┆ 0.3     ┆ 12.7    │
    # └─────┴─────────┴─────────┘