如何基于返回字典的UDF为PySpark DataFrame动态新增列？

问题背景 #

现有PySpark DataFrame，需要通过返回字典的UDF动态新增列，对应字典的键为列名，值为列值。示例代码如下：

基础DataFrame创建 #

# 导入包
from pyspark.sql import SparkSession
from pyspark.sql.types import MapType, StringType, IntegerType
from pyspark.sql.functions import udf, F, struct

# 创建SparkSession
spark = SparkSession.builder.appName('example').getOrCreate()

# 创建数据与DataFrame
data = [('John',25),('Smith',30),('Adam',58),('Henry',45)]
df = spark.createDataFrame(data, schema = ['Name', 'Age'])

返回字典的UDF定义与注册 #

def custom_udf(row, param1, param2):
    # 原函数逻辑（此处处理row的方式可能引发错误）
    ...
    return {
        "col1":0,
        "col2":1
    }

# 注册UDF（假设amodel和vectorizer是已定义的参数）
udf_output = udf(lambda row: custom_udf(row, param1=amodel, param2=vectorizer), MapType(StringType(), IntegerType()))

尝试的错误代码 #

df_output = df.withColumn("new_columns", udf_output(F.struct([df[x] for x in df.columns])))
for key, value in df_output.select("new_columns").collect().items(): 
       df_output = df_output.withColumn(key, F.lit(value))

运行后报错：An exception was thrown from a UDF: 'TypeError: sequence item 5: expected str instance, int found'

错误原因分析 #

1. UDF内部处理Row对象错误：传递给custom_udf的row是Spark的Row对象，而非字符串或序列。如果函数内将其当作序列操作（如索引取值、直接拼接字符串），会因Age列是整数类型引发类型不匹配错误。
2. 遍历逻辑错误：df_output.select("new_columns").collect()返回的是Row对象的列表，而非字典，调用.items()会直接报错；同时即使能遍历，F.lit(value)会把Driver端的单个值硬编码到所有行，逻辑完全错误。

正确解决方案 #

步骤1：修复UDF内部逻辑 #

确保正确处理Row对象，通过属性访问列值：

def custom_udf(row, param1, param2):
    # 正确获取Row中的列值
    name = row.Name
    age = row.Age
    # 加入你的业务逻辑（利用param1和param2）
    # 示例逻辑：col1为年龄+10，col2为名字长度
    col1_val = age + 10
    col2_val = len(name)
    return {
        "col1": col1_val,
        "col2": col2_val
    }

步骤2：添加存储Map的中间列 #

# 传递所有列组成的struct给UDF
df_output = df.withColumn("new_columns", udf_output(struct(*df.columns)))

步骤3：从Map中动态提取新列 #

方式1：已知字典键的情况下直接提取

df_output = df_output.withColumn("col1", df_output["new_columns"]["col1"])
df_output = df_output.withColumn("col2", df_output["new_columns"]["col2"])

方式2：动态获取键（未知键名时）

# 从第一条数据中提取Map的键（确保所有行返回的字典键一致）
sample_map = df_output.select("new_columns").first()["new_columns"]
keys = sample_map.keys()

# 遍历键，逐个添加新列
for key in keys:
    df_output = df_output.withColumn(key, df_output["new_columns"][key])

步骤4：（可选）删除中间列 #

df_output = df_output.drop("new_columns")

完整运行示例 #

# 导入包
from pyspark.sql import SparkSession
from pyspark.sql.types import MapType, StringType, IntegerType
from pyspark.sql.functions import udf, struct

# 创建SparkSession
spark = SparkSession.builder.appName('example').getOrCreate()

# 创建数据与DataFrame
data = [('John',25),('Smith',30),('Adam',58),('Henry',45)]
df = spark.createDataFrame(data, schema = ['Name', 'Age'])

# 模拟业务参数
amodel = None
vectorizer = None

# 修复后的UDF
def custom_udf(row, param1, param2):
    name = row.Name
    age = row.Age
    return {
        "col1": age + 10,
        "col2": len(name)
    }

# 注册UDF
udf_output = udf(lambda row: custom_udf(row, param1=amodel, param2=vectorizer), MapType(StringType(), IntegerType()))

# 添加中间列
df_output = df.withColumn("new_columns", udf_output(struct(*df.columns)))

# 动态提取新列
sample_map = df_output.select("new_columns").first()["new_columns"]
for key in sample_map.keys():
    df_output = df_output.withColumn(key, df_output["new_columns"][key])

# 删除中间列
df_output = df_output.drop("new_columns")

# 查看结果
df_output.show()

运行结果：

+-----+---+----+----+
| Name|Age|col1|col2|
+-----+---+----+----+
| John| 25|  35|   4|
|Smith| 30|  40|   5|
| Adam| 58|  68|   4|
|Henry| 45|  55|   5|
+-----+---+----+----+

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