用Polars快速计算余弦相似度矩阵

哈喽！大家好，很高兴又见面了，我是golang学习网的一名作者，今天由我给大家带来一篇《用Polars快速构建余弦相似度矩阵：数据到可视化全流程》，本文主要会讲到等等知识点，希望大家一起学习进步，也欢迎大家关注、点赞、收藏、转发! 下面就一起来看看吧！

本文详细介绍了如何在Polars中计算并构建余弦相似度矩阵。针对Polars的表达式限制，教程首先展示了如何通过with_row_index和join_where生成所有数据对组合，然后利用Polars原生的列表算术和表达式高效地计算余弦相似度，避免了自定义Python函数的性能瓶颈。最后，通过concat和pivot操作将结果转换为对称的矩阵形式，实现了类似相关矩阵的输出，强调了在Polars中利用原生表达式的重要性。

引言：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)），我们可以利用行索引来限制连接条件，只生成上三角部分的组合（包括对角线）。

1. 添加行索引: 使用with_row_index()为每一行添加一个唯一的索引。
2. 自连接并过滤: 使用join_where()进行自连接，并通过pl.col.index <= pl.col.index_right条件确保只生成索引较小的行与索引较大或相等的行进行组合。这在处理对称关系时非常高效。

# 转换为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      │
└─────┴───────┴──────────┘

现在我们得到了所有组合的余弦相似度，但它只包含了矩阵的上三角部分。

重塑为对称的相似度矩阵

为了将结果转换为一个完整的对称矩阵，我们需要：

1. 补充下三角部分: 余弦相似度是可交换的，即sim(A,B) = sim(B,A)。我们可以通过反转col和other列来生成下三角部分的数据。
2. 合并数据: 将原始结果（上三角）和反转结果（下三角）合并。
3. 透视: 使用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的强大功能和最佳实践来高效处理复杂的数据转换任务。

好了，本文到此结束，带大家了解了《用Polars快速计算余弦相似度矩阵》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多文章知识！