构建PolarsDataFrame余弦相似度矩阵教程

本文详细介绍了在Polars DataFrame中构建余弦相似度矩阵的实用方法，适用于数据分析和机器学习场景。针对列表类型列，文章阐述了如何利用`join_where`高效生成数据组合，并运用Polars原生表达式计算向量间的余弦相似度，避免了传统Python UDF带来的性能瓶颈。同时，文章还讲解了如何通过`pivot`操作将计算结果转化为清晰的相似度矩阵，便于理解和应用。该方法充分利用Polars的强大表达式引擎，优化内存管理，并兼容Polars 1.8.0及以上版本的新特性，为开发者提供了一个高效、专业的解决方案，以应对大规模向量相似度计算的需求。

本教程详细介绍了如何在 Polars DataFrame 中高效计算列表类型列之间的余弦相似度，并将其结果转换为一个类似相关系数矩阵的宽格式 DataFrame。文章将通过 join_where 生成数据组合，利用 Polars 原生表达式计算余弦相似度，并最终通过 pivot 操作构建出完整的对称相似度矩阵。

引言

在数据分析和机器学习领域，我们经常需要计算数据点之间的相似度。当数据以向量（或列表）的形式存储在 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() 是实现这一目标的强大工具。

1. 添加行索引: 使用 with_row_index() 为每一行添加一个唯一的索引。
2. 条件连接: 使用 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)）。我们可以通过以下步骤完成：

1. 复制并反转非对角线元素: 筛选出 col != other 的行，然后交换 col 和 other 列的值，形成反向的组合。
2. 合并结果: 将原始的 similarity_results 与反转后的结果合并。
3. 透视: 使用 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 的唯一值表示，矩阵中的每个元素代表相应两个向量的余弦相似度。

注意事项与性能优化

1. Polars 版本: 上述余弦相似度表达式利用了 Polars 1.8.0 及更高版本中引入的原生列表算术功能。如果使用较旧的 Polars 版本，可能需要采用不同的方法（例如使用 apply 配合 Python UDF，但这会牺牲性能）。强烈建议升级到最新版本的 Polars 以获得最佳性能和功能。
2. 避免 Python UDFs: 尽量避免在 Polars 中使用 Python 用户自定义函数（UDFs），尤其是在性能敏感的场景。Polars 的表达式引擎经过高度优化，能够利用多核并行计算，而 UDFs 会强制数据在 Polars 内部和 Python 解释器之间来回移动，导致性能下降。本教程中的方法完全避免了 UDFs。
3. 惰性计算 (.lazy()): 在处理大型数据集时，将 DataFrame 转换为惰性模式 (.lazy()) 可以让 Polars 优化查询计划，从而提高内存效率和执行速度。在最终 collect() 之前，Polars 不会实际执行计算。
4. 内存管理: 对于非常大的数据集，生成所有组合可能会消耗大量内存。join_where 配合 lazy() 已经相对高效，但仍需注意数据集大小。

总结

本教程展示了在 Polars 中构建余弦相似度矩阵的完整流程。通过巧妙地结合 with_row_index()、join_where() 生成数据组合，利用 Polars 原生表达式高效计算余弦相似度，并最终通过 pl.concat() 和 pivot() 将结果转换为易于理解的矩阵形式。这种方法不仅解决了直接使用 Python 函数作为聚合器时的错误，而且充分利用了 Polars 的高性能特性，为处理大规模向量相似度计算提供了专业且高效的解决方案。

本篇关于《构建PolarsDataFrame余弦相似度矩阵教程》的介绍就到此结束啦，但是学无止境，想要了解学习更多关于文章的相关知识，请关注golang学习网公众号！