PYSPARK：如何在 pyspark 数据框中找到两列的余弦相似度？

Posted 2023-04-13

【中文标题】PYSPARK：如何在 pyspark 数据框中找到两列的余弦相似度？

【英文标题】：PYSPARK: How to find cosine similarity of two columns in a pyspark dataframe?

【发布时间】：2017-05-09 21:33:55

【问题描述】：

如何在 pyspark 数据框中找到两列之间的余弦相似度？

假设我有一个 spark 数据框

|a |b |
+--+--|
|1 |4 |
|2 |5 |
|3 |6 |
+--+--+

现在我想知道 a 列中的值与 b 列中的值之间的余弦相似度是多少，即

cosine_similarity([1, 2, 3], [4, 5, 6]) 

【问题讨论】：

您能否再解释一下您的问题并给出示例数据框？您是指 2 列中每行元素之间的余弦相似度还是数据帧中 2 列之间的余弦相似度？

【参考方案1】：

我假设您想找到 2 列之间的相似性。说你有这个数据框：

df = spark.createDataFrame(pd.DataFrame([[1,2], [3,4]], columns=['a', 'b']))

制作简单的函数来获取数据框和两个列名。

import pyspark.sql.functions as func

def cosine_similarity(df, col1, col2):
    df_cosine = df.select(func.sum(df[col1] * df[col2]).alias('dot'), 
                          func.sqrt(func.sum(df[col1]**2)).alias('norm1'), 
                          func.sqrt(func.sum(df[col2] **2)).alias('norm2'))
    d = df_cosine.rdd.collect()[0].asDict()
    return d['dot']/(d['norm1'] * d['norm2'])

cosine_similarity(df, 'a', 'b') # output 0.989949

【讨论】：

是的，这个函数是用于 spark 数据帧的。