Compute Pairwise Group Metrics in Polars: Elementwise Alignment with In-Group Index, Join, and Aggregate

Learn a vectorized Polars pattern for pairwise group metrics: create an in-group index, join to align, and aggregate (e.g., dot product). Faster than loops.

Computing a single metric for every pair of groups is a common pattern in data analysis, but the straightforward nested loop over groups leaves performance on the table. If you work in Polars, there is a vectorized way to express pairwise operations that scales better and stays within the lazy/expression model. Below we walk through a concrete example with a dot product and show how to generalize the idea to other pairwise computations.

Problem setup

We start with a DataFrame that has a group identifier and a numeric series. The task is to produce a G × G matrix of pairwise metrics between all groups.

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)]),
    }
)

A direct approach computes the dot product of the "data" column for every pair of groups using a nested loop.

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

This works, but it serializes the computation and repeatedly filters the data. The question is how to move this into a single Polars expression that can exploit its parallel engine.

What’s really going on

“For each pair of groups” is a cartesian product problem. In Polars, a cartesian product is just a cross join. For example, the set of all pairs of group IDs looks like this:

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

However, for a dot product you don’t want to join every row with every other row. You want elementwise alignment: row 0 of group A with row 0 of group B, row 1 with row 1, and so on. If each group has N rows, then there should be N matches for each pair of groups, not N × N. That means we need a stable per-row index within each group and we need to join on that index, not on the raw records.

The Polars way: index within group + join + aggregation

The solution is to assign an index inside each group, use that index as the join key to align rows elementwise across groups, and then aggregate by the pair of groups using a dot product. This avoids the N × N blow-up and expresses the entire computation as one pipeline.

# assign an index per group
tbl_idx = tbl.with_columns(row_idx=pl.int_range(pl.len()).over("group_id"))

# pairwise elementwise alignment and aggregation
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"))
)

Conceptually, the join creates the cartesian pairs of groups but only matches rows that share the same within-group position. The final group_by rolls those aligned rows up to one value per pair using the dot expression.

Why this matters

Moving pairwise operations into expressions unlocks Polars’ parallel execution and avoids Python-level loops and repeated filters. The index-within-group trick prevents accidental cartesian explosions by ensuring a one-to-one row alignment. The same pattern is applicable whenever you need a G × G matrix from groupwise vectors and the operation is defined elementwise before reduction, such as dot products.

If your real target is a geometry-aware metric like Frechet distance, note that the polars-st plugin exposes a frechet_distance expression. You can adapt the same pairing pattern while using that expression where the dot product appears in the example.

Takeaways

Think of “all pairs” as a join. Decide whether you want a full cartesian product of rows or an elementwise alignment; if it’s the latter, create a deterministic within-group index and join on it. Aggregate on top of the joined pairs to get one value per group pair. This keeps the computation in Polars, reduces overhead from Python loops, and yields a clean, maintainable pipeline.