Spark：比较相似度

Question

我有一个数据框：

index |  Distribution |
+-----+--------------------+
|    0|  [1, 2, 0 ]   |
|    1|  [1, 5, 0 ]   |
|    2|  [1, 9, 0 ]   |
...

和一个清单：

a = [1, 0, 0]

现在我想在此数据框中添加一个新列，其中显示列表 a 与分布列中的元素之间的距离：

index |  Distribution | distan |
+-----+------------------------+
|    0|  [1, 2, 0 ]   |    2   |
|    1|  [1, 5, 0 ]   |    5   |
|    2|  [1, 9, 0 ]   |    9   |
...

我尝试使用

numpy.linalg.norm()

但我不知道如何将其循环到这个数据框中。

请问怎么做？非常感谢！

Answer 1

使用udf：

from scipy.spatial.distance import euclidean as euclidean_
from pyspark.sql.functions import array, lit

df = spark.createDataFrame(
    [(0, [1, 2, 0]), (1, [1, 5, 0]), (2, [1, 9, 0])],
    ["index", "distribution"])

euclidean = lambda x: udf(lambda y: euclidean_(x, y).tolist(), "double")

df.withColumn("dis", euclidean(a)("Distribution"))

## +-----+------------+---+
## |index|distribution|dis|
## +-----+------------+---+
## |    0|   [1, 2, 0]|2.0|
## |    1|   [1, 5, 0]|5.0|
## |    2|   [1, 9, 0]|9.0|
## +-----+------------+---+