Polars DataFrame计算余弦相似度触发AttributeError问题求助

问题背景 #

现有如下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)

需要计算col1各元素对应向量的两两余弦相似度，期望得到如下格式的结果：

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

尝试使用自定义numpy余弦相似度函数结合pivot方法计算：

from numpy.linalg import norm

cosine_similarity = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
df.pivot(on="col1", values="col2", index="col1", aggregate_function=cosine_similarity)

触发错误：

AttributeError: 'function' object has no attribute '_pyexpr'

错误原因 #

Polars的pivot方法中，aggregate_function参数要求传入Polars表达式（Expr），而非普通Python函数。pivot的设计是对分组后的列进行聚合操作，而非用于计算两两元素的相似度，直接传入自定义Python函数会导致Polars无法识别并抛出属性错误。

解决方法 #

方法一：基于Numpy矩阵运算（高效推荐） #

利用Numpy将向量转换为矩阵，批量计算余弦相似度矩阵，再转换为Polars DataFrame：

import polars as pl
import numpy as np
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)

# 将col2的向量转换为Numpy矩阵
vectors = np.array(df["col2"].to_list())

# 计算余弦相似度矩阵
norm_matrix = norm(vectors, axis=1)[:, np.newaxis]
cos_sim_matrix = vectors @ vectors.T / (norm_matrix @ norm_matrix.T)

# 构造结果DataFrame并格式化
result_df = pl.DataFrame(
    cos_sim_matrix,
    columns=df["col1"].to_list()
).with_columns(col1_col2=df["col1"]).select(["col1_col2"] + df["col1"].to_list())

# 保留两位小数
result_df = result_df.with_columns(pl.all().exclude("col1_col2").round(2))

print(result_df)

方法二：Polars交叉连接+Pivot（原生API） #

通过自交叉连接生成所有两两组合，计算相似度后再pivot为宽表：

import polars as pl
import numpy as np
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)

# 自交叉连接生成所有两两组合
cross_df = df.join(df, how="cross", suffix="_right")

# 计算每对组合的余弦相似度
cos_sim_df = cross_df.with_columns(
    pl.struct(["col2", "col2_right"]).map_elements(
        lambda x: (np.array(x["col2"]) @ np.array(x["col2_right"])) / (norm(x["col2"]) * norm(x["col2_right"])),
        return_dtype=pl.Float64
    ).alias("cos_sim")
)

# Pivot为宽表并格式化
result_df = cos_sim_df.pivot(
    index="col1",
    columns="col1_right",
    values="cos_sim"
).rename({"col1": "col1_col2"}).with_columns(pl.all().exclude("col1_col2").round(2))

print(result_df)

方法对比 #

• 方法一：利用Numpy矩阵运算，计算效率更高，适合处理大数据量的向量相似度计算。
• 方法二：完全基于Polars原生API实现，逻辑更贴合Polars的列式思维，但map_elements是逐行操作，数据量大时效率低于矩阵运算。

内容的提问来源于stack exchange，提问作者Okroshiashvili