在数据分析和机器学习领域，计算向量之间的相似度是一个常见任务，其中余弦相似度因其对向量长度不敏感的特性而广受欢迎。当需要在 Polars DataFrame 中计算所有向量对的余弦相似度并以矩阵形式呈现时，直接应用 Python 用户定义函数（UDF）往往会遇到性能瓶颈或兼容性问题。本教程将展示如何利用 Polars 的强大表达式引擎，高效、专业地完成这一任务。
问题背景与挑战
假设我们有一个 Polars DataFrame，其中包含一个标识符列 (col1) 和一个包含浮点数列表的向量列 (col2)：
import polars as pl
from numpy.linalg import norm # 原始问题中提到的numpy函数，但我们将用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)我们的目标是计算 col1 中每个组合的 col2 向量之间的余弦相似度，并生成一个类似相关矩阵的输出，例如：
┌─────────────────┬──────┬──────┬──────┬──────┐ │ col1_col2 ┆ a ┆ b ┆ c ┆ d │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ ╞═════════════════╪══════╪══════╪══════╪══════╡ │ a ┆ 1.0 ┆ 0.86 ┆ 0.83 ┆ 0.54 │ │ b ┆ 0.86 ┆ 1.0 ┆ 0.75 ┆ 0.41 │ │ c ┆ 0.83 ┆ 0.75 ┆ 1.0 ┆ 0.89 │ │ d ┆ 0.54 ┆ 0.41 ┆ 0.89 ┆ 1.0 │ └─────────────────┴──────┴──────┴──────┴──────┘
初学者可能会尝试定义一个 Python lambda 函数作为余弦相似度计算器，并将其直接传递给 pivot 方法。然而，Polars 的 pivot 方法期望一个 Polars 表达式作为聚合函数，而不是一个普通的 Python 函数。这会导致 AttributeError: 'function' object has no attribute '_pyexpr' 错误。因此，我们需要一种完全基于 Polars 表达式的解决方案。
步骤一：生成所有组合对
要计算所有向量对的相似度，首先需要生成这些对。Polars 提供了 with_row_index() 和 join_where() 方法，可以高效地完成此任务。with_row_index() 为 DataFrame 的每一行添加一个唯一的索引，而 join_where() 则允许我们基于条件连接 DataFrame 自身，从而生成所有可能的行组合。
为了避免重复计算（例如，计算 (a, b) 后无需再计算 (b, a)），我们可以在 join_where 中使用 pl.col.index
# 将DataFrame转换为LazyFrame以进行高效操作
lazy_df = df.with_row_index().lazy()
# 使用 join_where 生成所有组合对
# pl.col.index <= pl.col.index_right 确保我们只获取唯一的对（包括自身）
combinations_df = lazy_df.join_where(lazy_df, pl.col("index") <= pl.col("index_right")).collect()
print("\n生成的组合对 DataFrame:")
print(combinations_df)
输出的 combinations_df 将包含原始 DataFrame 的所有行对，每对数据都以 col1, col2 和 col1_right, col2_right 的形式呈现。
步骤二：使用 Polars 表达式计算余弦相似度
余弦相似度的数学公式定义为：
$$ \text{similarity} = \frac{A \cdot B}{|A| \cdot |B|} = \frac{\sum_{i=1}^{n} A_i Bi}{\sqrt{\sum{i=1}^{n} Ai^2} \cdot \sqrt{\sum{i=1}^{n} B_i^2}} $$
在 Polars 中，我们可以利用其强大的列表算术功能（自 Polars 1.8.0 版本起得到显著增强）和聚合函数来实现这个公式。
-
点积 (Numerator): (x * y).list.sum()
- x * y 对两个列表进行逐元素乘法。
- .list.sum() 计算乘积列表的和，得到点积。
-
向量模的乘积 (Denominator): (x * x).list.sum().sqrt() * (y * y).list.sum().sqrt()
- x * x 对列表 x 进行逐元素平方。
- .list.sum() 计算平方和。
- .sqrt() 计算平方和的平方根，得到向量 x 的模。
- 对向量 y 执行相同的操作，然后将两个模相乘。
将这些组合起来，我们可以构建一个 Polars 表达式来计算余弦相似度：
# 定义余弦相似度 Polars 表达式
cosine_similarity_expr = lambda x, y: (
(x * y).list.sum() / (
(x * x).list.sum().sqrt() * (y * y).list.sum().sqrt()
)
)
# 在组合对 DataFrame 上应用余弦相似度表达式
similarity_results = (
lazy_df.join_where(lazy_df, pl.col("index") <= pl.col("index_right"))
.select(
col = pl.col("col1"),
other = pl.col("col1_right"),
cosine = cosine_similarity_expr(
x = pl.col("col2"),
y = pl.col("col2_right")
)
)
.collect()
)
print("\n计算出的余弦相似度（长格式）:")
print(similarity_results)similarity_results DataFrame 现在包含了所有唯一对的余弦相似度，以长格式呈现。
步骤三：将结果转换为对称矩阵
由于我们在 join_where 中使用了 pl.col.index
- 复制并反转非对角线元素： 筛选出 col != other 的行，然后交换 col 和 other 列的值。
- 合并结果： 将原始结果与反转后的结果进行垂直拼接 (pl.concat)。
- 透视 (Pivot) 成矩阵： 使用 pivot 方法将长格式数据转换为宽格式的相似度矩阵。
# 复制并反转非对角线元素，以创建对称矩阵的下半部分
reversed_pairs = 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, reversed_pairs])
# 使用 pivot 方法将长格式数据转换为宽格式的相似度矩阵
correlation_matrix = full_similarity_data.pivot(
values="cosine",
index="col",
columns="other"
).sort("col") # 对行进行排序，使输出更规整
print("\n最终的余弦相似度矩阵:")
print(correlation_matrix)最终输出的 correlation_matrix 就是我们期望的余弦相似度矩阵，它是一个对称的方阵，其行和列由 col1 的值标识，单元格中包含相应的余弦相似度。
注意事项与最佳实践
- Polars 表达式的优势： 本教程的核心在于使用 Polars 的原生表达式 (x * y, list.sum(), sqrt()) 而非 Python UDF。Polars 表达式在内部经过优化，可以在其 Rust 后端高效执行，从而提供卓越的性能，尤其是在处理大型数据集时。避免使用 apply 或自定义 Python 函数是 Polars 中的一项重要最佳实践。
- Polars 版本要求： 列表算术（例如 x * y 对列表进行逐元素操作）在 Polars 1.8.0 及更高版本中得到了显著增强和优化。建议使用最新版本的 Polars 以获得最佳性能和功能。
- 内存管理： 对于非常大的 DataFrame，join_where 可能会生成一个非常大的中间 DataFrame。如果内存成为问题，可以考虑分块处理或优化数据类型。
- 可读性： 将复杂的表达式封装在 lambda 函数中可以提高代码的可读性和模块化，但请记住，这个 lambda 函数本身需要返回一个 Polars 表达式，而不是一个直接计算结果的 Python 函数。
通过遵循本教程的步骤，您可以在 Polars 中高效、专业地计算并构建余弦相似度矩阵，为您的数据分析工作提供强大支持。
