Spark 3.3.1（Databricks）中UDF处理整行数据报错求助

环境 #

Databricks + Spark 3.3.1

代码示例 #

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, StructType, StructField, IntegerType

# 定义输入DataFrame的Schema
input_schema = StructType([
    StructField("col1", IntegerType(), True),
    StructField("col2", IntegerType(), True)
])

# 定义接收整行数据的UDF
@udf(returnType=StringType())
def my_udf(row):
    col1 = row.col1
    col2 = row.col2
    result = col1 + col2
    return str(result)

# 创建示例DataFrame
data = [(1, 2), (3, 4), (5, 6)]
df = spark.createDataFrame(data, schema=input_schema)

# 应用UDF时触发错误
result_df = df.withColumn("result", my_udf(df))

报错信息 #

TypeError: Invalid argument, not a string or column: DataFrame[col1: int, col2: int] of type . For column literals, use 'lit', 'array', 'struct' or 'create_map' function.

问题原因 #

调用UDF时直接传入了整个DataFrame对象df，但Spark UDF要求参数必须是列对象，不能直接传递DataFrame。若要传递整行数据，需要先将所有列打包为一个struct类型的列。

解决方法 #

方法1：传递整行数据（匹配需求） #

使用struct函数将所有列打包成结构体列，再传入UDF，原UDF逻辑无需修改：

from pyspark.sql.functions import struct

# 修改UDF调用代码
result_df = df.withColumn("result", my_udf(struct(*df.columns)))

方法2：直接传递目标列（性能更优） #

若UDF仅需特定列，直接传入对应列参数即可，避免序列化整行数据带来的性能损耗：

# 重新定义UDF，接收列参数
@udf(returnType=StringType())
def my_udf(col1, col2):
    result = col1 + col2
    return str(result)

# 调用时传入目标列
result_df = df.withColumn("result", my_udf(df["col1"], df["col2"]))

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