本教程详细介绍了如何在 Polars DataFrame 中高效计算并构建余弦相似度矩阵。通过利用 Polars 的原生表达式和 join_where 方法，我们避免了使用低效的 Python UDF，从而实现了高性能的相似度计算。文章涵盖了从数据准备、生成组合、余弦相似度表达式的实现到最终矩阵转换的完整流程，帮助用户在 Polars 中专业处理向量相似度分析。
在数据分析和机器学习领域，计算向量之间的相似度是一个常见任务，其中余弦相似度因其对向量长度不敏感的特性而广受欢迎。当需要在 Polars DataFrame 中计算所有向量对的余弦相似度并以矩阵形式呈现时，直接应用 Python 用户定义函数（UDF）往往会遇到性能瓶颈或兼容性问题。本教程将展示如何利用 Polars 的强大表达式引擎，高效、专业地完成这一任务。
问题背景与挑战
假设我们有一个 Polars DataFrame，其中包含一个标识符列 (col1) 和一个包含浮点数列表的向量列 (col2)：
import polars as plfrom 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
# 将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 表达式计算余弦相似度
余弦相似度的数学公式定义为：
$$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 版本起得到显著增强）和聚合函数来实现这个公式。
点积 (Numerator): (x * y).list.sum()x * y 对两个列表进行逐元素乘法。.list.sum() 计算乘积列表的和，得到点积。向量模的乘积 (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
复制并反转非对角线元素： 筛选出 col != other 的行，然后交换 col 和 other 列的值。合并结果： 将原始结果与反转后的结果进行垂直拼接 (pl.concat)。透视 (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 表达式构建高效的余弦相似度矩阵的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1377041.html
微信扫一扫 
支付宝扫一扫 
