如何将DataFrame作为输入传入Spark UDF？适配PySpark执行现有Python函数

你的原函数是针对Pandas DataFrame编写的，要适配PySpark DataFrame，直接利用Spark的Schema信息就能高效筛选数值列，无需遍历行。改造后的函数如下：

def variabletype(spark_df):
    # 定义Spark支持的数值类型集合
    num_col_types = ['int', 'float', 'double', 'long']
    # 从DataFrame的dtypes中筛选数值型列
    numerical_columns = [col_name for col_name, col_type in spark_df.dtypes 
                         if any(num_type in col_type.lower() for num_type in num_col_types)]
    
    # 返回标记各列是否为数值型的DataFrame
    from pyspark.sql.functions import lit
    cols_list_df = spark_df.select(
        *[lit(col_name in numerical_columns).alias(f"{col_name}_is_numerical") 
          for col_name in spark_df.columns]
    )
    return cols_list_df

# 调用示例（df为你的PySpark DataFrame）
result_df = variabletype(df)
result_df.show()

Spark UDF（用户定义函数）是逐行处理数据的设计，它不能直接接收整个DataFrame作为输入——因为Spark会将DataFrame拆分为分布式分区，UDF仅能访问单个分区内的单条记录或指定列，无法获取整个DataFrame的全局信息。

如果你的需求是对每行数据做处理（而非全局列类型分析），正确操作方式如下：

• 传入单个或多个具体列作为UDF参数
• 若需访问整行数据，可先通过struct(*df.columns)将整行打包为结构体传入

示例代码：

from pyspark.sql.functions import udf, struct
from pyspark.sql.types import DoubleType, StringType

# 1. 传入指定列的UDF
@udf(returnType=DoubleType())
def calculate_sum(col1, col2):
    return col1 + col2

# 调用UDF处理指定列
df.withColumn("sum_result", calculate_sum(df["num_col1"], df["num_col2"])).show()

# 2. 处理整行数据的UDF
row_struct = struct(*df.columns)

@udf(returnType=StringType())
def process_row(row):
    # 通过row.col_name访问行内字段
    return f"col1值：{row.col1}, col2值：{row.col2}"

# 调用UDF处理整行
df.withColumn("row_processed", process_row(row_struct)).show()

总结：

• 针对整个DataFrame的全局操作（如列类型筛选），直接编写普通函数处理Spark DataFrame的Schema或元数据即可，无需使用UDF
• 仅当需要逐行/逐列处理数据时，才使用UDF，此时传入具体列或打包后的行结构体作为参数

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