引言
在数据分析和机器学习领域，我们经常需要计算数据点之间的相似度。当数据以向量（或列表）的形式存储在 dataframe 的列中时，余弦相似度是一种常用的度量标准。polars 作为一种高性能的 dataframe 库，提供了强大的表达式引擎来处理这类计算。然而，直接将自定义的 python 函数应用于 polars 的聚合操作（如 pivot）可能会遇到 attributeerror: 'function' object has no attribute '_pyexpr' 等问题，这通常是因为 polars 期望接收其内部表达式而不是普通的 python 函数。本文将展示如何利用 polars 的原生特性，优雅地解决这一问题，从而生成一个完整的余弦相似度矩阵。
数据准备
首先，我们定义一个包含列表数据的 Polars DataFrame，这是我们进行相似度计算的基础。
import polars as pl
from numpy.linalg import norm # 虽然这里引入了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 中每个唯一值（例如 'a', 'b', 'c', 'd'）对应的 col2 列表之间的余弦相似度，并最终生成一个交叉矩阵。
Polars 中的余弦相似度计算原理
余弦相似度的数学公式为： $ \text{cosine_similarity}(A, B) = \frac{A \cdot B}{|A| \cdot |B|} $ 其中，$A \cdot B$ 是向量 $A$ 和 $B$ 的点积，$|A|$ 和 $|B|$ 分别是向量 $A$ 和 $B$ 的欧几里得范数（L2 范数）。
在 Polars 中，我们可以将这个公式转化为表达式。值得注意的是，从 Polars 1.8.0 版本开始，Polars 引入了原生的列表算术操作，使得余弦相似度的计算更加高效和简洁。
# 定义 Polars 表达式形式的余弦相似度函数
def calculate_cosine_similarity_expr(x: pl.Expr, y: pl.Expr) -> pl.Expr:
"""
计算两个列表列之间的余弦相似度 Polars 表达式。
要求 Polars 版本 >= 1.8.0 以获得最佳性能。
"""
dot_product = (x * y).list.sum()
norm_x = (x * x).list.sum().sqrt()
norm_y = (y * y).list.sum().sqrt()
return dot_product / (norm_x * norm_y)
# 示例使用：
# cosine_similarity_expr = calculate_cosine_similarity_expr(pl.col("col2"), pl.col("col2_right"))这个表达式利用了 Polars 的列表乘法 (x * y) 来实现元素级别的乘积，然后通过 list.sum() 求和得到点积。欧几里得范数通过 (x * x).list.sum().sqrt() 来计算。这种方式完全在 Polars 的表达式引擎中执行，避免了 Python UDF 的性能开销。
生成数据组合
为了计算所有可能的 pairwise 相似度，我们需要将 DataFrame 中的每一行与所有其他行（包括自身）进行组合。with_row_index() 和 join_where() 是实现这一目标的强大工具。
- 添加行索引: 使用 with_row_index() 为每一行添加一个唯一的索引。
- 条件连接: 使用 join_where() 进行自连接，并设置条件 pl.col.index <= pl.col.index_right。这个条件确保我们只生成一次组合（例如，(a, b) 而不是 (a, b) 和 (b, a)），并且包含自相似度 (a, a)。
# 转换为 lazy DataFrame 以优化性能
lazy_df = df.with_row_index().lazy()
# 生成组合
combinations_df = lazy_df.join_where(lazy_df, pl.col.index <= pl.col.index_right).collect()
print("\n生成的所有组合 (部分):")
print(combinations_df.head())输出：
生成的所有组合 (部分): shape: (5, 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… │ └───────┴──────┴─────────────────────────────────┴─────────────┴────────────┴─────────────────────────────────┘
这个 DataFrame 包含了所有需要计算相似度的向量对。col2 和 col2_right 分别代表了组合中的两个向量。
计算所有组合的余弦相似度
现在，我们将上面定义的余弦相似度表达式应用于 combinations_df 中的 col2 和 col2_right 列。
# 计算余弦相似度
similarity_results = (
lazy_df.join_where(lazy_df, pl.col.index <= pl.col.index_right)
.select(
col="col1",
other="col1_right",
cosine=calculate_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 │ └─────┴───────┴──────────┘
similarity_results DataFrame 包含了每对 col1 值的余弦相似度。由于我们使用了 index <= index_right 的条件，所以它只包含了上三角矩阵和对角线上的值。
构建相似度矩阵
为了得到一个完整的对称相似度矩阵，我们需要处理非对角线元素的对称性（即 cosine(A, B) 等于 cosine(B, A)）。我们可以通过以下步骤完成：
- 复制并反转非对角线元素: 筛选出 col != other 的行，然后交换 col 和 other 列的值，形成反向的组合。
- 合并结果: 将原始的 similarity_results 与反转后的结果合并。
- 透视: 使用 pivot() 方法将数据从长格式转换为宽格式，形成最终的矩阵。
final_similarity_matrix = (
pl.concat(
[
similarity_results,
# 筛选非对角线元素，并反转 col 和 other
similarity_results.filter(pl.col.col != pl.col.other)
.select(col="other", other="col", cosine="cosine")
]
)
.pivot(
values="cosine",
index="col",
columns="other"
)
)
print("\n最终的余弦相似度矩阵:")
print(final_similarity_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 │ └─────┴──────────┴──────────┴──────────┴──────────┘
现在我们得到了一个与期望输出完全一致的余弦相似度矩阵，其中行和列都由 col1 的唯一值表示，矩阵中的每个元素代表相应两个向量的余弦相似度。
注意事项与性能优化
- Polars 版本: 上述余弦相似度表达式利用了 Polars 1.8.0 及更高版本中引入的原生列表算术功能。如果使用较旧的 Polars 版本，可能需要采用不同的方法（例如使用 apply 配合 Python UDF，但这会牺牲性能）。强烈建议升级到最新版本的 Polars 以获得最佳性能和功能。
- 避免 Python UDFs: 尽量避免在 Polars 中使用 Python 用户自定义函数（UDFs），尤其是在性能敏感的场景。Polars 的表达式引擎经过高度优化，能够利用多核并行计算，而 UDFs 会强制数据在 Polars 内部和 Python 解释器之间来回移动，导致性能下降。本教程中的方法完全避免了 UDFs。
- 惰性计算 (.lazy()): 在处理大型数据集时，将 DataFrame 转换为惰性模式 (.lazy()) 可以让 Polars 优化查询计划，从而提高内存效率和执行速度。在最终 collect() 之前，Polars 不会实际执行计算。
- 内存管理: 对于非常大的数据集，生成所有组合可能会消耗大量内存。join_where 配合 lazy() 已经相对高效，但仍需注意数据集大小。
总结
本教程展示了在 Polars 中构建余弦相似度矩阵的完整流程。通过巧妙地结合 with_row_index()、join_where() 生成数据组合，利用 Polars 原生表达式高效计算余弦相似度，并最终通过 pl.concat() 和 pivot() 将结果转换为易于理解的矩阵形式。这种方法不仅解决了直接使用 Python 函数作为聚合器时的错误，而且充分利用了 Polars 的高性能特性，为处理大规模向量相似度计算提供了专业且高效的解决方案。
