如何为PySpark DataFrame添加欧氏距离列？

我找不到一种简单方法，用外部数组和PySpark列中的数组创建基于欧氏距离的新列。现有以下DataFrame和向量：

rdd = sc.parallelize([(0,[0.1, 0.2, 0.3]), (1,[0.2, 0.3, 0.1]), (2,[0.1, 0.3, 0.2])])
df = sqlContext.createDataFrame(rdd, ["id", "values"])

vector = [0.3, 0.2, 0.1]

我想要计算给定vector与每行values之间的欧氏距离，并将其作为新列。刚接触PySpark，尝试了以下代码但报错：

# 定义计算欧氏距离的UDF
def euclidean_distance(values, query_vector):
    return sqrt(sum((v - q) ** 2 for v, q in zip(values, query_vector)))

udf_euclidean_distance = F.udf(euclidean_distance)

# 通过UDF处理列
df = df.withColumn('similarity_score', 
                   udf_euclidean_distance(col('values'), lit(vector)))

错误信息：

feature not supported.

问题原因及解决方法 #

问题出在lit()无法直接传递数组作为UDF参数，且Python UDF直接使用本地函数存在分布式环境兼容性问题。以下是两种可行解决方法：

方案一：使用Spark内置函数实现（推荐）

用Spark原生数组操作函数计算，性能远高于Python UDF，适合大数据场景：

from pyspark.sql import functions as F
from pyspark.sql.types import FloatType

# 把外部向量转为Spark数组列
query_vector = F.array([F.lit(x) for x in vector])

# 计算元素平方差、求和后开平方
df = df.withColumn("similarity_score",
                   F.sqrt(
                       F.sum(
                           F.transform(
                               F.arrays_zip("values", query_vector),
                               lambda x: (x["values"] - x["0"])**2
                           )
                       ).cast(FloatType())
                   )
                  )

df.show()

方案二：修正自定义UDF的参数传递

如果必须用UDF，通过闭包把外部向量传入UDF，避免lit()传数组的问题：

from pyspark.sql import functions as F
from math import sqrt
from pyspark.sql.types import FloatType

# 用闭包绑定外部向量
def euclidean_distance(query_vector):
    def inner(values):
        return sqrt(sum((v - q)**2 for v, q in zip(values, query_vector)))
    return F.udf(inner, FloatType())

# 创建绑定了目标向量的UDF
udf_euclidean = euclidean_distance(vector)

# 应用UDF生成新列
df = df.withColumn("similarity_score", udf_euclidean(F.col("values")))

df.show()

内容的提问来源于stack exchange，提问作者T_d