本教程详细介绍了如何在 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 plfrom 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
# 转换为 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
构建相似度矩阵
为了得到一个完整的对称相似度矩阵，我们需要处理非对角线元素的对称性（即 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 的高性能特性，为处理大规模向量相似度计算提供了专业且高效的解决方案。
以上就是Polars DataFrame 余弦相似度矩阵的构建方法的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1376735.html
微信扫一扫 
支付宝扫一扫 
