PySpark中通过两列自定义函数新增列时UDF报错如何解决

你遇到的PickleException: expected zero arguments for construction of ClassDict (for numpy.dtype)错误，核心原因是自定义Python UDF的返回值是numpy专属数值类型（比如numpy.float64），PySpark的序列化组件无法直接识别numpy自定义类型，只能识别Python原生基础类型、Spark SQL内置数据类型，在任务执行序列化返回值阶段直接失败。
另外原有UDF还有一个隐藏兼容问题：如果features和center是Spark存储的列表/Vector类型，直接用减号计算会触发类型错误，需要先转成numpy数组才能做向量减法。

方案1：修正自定义UDF（兼容原有写法） #

显式声明UDF的返回类型，同时把numpy计算得到的结果转成Python原生float类型，补全数组转换逻辑：

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

@udf(returnType=FloatType())
def dist(feat, cent):
    # 先把输入转成numpy数组再做计算，返回前转成Python原生float
    return float(np.linalg.norm(np.array(feat) - np.array(cent)))

# 调用逻辑和原有写法一致
df.withColumn("distance", dist(col("features"), col("center"))).show()

方案2：使用Spark内置函数（性能最优，推荐） #

Python UDF需要在JVM和Python进程之间做数据序列化传输，性能比Spark内置函数低1~2个数量级，欧氏距离计算完全可以用内置函数实现，从根源避免序列化问题：
如果你的features和center是数组类型（和示例结构一致），直接用SQL内置聚合函数实现：

from pyspark.sql.functions import sqrt, expr

df = df.withColumn(
    "distance",
    sqrt(
        expr("""
        aggregate(
            arrays_zip(features, center),
            0D,
            (acc, pair) -> acc + power(pair.features - pair.center, 2)
        )
        """)
    )
)
df.show()

如果两列是MLlib的Vector类型，可以先转成数组再用上面的逻辑计算，或者直接使用MLlib自带的距离工具。

• 编写PySpark Python UDF时必须显式指定returnType，不要依赖框架自动类型推断，减少类型匹配异常
• UDF中所有numpy、pandas等第三方库计算得到的标量结果，返回前都要转成Python原生的int/float/str/bool等基础类型，否则会触发同类序列化错误
• 数据量较大的场景优先使用Spark内置函数，避免Python UDF带来的性能损耗

内容的提问来源于stack exchange，提问作者ai.jennetta