Как считать попарные метрики в Polars без циклов: индекс в группе, join и dot

Показываем векторизованный подход в Polars для G×G попарных метрик: индекс внутри группы, join по позициям и агрегация dot. Быстрее циклов и масштабируется

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

Постановка задачи

Начинаем с DataFrame, где есть идентификатор группы и числовой ряд. Нужно получить матрицу G × G попарных метрик между всеми группами.

import polars as pl
import numpy as np

points_per_group = 10
num_groups = 3

tbl = pl.DataFrame(
    {
        "group_id": np.concatenate([[g] * points_per_group for g in range(num_groups)]),
        "data": np.concatenate([np.random.rand(points_per_group) for _ in range(num_groups)]),
    }
)

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

def pairwise_group_metric(frame: pl.DataFrame):
    ids = frame["group_id"].unique(maintain_order=True)
    out = np.zeros((frame["group_id"].n_unique(), frame["group_id"].n_unique()))
    for i, a in enumerate(ids):
        part_a = frame.filter(pl.col("group_id") == a)
        for j, b in enumerate(ids):
            part_b = frame.filter(pl.col("group_id") == b)
            out[i][j] = (part_a["data"] * part_b["data"]).sum()
    return out

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

Что здесь на самом деле происходит

«Для каждой пары групп» — это задача декартова произведения. В Polars декартово произведение — это просто cross join. Например, множество всех пар идентификаторов групп выглядит так:

pairs = pl.DataFrame({"group_id": range(3)})
pairs.join(pairs, how="cross")

Однако для скалярного произведения не нужно соединять каждую строку с каждой. Нужна покомпонентная стыковка: строка 0 группы A со строкой 0 группы B, строка 1 со строкой 1 и т. д. Если в каждой группе N строк, то для каждой пары групп должно быть N сопоставлений, а не N × N. Значит, нам нужен стабильный построчный индекс внутри группы и соединять нужно по этому индексу, а не по исходным записям.

Подход Polars: индекс внутри группы + join + агрегация

Решение: назначить индекс внутри каждой группы, использовать его как ключ соединения для покомпонентного выравнивания строк между группами, а затем агрегировать по паре групп с помощью скалярного произведения. Это предотвращает взрыв N × N и выражает всё вычисление одним конвейером.

# назначаем индекс внутри каждой группы
tbl_idx = tbl.with_columns(row_idx=pl.int_range(pl.len()).over("group_id"))

# попарное покомпонентное выравнивание и агрегация
result = (
    tbl_idx.join(tbl_idx, on="row_idx")
    .group_by("group_id", "group_id_right", maintain_order=True)
    .agg(pl.col("data").dot("data_right"))
)

Идея такая: join формирует декартовы пары групп, но сопоставляет только строки с одинаковой позицией внутри группы. Финальный group_by сворачивает выровненные строки в одно значение на пару групп через выражение dot.

Почему это важно

Перенос попарных операций в выражения раскрывает параллельное исполнение Polars и избавляет от циклов на уровне Python и повторных фильтраций. Приём с индексом внутри группы предотвращает случайные декартовы взрывы, гарантируя выравнивание «один к одному». Тот же шаблон применим всякий раз, когда нужна матрица G × G из групповых векторов и операция задаётся покомпонентно до редукции — как в случае со скалярным произведением.

Если вам нужна метрика с учётом геометрии, например расстояние Фреше, обратите внимание: плагин polars-st предоставляет выражение frechet_distance. Можно использовать тот же шаблон попарного сопоставления, подставив это выражение вместо скалярного произведения.

Выводы

Смотрите на «все пары» как на join. Решите, нужен ли полный декартов продукт строк или покомпонентное выравнивание; если второе — создайте детерминированный индекс внутри группы и соединяйте по нему. Затем агрегируйте поверх соединённых пар, чтобы получить одно значение на пару групп. Так вычисления остаются в Polars, уменьшаются накладные расходы циклов Python, а конвейер получается чистым и удобным в сопровождении.