如何在PySpark 3.3.0中将接收DataFrame的复杂函数注册为UDF

首先明确：PySpark UDF无法直接接收整个DataFrame作为参数，UDF的设计目标是处理行级或单/多列的元素级操作，而非全局DataFrame的聚合统计。要借助Spark的分布式处理优势，应该用Spark原生API实现统计逻辑，而非强行注册UDF。

实现步骤与代码示例 #

下面是直接基于Spark API编写的统计函数，能高效计算各列的Null值数量及占比：

1. 导入依赖 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum, when, lit, stack

2. 定义统计函数 #

def calculate_null_stats(df):
    total_rows = df.count()
    # 处理空DataFrame的边界情况
    if total_rows == 0:
        return df.sparkSession.createDataFrame(
            [], 
            schema="column_name string, null_count long, null_ratio double"
        )
    
    # 构造每列的Null计数、占比聚合表达式
    agg_exprs = []
    for col_name in df.columns:
        null_count = sum(when(col(col_name).isNull(), 1).otherwise(0)).alias(f"{col_name}_null_count")
        null_ratio = (null_count / lit(total_rows)).alias(f"{col_name}_null_ratio")
        agg_exprs.extend([null_count, null_ratio])
    
    # 执行聚合计算
    agg_result = df.agg(*agg_exprs)
    
    # 将宽表结果重塑为长表（列名、Null计数、Null占比）
    stack_args = []
    for col_name in df.columns:
        stack_args.extend([f"'{col_name}'", f"{col_name}_null_count", f"{col_name}_null_ratio"])
    
    result_df = agg_result.select(
        stack(len(df.columns), *stack_args).alias("column_name", "null_count", "null_ratio")
    )
    
    return result_df

3. 使用示例 #

# 初始化SparkSession
spark = SparkSession.builder.appName("NullStats").getOrCreate()

# 测试数据
test_data = [
    (1, None, "a"),
    (None, 2, None),
    (3, None, "c")
]
test_df = spark.createDataFrame(test_data, schema=["col1", "col2", "col3"])

# 调用函数
null_stats_df = calculate_null_stats(test_df)
null_stats_df.show()

执行后输出：

+-----------+----------+-------------------+
|column_name|null_count|         null_ratio|
+-----------+----------+-------------------+
|       col1|         1|0.3333333333333333|
|       col2|         2|0.6666666666666666|
|       col3|         1|0.3333333333333333|
+-----------+----------+-------------------+

为什么不用UDF？ #

1. UDF的设计限制：PySpark的标量UDF处理单/多行数据，分组UDF处理分组内数据，均不支持接收整个DataFrame作为输入。
2. 性能损耗：即使通过某种方式将DataFrame传入UDF（比如转为Pandas对象），也会将数据拉到Driver端处理，完全失去Spark的分布式计算优势，甚至引发内存溢出。
3. 原生API更高效：Spark内置的聚合函数经过优化，能在集群分布式执行，性能远优于自定义UDF。

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