Polars列表余弦相似度计算与矩阵生成教程

本教程详细介绍了如何在Polars DataFrame中对列表（List）类型的列进行两两余弦相似度计算，并将结果整理成一个对称的相似度矩阵。通过利用Polars的内置表达式、join_where生成组合以及pivot操作，我们能够高效地处理列表数据并避免使用性能较低的Python UDFs，从而实现类似于相关性矩阵的输出。

在数据分析和机器学习领域，计算向量之间的相似度是常见的任务。当数据存储在Polars DataFrame的列表（List）类型列中时，我们可能需要计算这些列表值之间的两两余弦相似度，并以矩阵形式展示结果，类似于相关性矩阵。本教程将指导您如何高效地完成这一任务，避免直接使用Python用户定义函数（UDFs）可能带来的性能问题。

1. 数据准备

首先，我们创建一个包含列表类型数据的Polars DataFrame作为示例：

import polars as pl

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'）对应的 col2 列表之间的余弦相似度，并最终生成一个4x4的相似度矩阵。

2. 生成所有组合

为了计算所有可能的两两相似度，我们需要将DataFrame与自身进行连接，以生成所有唯一的配对。这里我们将使用 with_row_index() 为每行添加一个索引，然后通过 join_where() 进行条件连接，确保只生成上三角矩阵（包括对角线）的组合，避免重复计算。

# 进入Lazy模式以提高性能
df_lazy = df.with_row_index().lazy()

# 使用join_where生成所有组合，只保留index <= index_right的部分
combinations_df = df_lazy.join_where(df_lazy, pl.col.index <= pl.col.index_right).collect()

print("\n生成的组合DataFrame:")
print(combinations_df)

输出：

生成的组合DataFrame:
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    ┆ [-