PySpark创建接收多列参数的UDF运行报错，如何正确实现？

问题根因 #

你遇到的任务中止报错主要由两个原因导致：

• 自定义的haversine函数使用numpy计算，返回值为numpy.float类型，和你注册UDF时指定的FloatType()要求的Python原生浮点类型不兼容，序列化过程出错
• 若Spark集群的worker节点未安装numpy依赖，也会触发执行阶段的任务失败

多参数UDF正确实现方式 #

方案1：修改原有Python UDF #

调整返回值类型，补充空值判断避免脏数据报错：

import numpy as np
from pyspark.sql.types import FloatType
from pyspark.sql.functions import udf

def haversine(lon1, lat1, lon2, lat2):
    # 空值拦截，避免脏数据导致执行失败
    if None in (lon1, lat1, lon2, lat2):
        return None
    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])
    newlon = lon2 - lon1
    newlat = lat2 - lat1
    haver_formula = np.sin(newlat/2.0)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(newlon/2.0)**2
    dist = 2 * np.arcsin(np.sqrt(haver_formula))
    miles = 3958 * dist 
    # 转换为Python原生浮点类型，匹配Spark的FloatType要求
    return float(miles)

# 注册UDF
udf_haversine = udf(haversine, FloatType())
# 调用UDF计算距离
data = data.withColumn("distance", udf_haversine("long", "lat", "merch_long","merch_lat"))

方案2：使用Spark内置函数实现（更优） #

完全使用Spark原生函数实现逻辑，避免Python UDF的序列化开销，执行效率提升5~10倍，无需依赖worker节点的numpy环境：

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

def haversine_spark(lon1_col, lat1_col, lon2_col, lat2_col):
    lon1_rad = F.radians(lon1_col)
    lat1_rad = F.radians(lat1_col)
    lon2_rad = F.radians(lon2_col)
    lat2_rad = F.radians(lat2_col)
    
    newlon = lon2_rad - lon1_rad
    newlat = lat2_rad - lat1_rad
    
    haver_formula = F.sin(newlat/2.0)**2 + F.cos(lat1_rad) * F.cos(lat2_rad) * F.sin(newlon/2.0)**2
    dist = 2 * F.asin(F.sqrt(haver_formula))
    miles = 3958 * dist
    return miles.cast(FloatType())

# 直接调用函数计算距离，无需注册UDF
data = data.withColumn("distance", haversine_spark("long", "lat", "merch_long","merch_lat"))

注意事项 #

• 若使用Python UDF方案，需保证所有Spark worker节点都已安装numpy依赖
• 传入的四个字段必须为数值类型，若为字符串类型需先通过cast(FloatType())完成类型转换

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