使用 Polars 表达式构建高效的余弦相似度矩阵

在数据分析和机器学习领域，计算向量之间的相似度是一个常见任务，其中余弦相似度因其对向量长度不敏感的特性而广受欢迎。当需要在 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 <= pl.col.index_right 条件，确保只生成上三角部分的组合（包括对角线上的自身组合）。

# 将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 表达式计算余弦相似度

余弦相似度的数学公式定义为：

乾坤圈新媒体矩阵管家

新媒体账号、门店矩阵智能管理系统

17

查看详情

$$ \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 版本起得到显著增强）和聚合函数来实现这个公式。

1. 点积 (Numerator): (x * y).list.sum()
   • x * y 对两个列表进行逐元素乘法。
   • .list.sum() 计算乘积列表的和，得到点积。
2. 向量模的乘积 (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 <= pl.col.index_right 条件，similarity_results 只包含了矩阵的上三角部分（包括对角线）。为了构建一个完整的对称矩阵，我们需要添加下三角部分的条目。这可以通过以下步骤完成：

1. 复制并反转非对角线元素： 筛选出 col != other 的行，然后交换 col 和 other 列的值。
2. 合并结果： 将原始结果与反转后的结果进行垂直拼接 (pl.concat)。
3. 透视 (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 中高效、专业地计算并构建余弦相似度矩阵，为您的数据分析工作提供强大支持。

以上就是使用 Polars 表达式构建高效的余弦相似度矩阵的详细内容，更多请关注php中文网其它相关文章！