PySpark: доступ к векторным элементам в sql

Question

У меня есть искровой фрейм данных, в котором есть столбец с именем 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-запросом выше?

Answer 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|
# +-------------+-----------------+

Если вы хотите разобрать векторы, вы можете использовать Как разбить Вектор на столбцы - используя PySpark