在 pyspark 中查找类型为 array<double> 的两列之间的余弦相似度答案

【问题标题】：Find cosine similarity between two columns of type array<double> in pyspark在 pyspark 中查找类型为 array<double> 的两列之间的余弦相似度
【发布时间】：2022-01-16 13:30:55
【问题描述】：

我正在尝试在 pyspark 数据框中找到两列类型数组之间的余弦相似度，并将余弦相似度添加为第三列，如下所示

Col1	Col2	Dot Prod
[0.5, 0.6 ... 0.7]	[0.5, 0.3 .... 0.1]	dotProd(Col1, Col2)

我目前的实现是：

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'])

但我猜上面的代码只适用于具有整数值的列。无论如何我可以扩展上述函数以实现数组列的类似行为

【问题讨论】：

标签： python apache-spark pyspark

【解决方案1】：

是的，上面的代码是针对数字的，而不是针对数字数组的。

您可以将数字数组转换为 pyspark Vectors https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.linalg.Vectors.html

然后调用usedense和dot函数。

例子

from pyspark.ml.linalg import Vectors

x = Vectors.dense([1,2,3])
y = Vectors.dense([4,5,6])

x.dot(y)

同样，你也可以使用pyspark中Vector的norm函数进行归一化。

【讨论】：

【解决方案2】：

对于double的数组，可以使用aggregate函数。

df = spark.createDataFrame([[[0.1, 0.5, 2.0, 1.0], [3.0, 2.4, 0.2, 1.1]]], ['Col1', 'Col2'])
df.show()

+--------------------+--------------------+
|                Col1|                Col2|
+--------------------+--------------------+
|[0.1, 0.5, 2.0, 1.0]|[3.0, 2.4, 0.2, 1.1]|
+--------------------+--------------------+

df.withColumn('dot',   f.expr('aggregate(arrays_zip(Col1, Col2), 0D, (acc, x) -> acc + (x.Col1 * x.Col2))')) \
  .withColumn('norm1', f.expr('sqrt(aggregate(Col1, 0D, (acc, x) -> acc + (x * x)))')) \
  .withColumn('norm2', f.expr('sqrt(aggregate(Col2, 0D, (acc, x) -> acc + (x * x)))')) \
  .withColumn('cosine', f.expr('dot / (norm1 * norm2)')) \
  .show(truncate=False)

+--------------------+--------------------+---+-----------------+-----------------+------------------+
|Col1                |Col2                |dot|norm1            |norm2            |cosine            |
+--------------------+--------------------+---+-----------------+-----------------+------------------+
|[0.1, 0.5, 2.0, 1.0]|[3.0, 2.4, 0.2, 1.1]|3.0|2.293468988235943|4.001249804748511|0.3269133956691457|
+--------------------+--------------------+---+-----------------+-----------------+------------------+

【讨论】：