PySpark：將RDD轉換為數據框中的列

Question

我有一個spark數據框，通過它可以計算行與給定的corcordates之間的歐幾里得距離。 我在這里重新創建結構相似的數據幀“ 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

除了創建新的rdd之外，還可以使用UDF ：

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()