使用 Polars 计算 DataFrame 的相关矩阵：余弦相似度方法详解-Python教程-PHP中文网

使用 Polars 计算 DataFrame 的相关矩阵：余弦相似度方法详解

心靈之曲

发布： 2025-10-10 14:19:17

原创

289人浏览过

使用 polars 计算 dataframe 的相关矩阵：余弦相似度方法详解

本文档详细介绍了如何使用 Polars 库计算 DataFrame 中各列之间的余弦相似度，并将其以相关矩阵的形式呈现。通过 join_where 方法生成列组合，利用 Polars 表达式计算余弦相似度，最后使用 pivot 方法将结果转换为矩阵形式，方便进行数据分析和挖掘。

前提条件

确保你已经安装了 Polars 库。可以使用 pip 进行安装：

pip install polars

登录后复制

数据准备

首先，我们创建一个 Polars DataFrame，其中包含字符串列 col1 和列表列 col2。col2 列包含数值列表，我们将基于这些列表计算余弦相似度。

import polars as pl
from numpy.linalg import norm

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(df)

登录后复制

输出：

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… │
└──────┴─────────────────────────────────┘

登录后复制

生成列组合

为了计算每对列之间的余弦相似度，我们需要生成所有可能的列组合。我们可以使用 join_where 方法来实现这一点。首先，添加一个行索引，然后使用 join_where 将 DataFrame 与自身连接，条件是左侧的索引小于等于右侧的索引，以避免重复计算。

df = df.with_row_index().lazy()

combinations_df = df.join_where(df, pl.col("index") <= pl.col("index_right")).collect()

print(combinations_df)

登录后复制

输出：

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 表达式，我们可以高效地计算余弦相似度。

cosine_similarity = lambda x, y: (
    (x * y).list.sum() / (
        (x * x).list.sum().sqrt() * (y * y).list.sum().sqrt()
    )
)

登录后复制

现在，我们可以使用这个函数来计算每对列之间的余弦相似度。

通义万相

通义万相，一个不断进化的AI艺术创作大模型

596

查看详情

out = (
   combinations_df
     .select(
        col = "col1",
        other = "col1_right",
        cosine = cosine_similarity(
           x = pl.col("col2"),
           y = pl.col("col2_right")
        )
     )
)

print(out)

登录后复制

输出：

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      │
└─────┴───────┴──────────┘

登录后复制

转换为相关矩阵

为了将结果转换为相关矩阵的形式，我们需要将上面的结果进行透视。首先，我们需要将 out DataFrame 中 col 和 other 列互换，然后与原始的 out DataFrame 进行垂直拼接，最后使用 pivot 方法进行透视。

result = pl.concat(
   [
      out, 
      out.filter(pl.col("col") != pl.col("other")).select(col="other", other="col", cosine="cosine")
   ]
).collect().pivot(values="cosine", index="col", columns="other")

print(result)

登录后复制

输出：

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      │
└─────┴──────────┴──────────┴──────────┴──────────┘

登录后复制

现在，result DataFrame 就是我们想要的相关矩阵，其中每个值表示对应列之间的余弦相似度。

总结

本文档介绍了如何使用 Polars 库计算 DataFrame 中各列之间的余弦相似度，并将其以相关矩阵的形式呈现。通过 join_where 方法生成列组合，利用 Polars 表达式计算余弦相似度，最后使用 pivot 方法将结果转换为矩阵形式。这种方法可以应用于各种数据分析和挖掘任务，例如推荐系统、文本相似度计算等。

注意事项：

确保你的 Polars 版本支持列表算术运算。如果你的 Polars 版本低于 1.8.0，请升级到最新版本。
在处理大型 DataFrame 时，可以考虑使用 lazy evaluation 来提高性能。
余弦相似度是一种常用的相似度度量方法，但它只考虑向量之间的角度，不考虑向量的长度。在某些情况下，可能需要使用其他的相似度度量方法。

以上就是使用 Polars 计算 DataFrame 的相关矩阵：余弦相似度方法详解的详细内容，更多请关注php中文网其它相关文章！