[Python] Is there a polars operation to apply a function over each pair of groups?

I have a polars data frame which could be generated like so:

import polars as pl
import numpy as np

num_points = 10
group_count = 3

df = pl.DataFrame(
{
"group_id": np.concatenate([*num_points for i in range(group_count)]),
"data": np.concatenate([np.random.rand(num_points) for _ in range(group_count)]),
}
)


and I want to apply some operation that takes in 2 groups, and outputs a single number, over all the groups. This operation is not a polars native expression, and is shapely.frechet_distance , but for a good starting point I would like to see if that is possible to do fast using a polars expression, like taking the dot product.

An easy solution is to loop over all the groups like so:

def distance_between_groups(df: pl.DataFrame):
"""
for each pair of groups g1, g2, returns some aggregate function of their data,
such that we get a G x G matrix of the resulting aggregations.
"""
groups = df['group_id'].unique(maintain_order=True)
res = np.zeros((df['group_id'].n_unique(), df['group_id'].n_unique()))
for i, g1 in enumerate(groups):
g1_df = df.filter(pl.col('group_id') == g1)
for j, g2 in enumerate(groups):
g2_df = df.filter(pl.col('group_id') == g2)
res[j] = (g1_df['data'] * g2_df['data']).sum()

return res



but I am looking for something more efficient, that uses the parallelism of polars.

Continue reading...