[Python] Find pairs of keys for rows that have at least one property in common

I'm using polars with a data frame whose schema looks like this:

Schema({'CustomerID': String, 'StockCode': String, 'Total': Int64})


interpreted as "Customer CustomerID bought Total of product StockCode." I'm looking for an efficient way to generate all unique pairs of customer IDs such that the two customers purchased at least one of the same product. So, given:

import polars as pl

df = pl.from_repr("""
┌─────────────┬────────────┬───────┐
│ CustomerID ┆ StockCode ┆ Total │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ f64 │
╞═════════════╪════════════╪═══════╡
│ A ┆ 123 ┆ 45.78 │
│ A ┆ 140 ┆ 10.26 │
│ B ┆ 125 ┆ 99.62 │
│ B ┆ 128 ┆ 23.65 │
│ B ┆ 140 ┆ 92.95 │
│ C ┆ 123 ┆ 45.78 │
│ D ┆ 145 ┆ 7.58 │
└─────────────┴────────────┴───────┘
""")


the algorithm should produce:

[
( 'A', 'B' ), # Because of product 140
( 'A', 'C' ) # Because of product 123
]


I'll then use this to compute a pairwise similarity measure (e.g. dot product or cosine) and then, as efficiently as possible, produce a data frame that looks like:

Schema({'ID1': String, 'ID2': String, 'Similarity': Float64})


likely with a threshold for minimum similarity. Doing something like:

from itertools import combinations

row_keys = pl.Series(
df.select(category_col)
.unique(
).drop_nulls(
).collect()
).to_list()

pairs = combinations(row_keys, 2)


generates all possible pairs, but is hugely inefficient.

Continue reading...