如何在 pyspark 数据帧中找到两列之间的余弦相似性?
假设我有一个火花数据帧
|a |b |
+--+--|
|1 |4 |
|2 |5 |
|3 |6 |
+--+--+
现在我想知道 a 列中的值和 b 列中的值之间的余弦相似性是什么,即
cosine_similarity([1, 2, 3], [4, 5, 6])
我假设您想找到 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