PySpark 数据清洗：UDF 与窗口函数应用指南

在 PySpark 数据清洗中，用户自定义函数（UDF） 和 窗口函数 是处理复杂逻辑的核心工具。下面通过结构化示例说明其应用场景与实现方法。

1. UDF 自定义函数

作用：扩展原生函数能力，实现自定义转换逻辑。
示例场景：清洗文本字段，提取有效信息。

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 定义UDF：提取邮箱域名
def extract_domain(email):
    return email.split("@")[-1] if "@" in email else None

# 注册UDF（指定返回类型）
extract_domain_udf = udf(extract_domain, StringType())

# 应用UDF
df_clean = df.withColumn("email_domain", extract_domain_udf("email"))

关键点：

• 使用 udf() 包装 Python 函数，声明返回数据类型（如 StringType()）
• 避免复杂 UDF（影响性能），优先使用内置函数

2. 窗口函数

作用：在数据分区内执行聚合计算，保留原始行。
示例场景：计算用户消费金额的移动平均。

from pyspark.sql.window import Window
from pyspark.sql.functions import avg, col

# 定义窗口：按用户ID分区，按时间戳排序
window_spec = Window.partitionBy("user_id").orderBy("timestamp").rowsBetween(-2, 0)

# 应用窗口函数：计算3行移动平均
df_result = df.withColumn("moving_avg", avg(col("amount")).over(window_spec))

窗口参数解析：

3. UDF + 窗口函数联合应用

场景：识别异常交易（金额超过分区内均值3个标准差）

# 步骤1：用窗口函数计算分区统计量
window_stats = Window.partitionBy("user_id")
df = df.withColumn("mean_amount", avg("amount").over(window_stats)) \
       .withColumn("std_amount", stddev("amount").over(window_stats))

# 步骤2：定义UDF标记异常
def is_anomaly(amount, mean, std):
    return (abs(amount - mean) > 3 * std) if std > 0 else False

is_anomaly_udf = udf(is_anomaly, BooleanType())

# 步骤3：应用联合逻辑
df_result = df.withColumn("is_anomaly", is_anomaly_udf("amount", "mean_amount", "std_amount"))

性能优化建议

1. UDF 替代方案：
   • 优先使用 pyspark.sql.functions 内置函数（如 regexp_extract()）
   • 复杂逻辑改用 Pandas UDF（向量化计算）
2. 窗口函数优化：
   • 限制窗口范围（避免全分区扫描）
   • 对排序字段建立索引
3. 执行计划检查：

   df_result.explain()  # 查看执行计划，避免全表扫描
   

关键公式：移动平均计算可表示为 $$ \bar{x}t = \frac{1}{n} \sum{i=t-k}^{t} x_i $$ 其中 $k$ 为窗口大小。