UDF转换输出

Question

在Pyspark上，我定义了一个UDF如下：

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

def closest_point(point, points):
    """ Find closest point from a list of points. """
    return points[cdist([point], points).argmin()]

udf_closest_point = udf(closest_point)

dfC1 = dfC1.withColumn("closest", udf_closest_point(dfC1.point, dfC1.points))

我的数据是这样的：

• 点= 0.2,0.5或0.1,0.6 -浮点数数组
• 点=[0，1，1，0，1，1，1，0]-浮点数数组
• 最近=例如，'0，1‘-字符串(它是从点转换成字符串的值之一)

对于我的UDF，我应该更改什么来返回一个浮点数数组，而不是一个字符串？

Answer 1

可以将UDF的返回类型指定为浮点数ArrayType(FloatType())数组。

from pyspark.sql.types import ArrayType, FloatType
udf_closest_point = udf(closest_point, ArrayType(FloatType()))