[Python] Polars - Perform matrix inner product on lazy frames to produce sparse representation...

Suppose we have a polars dataframe like:

df = pl.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]}).lazy()

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


I would like to X^TX the matrix while preserving the sparse matrix format for arrow* - in pandas I would do something like:

pdf = df.collect().to_pandas()
numbers = pdf[["a", "b"]]
(numbers.T @ numbers).melt(ignore_index=False)

variable value
a a 14
b a 26
a b 26
b b 50


I did something like this in polars:

df.select(
(pl.col("a") * pl.col("a")).sum().alias("aa"),
(pl.col("a") * pl.col("b")).sum().alias("ab"),
(pl.col("b") * pl.col("a")).sum().alias("ba"),
(pl.col("b") * pl.col("b")).sum().alias("bb"),
).unpivot().collect()

shape: (4, 2)
┌──────────┬───────┐
│ variable ┆ value │
│ --- ┆ --- │
│ str ┆ i64 │
╞══════════╪═══════╡
│ aa ┆ 14 │
│ ab ┆ 26 │
│ ba ┆ 26 │
│ bb ┆ 50 │
└──────────┴───────┘


Which is almost there but not quite. This is a hack to get around the fact that I can't store lists as the column names (and then I could unnest them to become two different columns representing the x and y axis of the matrix). Is there a way to get the same format as shown in the pandas example?

*arrow is a columnar data format which means it's performant when scaled across rows but not across columns, which is why I think the sparse matrix representation is better if I want to use the results of the gram matrix chained with pl.LazyFrames later down the graph. I could be wrong though!

Continue reading...