PySpark:访问sql中的向量元素

时间:2018-05-15 09:02:54

标签: python apache-spark pyspark k-means

我的spark数据框有一个名为features的列,用于保存数据向量。此列是pyspark&n; StandardScaler对象的输出。我在这里创建一个类似于我的数据集。

# create sample data
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]|
+-------------+

我想找到每个向量与特定聚类中心(相同长度的数组)之间的欧几里德距离。假设集群中心是:

cluster_center_0 = np.array([0.6, 0.7, 0.8])

我如何实现这一目标?我尝试创建一个SQL查询,希望我可以使用OFFSET访问向量中的元素,从那里可以很容易地计算距离。但那并没有成功。这是我使用的查询。不幸的是它没有用,我对sql的知识非常有限

SELECT aml_cluster_inpt_features
aml_cluster_inpt_features[OFFSET(0)] AS offset_0,
aml_cluster_inpt_features[OFFSET(1)] AS offset_1,
aml_cluster_inpt_features[OFFSET(2)] AS offset_2,
aml_cluster_inpt_features[OFFSET(3)] AS offset_3,
FROM event_rate_holder

有更简单的方法吗?如果没有,我是否正朝着正确的方向前进?上面的SQL查询?

1 个答案:

答案 0 :(得分:2)

只需使用UDF:

from pyspark.sql.functions import udf
from scipy.spatial import distance

def euclidean(v1):
    @udf("double")
    def _(v2):
        return distance.euclidean(v1, v2) if v2 is not None else None
    return _


center = np.array([0.6, 0.7, 0.8])

df_vector.withColumn("dist", euclidean(center)("features")).show()
# +-------------+-----------------+
# |     features|             dist|
# +-------------+-----------------+
# |[1.0,2.0,3.0]|2.586503431275513|
# |[4.0,5.0,6.0]|7.555792479945437|
# +-------------+-----------------+

如果您想反汇编矢量,可以使用How to split Vector into columns - using PySpark