如何将PySpark DataFrame整行传入UDF并按列名访问值？

问题原因 #

你遇到的空值问题，核心是没有为UDF指定返回数据类型。Spark无法自动推断结构体（struct）传入后的UDF返回类型，导致结果全部为空。

解决方案一：指定UDF返回类型（基础实现） #

通过显式指定UDF的返回类型，结合struct打包整列传入，就能实现通过列名访问Row对象的需求：

from pyspark.sql.functions import col, struct, udf
from pyspark.sql.types import IntegerType

# 初始化DataFrame
df = spark.createDataFrame([(2, 25, 4, 10), (6, 30, 9, 11), (8, 35, 12, 18)], ["A", "B", "C", "D"])

# 定义处理整行的函数
def func(row):
    return row.A + row.B

# 定义UDF时显式指定返回类型
udf_func = udf(func, returnType=IntegerType())

# 将整列打包为struct传入UDF，生成新列
df = df.withColumn("new_column", udf_func(struct([col(x) for x in df.columns])))

# 查看结果
df.show()

执行后会得到正确的计算结果：

+---+---+---+---+----------+
|  A|  B|  C|  D|new_column|
+---+---+---+---+----------+
|  2| 25|  4| 10|        27|
|  6| 30|  9| 11|        36|
|  8| 35| 12| 18|        43|
+---+---+---+---+----------+

解决方案二：使用pandas_udf（更贴近Pandas体验+更高性能） #

如果想要更贴近Pandas apply(axis=1)的写法，推荐使用pandas_udf，它基于Arrow批量处理数据，性能比普通UDF更优：

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType
import pandas as pd

# 初始化DataFrame
df = spark.createDataFrame([(2, 25, 4, 10), (6, 30, 9, 11), (8, 35, 12, 18)], ["A", "B", "C", "D"])

# 定义pandas_udf，输入为Pandas Series（对应整行）
@pandas_udf(IntegerType())
def func_pandas(row: pd.Series) -> int:
    return row["A"] + row["B"]

# 直接传入所有列，语法和Pandas apply类似
df = df.withColumn("new_column", func_pandas(*df.columns))

# 查看结果
df.show()

关键注意事项 #

• 普通UDF必须指定返回类型：Spark对复杂输入（如struct打包的Row）无法自动推断返回类型，必须显式声明。
• 优先选择pandas_udf：对于大数据量场景，pandas_udf的批量处理模式比逐行处理的普通UDF性能高很多。
• PySpark Row对象的访问：传入UDF的Row对象支持.或[]两种访问方式（如row.A等价于row["A"]）。

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