计算向量与不同火花数据框的距离

Question

我有两个火花数据框架：

> df1
+--------+-----------------------------+
|   word |                     init_vec|
+--------+-----------------------------+
|  venus |[-0.235, -0.060, -0.609, ...]|
+--------+-----------------------------+


> df2
+-----------------------------+-----+
|                    targ_vec |   id|
+-----------------------------+-----+
|[-0.272, -0.070, -0.686, ...]| 45ha|
+-----------------------------+-----+
|[-0.234, -0.060, -0.686, ...]| 98pb|
+-----------------------------+-----+
|[-0.562, -0.334, -0.981, ...]| c09j|
+-----------------------------+-----+

我需要从df1和targ_vec of of of of of of of of df2 init_vec之间的欧几里得距离。 > dataframe，然后将最接近3个最接近的向量返回到init_vec。

    > desired_output
    +--------+-----------------------------+-----+----------+
    |   word |                     targ_vec|   id|  distance|
    +--------+-----------------------------+-----+----------+
    |  venus |[-0.234, -0.060, -0.686, ...]| 98pb|some_value|
    +--------+-----------------------------+-----+----------+
    |  venus |[-0.221, -0.070, -0.613, ...]| tg67|some_value|
    +--------+-----------------------------+-----+----------+
    |  venus |[-0.240, -0.091, -0.676, ...]| jhg6|some_value|
    +--------+-----------------------------+-----+----------+

我需要使用Pyspark实施此功能。

Answer 1

在DF1和DF2之间的交叉接合以将DF1.Init_Vec添加到DF2的所有行中：

df1 = (df1
      .withColumn('distance', f.sqrt(f.expr('aggregate(transform(targ_vec, (element, idx) -> power(abs(element - element_at(init_vec, cast(idx + 1 as int))), 2)), cast(0 as double), (acc, value) -> acc + value)')))
     )

然后，您可以对数据框进行排序，并将3行保持在最小距离值的情况下。