引言：Polars中余弦相似度矩阵构建的挑战
在数据分析和机器学习中，余弦相似度是一种衡量两个非零向量之间角度的度量，广泛应用于文本相似度、推荐系统等领域。当我们需要计算数据集中所有向量对的余弦相似度并以矩阵形式呈现时，polars作为高性能的数据处理库提供了强大的能力。然而，直接使用python自定义函数（udf）与polars的pivot方法结合时，可能会遇到attributeerror: 'function' object has no attribute '_pyexpr'等错误，这表明polars的pivot聚合函数期望的是polars的表达式，而非普通的python函数。
本教程将展示如何利用Polars的内置功能和表达式，高效且优雅地构建一个余弦相似度矩阵，避免UDF的限制，并充分发挥Polars的性能优势。
数据准备
首先，我们定义一个包含向量数据的Polars DataFrame。col1作为标识符，col2存储列表形式的向量数据。
import polars as pl
from numpy.linalg import norm # 尽管这里导入了，但后续我们将使用Polars原生表达式
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("原始DataFrame:")
print(df)输出：
原始DataFrame: 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… │ └──────┴─────────────────────────────────┘
生成所有组合对
要计算所有col1组合之间的余弦相似度，我们需要将DataFrame自身连接，以创建所有可能的向量对。为了避免重复计算（因为余弦相似度是对称的，即sim(a,b) = sim(b,a)），我们可以利用行索引来限制连接条件，只生成上三角部分的组合（包括对角线）。
- 添加行索引: 使用with_row_index()为每一行添加一个唯一的索引。
- 自连接并过滤: 使用join_where()进行自连接，并通过pl.col.index
# 转换为lazy模式以利用Polars的优化
df_lazy = df.with_row_index().lazy()
# 生成组合
combinations = df_lazy.join_where(df_lazy, pl.col.index <= pl.col.index_right).collect()
print("\n生成的组合DataFrame:")
print(combinations)输出：
生成的组合DataFrame: 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表达式计算余弦相似度
余弦相似度的公式为： cos(theta) = (A . B) / (||A|| * ||B||) 其中：
- A . B 是向量A和B的点积。
- ||A|| 是向量A的欧几里得范数（L2范数）。
- ||B|| 是向量B的欧几里得范数。
在Polars中，我们可以利用其强大的表达式系统和列表算术来高效地实现这个公式：
- 点积: (x * y).list.sum()。Polars的列表算术允许直接对列表中的元素进行逐元素乘法，然后list.sum()计算列表元素的总和。
- 范数: (x * x).list.sum().sqrt()。先将向量元素平方，求和，再开方。
我们可以定义一个lambda函数，它接受两个Polars表达式作为输入，并返回一个表示余弦相似度计算的Polars表达式。
cosine_similarity_expr = lambda x, y: (
(x * y).list.sum() / (
(x * x).list.sum().sqrt() * (y * y).list.sum().sqrt()
)
)
# 在combinations DataFrame上应用余弦相似度表达式
similarity_results = (
combinations.select(
col = "col1",
other = "col1_right",
cosine = cosine_similarity_expr(
x = pl.col.col2,
y = pl.col.col2_right
)
)
).collect()
print("\n余弦相似度计算结果（上三角部分）:")
print(similarity_results)输出：
余弦相似度计算结果（上三角部分）: 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 │ └─────┴───────┴──────────┘
现在我们得到了所有组合的余弦相似度，但它只包含了矩阵的上三角部分。
重塑为对称的相似度矩阵
为了将结果转换为一个完整的对称矩阵，我们需要：
- 补充下三角部分: 余弦相似度是可交换的，即sim(A,B) = sim(B,A)。我们可以通过反转col和other列来生成下三角部分的数据。
- 合并数据: 将原始结果（上三角）和反转结果（下三角）合并。
- 透视: 使用pivot操作将other列的值作为新列，col作为索引，cosine作为值，从而构建最终的矩阵。
# 补充下三角部分：筛选出 col != other 的行，然后交换 col 和 other
lower_triangle = similarity_results.filter(pl.col.col != pl.col.other).select(
col=pl.col.other,
other=pl.col.col,
cosine=pl.col.cosine
)
# 合并上三角和下三角数据
full_similarity_data = pl.concat([
similarity_results,
lower_triangle
])
# 使用pivot操作生成最终的矩阵
correlation_matrix = full_similarity_data.pivot(
index="col",
columns="other",
values="cosine"
)
print("\n最终的余弦相似度矩阵:")
print(correlation_matrix)输出：
最终的余弦相似度矩阵: 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 │ └─────┴──────────┴──────────┴──────────┴──────────┘
现在我们得到了一个完整的、对称的余弦相似度矩阵，其结构与相关矩阵类似。
总结与注意事项
- 避免UDFs: 在Polars中，尽可能使用其原生的表达式和方法进行数据操作。自定义Python函数（UDF）通常会导致性能下降，因为它们会强制Polars将数据从其内部优化格式转换到Python对象，然后再转换回来。本教程通过将余弦相似度逻辑完全表达为Polars表达式，避免了这种性能开销。
- 列表算术: Polars的列表（list）类型支持丰富的算术操作和聚合函数（如list.sum()），这使得处理向量数据变得非常方便和高效。
- join_where的高效性: join_where结合索引条件是生成组合对的强大工具，特别适用于对称关系，可以有效减少需要计算的对数。
- pivot的灵活性: pivot方法是重塑数据、从长格式转换为宽格式的关键工具，尤其适用于构建矩阵。
通过上述步骤，我们不仅成功地在Polars中构建了余弦相似度矩阵，还展示了如何利用Polars的强大功能和最佳实践来高效处理复杂的数据转换任务。
