本文档详细介绍了如何使用 Polars 库计算 DataFrame 中各列之间的余弦相似度，并将其以相关矩阵的形式呈现。通过 join_where 方法生成列组合，利用 Polars 表达式计算余弦相似度，最后使用 pivot 方法将结果转换为矩阵形式，方便进行数据分析和挖掘。
前提条件
确保你已经安装了 Polars 库。可以使用 pip 进行安装：
pip install polars
数据准备
首先，我们创建一个 Polars DataFrame，其中包含字符串列 col1 和列表列 col2。col2 列包含数值列表，我们将基于这些列表计算余弦相似度。
import polars as pl
from numpy.linalg import norm
data = {
"col1": ["a", "b", "c", "d"],
"col2": [[-0.06066, 0.072485, 0.548874, 0.158507],
[-0.536674, 0.10478, 0.926022, -0.083722],
[-0.21311, -0.030623, 0.300583, 0.261814],
[-0.308025, 0.006694, 0.176335, 0.533835]],
}
df = pl.DataFrame(data)
print(df)输出：
shape: (4, 2) ┌──────┬─────────────────────────────────┐ │ col1 ┆ col2 │ │ --- ┆ --- │ │ str ┆ list[f64] │ ╞══════╪═════════════════════════════════╡ │ a ┆ [-0.06066, 0.072485, … 0.15850… │ │ b ┆ [-0.536674, 0.10478, … -0.0837… │ │ c ┆ [-0.21311, -0.030623, … 0.2618… │ │ d ┆ [-0.308025, 0.006694, … 0.5338… │ └──────┴─────────────────────────────────┘
生成列组合
为了计算每对列之间的余弦相似度，我们需要生成所有可能的列组合。我们可以使用 join_where 方法来实现这一点。首先，添加一个行索引，然后使用 join_where 将 DataFrame 与自身连接，条件是左侧的索引小于等于右侧的索引，以避免重复计算。
df = df.with_row_index().lazy()
combinations_df = df.join_where(df, pl.col("index") <= pl.col("index_right")).collect()
print(combinations_df)输出：
shape: (10, 6) ┌───────┬──────┬─────────────────────────────────┬─────────────┬────────────┬─────────────────────────────────┐ │ index ┆ col1 ┆ col2 ┆ index_right ┆ col1_right ┆ col2_right │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ u32 ┆ str ┆ list[f64] ┆ u32 ┆ str ┆ list[f64] │ ╞═══════╪══════╪═════════════════════════════════╪═════════════╪════════════╪═════════════════════════════════╡ │ 0 ┆ a ┆ [-0.06066, 0.072485, … 0.15850… ┆ 0 ┆ a ┆ [-0.06066, 0.072485, … 0.15850… │ │ 0 ┆ a ┆ [-0.06066, 0.072485, … 0.15850… ┆ 1 ┆ b ┆ [-0.536674, 0.10478, … -0.0837… │ │ 0 ┆ a ┆ [-0.06066, 0.072485, … 0.15850… ┆ 2 ┆ c ┆ [-0.21311, -0.030623, … 0.2618… │ │ 0 ┆ a ┆ [-0.06066, 0.072485, … 0.15850… ┆ 3 ┆ d ┆ [-0.308025, 0.006694, … 0.5338… │ │ 1 ┆ b ┆ [-0.536674, 0.10478, … -0.0837… ┆ 1 ┆ b ┆ [-0.536674, 0.10478, … -0.0837… │ │ 1 ┆ b ┆ [-0.536674, 0.10478, … -0.0837… ┆ 2 ┆ c ┆ [-0.21311, -0.030623, … 0.2618… │ │ 1 ┆ b ┆ [-0.536674, 0.10478, … -0.0837… ┆ 3 ┆ d ┆ [-0.308025, 0.006694, … 0.5338… │ │ 2 ┆ c ┆ [-0.21311, -0.030623, … 0.2618… ┆ 2 ┆ c ┆ [-0.21311, -0.030623, … 0.2618… │ │ 2 ┆ c ┆ [-0.21311, -0.030623, … 0.2618… ┆ 3 ┆ d ┆ [-0.308025, 0.006694, … 0.5338… │ │ 3 ┆ d ┆ [-0.308025, 0.006694, … 0.5338… ┆ 3 ┆ d ┆ [-0.308025, 0.006694, … 0.5338… │ └───────┴──────┴─────────────────────────────────┴─────────────┴────────────┴─────────────────────────────────┘
计算余弦相似度
定义一个函数来计算两个向量之间的余弦相似度。 利用 Polars 表达式，我们可以高效地计算余弦相似度。
cosine_similarity = lambda x, y: (
(x * y).list.sum() / (
(x * x).list.sum().sqrt() * (y * y).list.sum().sqrt()
)
)现在，我们可以使用这个函数来计算每对列之间的余弦相似度。
out = (
combinations_df
.select(
col = "col1",
other = "col1_right",
cosine = cosine_similarity(
x = pl.col("col2"),
y = pl.col("col2_right")
)
)
)
print(out)输出：
shape: (10, 3) ┌─────┬───────┬──────────┐ │ col ┆ other ┆ cosine │ │ --- ┆ --- ┆ --- │ │ str ┆ str ┆ f64 │ ╞═════╪═══════╪══════════╡ │ a ┆ a ┆ 1.0 │ │ a ┆ b ┆ 0.856754 │ │ a ┆ c ┆ 0.827877 │ │ a ┆ d ┆ 0.540282 │ │ b ┆ b ┆ 1.0 │ │ b ┆ c ┆ 0.752199 │ │ b ┆ d ┆ 0.411564 │ │ c ┆ c ┆ 1.0 │ │ c ┆ d ┆ 0.889009 │ │ d ┆ d ┆ 1.0 │ └─────┴───────┴──────────┘
转换为相关矩阵
为了将结果转换为相关矩阵的形式，我们需要将上面的结果进行透视。首先，我们需要将 out DataFrame 中 col 和 other 列互换，然后与原始的 out DataFrame 进行垂直拼接，最后使用 pivot 方法进行透视。
result = pl.concat(
[
out,
out.filter(pl.col("col") != pl.col("other")).select(col="other", other="col", cosine="cosine")
]
).collect().pivot(values="cosine", index="col", columns="other")
print(result)输出：
shape: (4, 5) ┌─────┬──────────┬──────────┬──────────┬──────────┐ │ col ┆ a ┆ b ┆ c ┆ d │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ ╞═════╪══════════╪══════════╪══════════╪══════════╡ │ a ┆ 1.0 ┆ 0.856754 ┆ 0.827877 ┆ 0.540282 │ │ b ┆ 0.856754 ┆ 1.0 ┆ 0.752199 ┆ 0.411564 │ │ c ┆ 0.827877 ┆ 0.752199 ┆ 1.0 ┆ 0.889009 │ │ d ┆ 0.540282 ┆ 0.411564 ┆ 0.889009 ┆ 1.0 │ └─────┴──────────┴──────────┴──────────┴──────────┘
现在，result DataFrame 就是我们想要的相关矩阵，其中每个值表示对应列之间的余弦相似度。
总结
本文档介绍了如何使用 Polars 库计算 DataFrame 中各列之间的余弦相似度，并将其以相关矩阵的形式呈现。通过 join_where 方法生成列组合，利用 Polars 表达式计算余弦相似度，最后使用 pivot 方法将结果转换为矩阵形式。这种方法可以应用于各种数据分析和挖掘任务，例如推荐系统、文本相似度计算等。
注意事项：
- 确保你的 Polars 版本支持列表算术运算。如果你的 Polars 版本低于 1.8.0，请升级到最新版本。
- 在处理大型 DataFrame 时，可以考虑使用 lazy evaluation 来提高性能。
- 余弦相似度是一种常用的相似度度量方法，但它只考虑向量之间的角度，不考虑向量的长度。在某些情况下，可能需要使用其他的相似度度量方法。
