PySpark 数据清洗：UDF 与窗口函数详解

在 PySpark 数据清洗中，UDF（用户自定义函数） 和 窗口函数 是两大核心工具，适用于不同场景的数据处理需求。下面通过具体示例说明其用法：

1. UDF 自定义函数

适用场景：当内置函数无法满足复杂清洗逻辑时（如文本解析、自定义计算等）

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

# 创建Spark会话
spark = SparkSession.builder.appName("DataCleaning").getOrCreate()

# 示例数据：用户日志
data = [("user1", "2023-01-01 08:30:45"), 
        ("user2", "2023-01-01T14:25:30Z")]
df = spark.createDataFrame(data, ["user_id", "raw_timestamp"])

# 定义UDF：统一时间格式
def clean_timestamp(ts):
    """处理多种时间格式：替换'T'和'Z'"""
    return ts.replace("T", " ").replace("Z", "") if ts else None

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

# 应用UDF
df_clean = df.withColumn("clean_timestamp", clean_ts_udf("raw_timestamp"))
df_clean.show(truncate=False)

输出：

+-------+---------------------+---------------------+
|user_id|raw_timestamp        |clean_timestamp      |
+-------+---------------------+---------------------+
|user1  |2023-01-01 08:30:45  |2023-01-01 08:30:45  |
|user2  |2023-01-01T14:25:30Z|2023-01-01 14:25:30  |
+-------+---------------------+---------------------+

2. 窗口函数

适用场景：需要跨行计算时（如排名、移动平均、累积值等）

from pyspark.sql import Window
from pyspark.sql.functions import row_number, col

# 示例数据：销售记录
sales_data = [("A", 100), ("A", 150), ("B", 80), ("B", 200)]
df_sales = spark.createDataFrame(sales_data, ["product", "amount"])

# 定义窗口规范：按产品分区，按金额降序
window_spec = Window.partitionBy("product").orderBy(col("amount").desc())

# 窗口函数应用：计算产品内销售排名
df_ranked = df_sales.withColumn("rank", row_number().over(window_spec))
df_ranked.show()

输出：

+-------+------+----+
|product|amount|rank|
+-------+------+----+
|      B|   200|   1|
|      B|    80|   2|
|      A|   150|   1|
|      A|   100|   2|
+-------+------+----+

3. 结合使用案例：异常值检测

场景：识别超出平均销售额2倍标准差的数据点

from pyspark.sql.functions import avg, stddev

# 定义统计窗口（不分区）
stats_window = Window.partitionBy()

# 计算全局统计量
df_stats = df_sales.withColumn("global_avg", avg("amount").over(stats_window)) \
                   .withColumn("global_std", stddev("amount").over(stats_window))

# 定义异常检测UDF
def is_outlier(amount, avg_val, std_val):
    threshold = avg_val + 2 * std_val
    return amount > threshold

# 注册UDF
outlier_udf = udf(is_outlier, BooleanType())

# 检测异常值
result = df_stats.withColumn("is_outlier", 
    outlier_udf("amount", "global_avg", "global_std"))
result.show()

关键选择指南

最佳实践：优先使用内置函数和窗口函数，仅在必要时使用UDF。对于复杂UDF，考虑Pandas UDF（向量化UDF）提升性能。