本教程详细介绍了如何在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 ┆ [-
