Question

我有一个火花数据框，我用它计算一行和一组给定的corrdinates之间的欧几里德距离。我在这里重新创建一个结构相似的数据框'df_vector'以便更好地解释。

from pyspark.ml.feature import VectorAssembler
arr = [[1,2,3], [4,5,6]]
df_example = spark.createDataFrame(arr, ['A','B','C'])
assembler = VectorAssembler(inputCols=[x for x in df_example.columns],outputCol='features')
df_vector = assembler.transform(df_example).select('features') 

>>> df_vector.show()
+-------------+
|     features|
+-------------+
|[1.0,2.0,3.0]|
|[4.0,5.0,6.0]|
+-------------+

>>> df_vector.dtypes
[('features', 'vector')]

正如您所见，features列是一个向量。实际上，我将此向量列作为StandardScaler的输出。无论如何，由于我需要计算欧几里德距离，我会做以下

rdd = df_vector.select('features').rdd.map(lambda r: np.linalg.norm(r-b))

其中

b = np.asarray([0.5,1.0,1.5])

我有我需要的所有计算但我需要rdd作为df_vector中的列。我该怎么做呢？

Answer 1

您可以使用UDF：

，而不是创建新的rdd

norm_udf = udf(lambda r: np.linalg.norm(r - b).tolist(), FloatType())
df_vector.withColumn("norm", norm_udf(df.features))

确保在工作节点上定义numpy。

Answer 2

解决性能问题的一种方法可能是使用mapPartitions。在分区级别，想法是将features转换为数组，然后计算整个数组的范数（因此隐式使用numpy向量化）。然后做一些家务，以获得你想要的形式。对于大型数据集，这可能会提高性能：

这是在分区级别计算范数的函数：

from pyspark.sql import Row
def getnorm(vectors):
    # convert vectors into numpy array
    vec_array=np.vstack([v['features'] for v in vectors])
    # calculate the norm
    norm=np.linalg.norm(vec_array-b, axis=1)
    # tidy up to get norm as a column
    output=[Row(features=x, norm=y) for x,y in zip(vec_array.tolist(), norm.tolist())]
    return(output)

使用mapPartitions应用此项会得到一个行的RDD，然后可以将其转换为DataFrame：

df_vector.rdd.mapPartitions(getnorm).toDF()

PySpark：将RDD转换为数据框

2 个答案: