Spark DataFrame调用自定义UDF失败问题求助

错误原因分析 #

• 第一种写法直接调用Python函数key_to_val(sdf.id)：Python函数仅能处理单个值，无法直接接收Spark的Column对象；且注册UDF后可能出现变量名冲突，导致'list' object is not callable错误。
• 第二种写法的问题：变量名错误（stocks_sdf应为sdf），原函数返回字符串"Null"而非Python的None（无法被Spark识别为SQL NULL），UDF调用方式也未正确适配Column对象。

修正后的解决方案 #

方案1：使用已注册的SQL UDF（通过expr调用） #

先修正函数返回值，将字符串"Null"改为None，让Spark识别为SQL NULL：

%%spark
from pyspark.sql.types import StringType

d = {'I':'Ice', 'U':'UN', 'T':'Tick'}

def key_to_val(k):
    if k in d:
        return d[k]
    else:
        return None  # 返回Python None对应SQL NULL

spark.udf.register('key_to_val', key_to_val, StringType())

再通过expr调用注册好的UDF：

from pyspark.sql.functions import expr

sdf = sdf.withColumn('id', expr('key_to_val(id)'))
sdf.show()

方案2：正确使用PySpark UDF #

直接定义并调用PySpark UDF，无需注册为SQL UDF：

%%spark
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType

d = {'I':'Ice', 'U':'UN', 'T':'Tick'}

def key_to_val(k):
    return d.get(k, None)  # 用字典get方法简化逻辑，不存在则返回None

key_to_val_udf = udf(key_to_val, StringType())

# 用col('id')传入列对象，确保变量名正确
sdf = sdf.withColumn('id', key_to_val_udf(col('id')))
sdf.show()

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

避免UDF开销，用create_map和coalesce实现字典映射：

%%spark
from pyspark.sql.functions import create_map, coalesce, lit
from itertools import chain

# 将字典转为键值对列表，用于构建映射表达式
map_expr = create_map([lit(x) for x in chain(*d.items())])

# 用coalesce获取映射值，无匹配则返回NULL
sdf = sdf.withColumn('id', coalesce(map_expr[col('id')], lit(None)))
sdf.show()

最终输出 #

执行任一方案后，将得到符合预期的结果：

+----+------------+--------------+
|id  |date        |Num           |
+----+------------+--------------+
|Ice |2012-01-03  |1             |
|null|2013-01-11  |2             |
+----+------------+--------------+

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