PySpark：在 sql 中访问向量元素

Posted 2023-03-12

【中文标题】PySpark：在 sql 中访问向量元素

【英文标题】：PySpark: accessing vector elements in sql

【发布时间】：2018-10-25 01:14:54

【问题描述】：

我的 spark 数据框有一个名为 features 的列，其中包含数据向量。此列是pyspark 的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】：

只需使用 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

【讨论】：

感谢您的回答。这工作正常。但是有没有办法我们可以避免 udf 并且仍然得到结果？我已经阅读（并经历过）udf 效率不高并且减慢了进程。我们使用的数据量是每个数据集大约 1000 万条记录，并且有多个数据集，所以我不确定 udf 是否适合这里

没有。向量没有本机实现，因此无法直接访问。这可能会在未来发生变化（尽管我不期望性能提升，只有更好的 API）。