创建距离中心的列

Question

我正在使用pyspark 运行Kmeans 算法。输入是长度为 20 的 Vector（文本 verbatims 上的 word2vec 的输出）。然后我转换我的输入dataframe 以获得与每个verbatim 关联的预测中心。

from pyspark.ml.clustering import KMeans

n_centres = 14
kmeans = KMeans().setK(n_centres).setSeed(1)
model = kmeans.fit(df)
df_pred = model.transform(df)

我有以下结果：

df_pred.show()

+--------------------+----------+
|            features|prediction|
+--------------------+----------+
|[-0.1879145856946...|        13|
|[-0.4428333640098...|         6|
|[0.00466226078569...|         9|
|[0.09467326601346...|        12|
|[-0.0388545106080...|         5|
|[-0.1805213503539...|        13|
|[0.08455141757925...|         3|
+--------------------+----------+

我想在我的数据框中添加一列，其中包含要素数组与其关联的中心之间的距离。我知道我可以得到中心的坐标，我知道如何计算向量和中心之间的距离：

model.clusterCenters()[3] # to get the coordinates of cluster number 3
v1.squared_distance(center_vect) # euclidean distance between v1 and the center center_vect

但我不知道如何将此计算的结果添加为列。 udf 或 map 似乎是一个解决方案，但我不断收到如下错误：PicklingError: Could not serialize object...。

Answer 1

您认为需要使用 UDF 是正确的。这是一个在类似情况下如何工作的示例：

>>> import random
>>> from pyspark.sql.functions import udf
>>> centers = {1: 2, 2: 3, 3: 4, 4:5, 5:6}
>>> choices = [1, 2, 3, 4,5]
>>> l = [(random.random(), random.choice(choices)) for i in range(10)]
>>> df = spark.createDataFrame(df, ['features', 'prediction'])
>>> df.show()
+-------------------+----------+
|           features|prediction|
+-------------------+----------+
| 0.4836744206538728|         3|
|0.38698675915124414|         4|
|0.18612684714681604|         3|
| 0.5056159922655895|         1|
| 0.7825023909896331|         4|
|0.49933715239708243|         5|
| 0.6673811293962939|         4|
| 0.7010166164833609|         3|
| 0.6867109795526414|         5|
|0.21975859257732422|         3|
+-------------------+----------+
>>> dist = udf(lambda features, prediction: features - centers[prediction])
>>> df.withColumn('dist', dist(df.features, df.prediction)).show()
+-------------------+----------+-------------------+
|           features|prediction|               dist|
+-------------------+----------+-------------------+
| 0.4836744206538728|         3| -3.516325579346127|
|0.38698675915124414|         4| -4.613013240848756|
|0.18612684714681604|         3| -3.813873152853184|
| 0.5056159922655895|         1|-1.4943840077344106|
| 0.7825023909896331|         4| -4.217497609010367|
|0.49933715239708243|         5| -5.500662847602918|
| 0.6673811293962939|         4|-4.3326188706037065|
| 0.7010166164833609|         3| -3.298983383516639|
| 0.6867109795526414|         5| -5.313289020447359|
|0.21975859257732422|         3| -3.780241407422676|
+-------------------+----------+-------------------+

您可以将我创建 UDF 的行更改为如下所示：

dist = udf(lambda features, prediction: features.squared_distance(model.clusterCenters()[prediction]))

由于我没有可用的实际数据，我希望这是正确的！