为6k行PySpark DataFrame添加特征耗时8小时未完成，求排查

核心问题（6000行数据跑8小时的根源） #

• 错误操作Spark Session：输入的df已经是Spark DataFrame，说明Spark Session早已初始化，函数里重复执行getOrCreate()最后还stop()，不仅会引发资源冲突，还会让返回的DataFrame因Session关闭无法正常计算。
• 多余的全局排序：df.orderBy("row_num").drop("row_num")是全局排序操作，PySpark的全局排序需要全量shuffle数据，哪怕只有6000行也会大幅拖慢速度——窗口函数已经按分区+时间戳完成排序，分区内顺序完全满足需求，这一步纯冗余。
• 分区列名不匹配：函数文档标注输入列是Location，但代码里用kp作为分区列，如果原DataFrame没有kp列，会导致全表作为单个分区执行窗口计算，性能直接暴跌。
• 重复定义Window Spec：每个滚动均值都重新定义几乎一致的Window对象，增加了Spark的查询计划解析成本，应该复用基础Window结构。
• 循环生成lag列的低效写法：循环中每次调用withColumn都会生成新的DataFrame节点，批量生成列能减少计划构建开销。

修正后的代码 #

from pyspark.sql.window import Window
from pyspark.sql import functions as F

def add_rain_features(df, lag_hours=7):
    """
    预处理输入的PySpark DataFrame，为污泥缓冲区预测模型生成特征。

    参数:
    df (pyspark.sql.DataFrame): 输入DataFrame，需包含以下列:
        - Location (站点位置)
        - Timestamp (时间戳)
        - Value (污泥体积，立方米)
        - Min of the sludge buffer (污泥缓冲区最小值)
        - Max of the sludge buffer (污泥缓冲区最大值)
        - Rain 1h (该时间点的小时降雨量，毫米)
        - Value_diff (与上一小时相比的污泥体积变化量，立方米)
    lag_hours (int): 降雨量滞后特征的小时数

    返回:
    pyspark.sql.DataFrame: 包含滞后特征和滚动统计特征的处理后DataFrame
    """
    # 注意：不要在函数内创建/停止Spark Session，输入df已关联活跃Session
    # 使用和文档一致的分区列，若实际列是kp则修改文档描述
    window_spec = Window.partitionBy("Location").orderBy("Timestamp")
    
    # 批量生成滞后特征，避免循环多次构建查询计划
    lag_cols = [
        F.lag(F.col("Rain 1h"), lag).over(window_spec).alias(f'Rain_{lag}h_ago')
        for lag in range(1, lag_hours + 1)
    ]
    df = df.select("*", *lag_cols)
    
    # 复用基础Window，仅修改滚动范围
    base_rolling_window = Window.partitionBy("Location").orderBy("Timestamp")
    df = df.withColumn("Rain_mean_2h", F.avg("Rain 1h").over(base_rolling_window.rowsBetween(-1, 0))) \
           .withColumn("Rain_mean_3h", F.avg("Rain 1h").over(base_rolling_window.rowsBetween(-2, 0))) \
           .withColumn("Rain_mean_4h", F.avg("Rain 1h").over(base_rolling_window.rowsBetween(-3, 0))) \
           .withColumn("Rain_mean_5h", F.avg("Rain 1h").over(base_rolling_window.rowsBetween(-4, 0))) \
           .withColumn("Rain_mean_6h", F.avg("Rain 1h").over(base_rolling_window.rowsBetween(-5, 0)))
    
    return df

额外优化建议 #

• 确认Timestamp列是时间戳类型（而非字符串），否则排序会出错且性能差。
• 如果Location的基数很小（比如只有几个站点），可以提前缓存数据：df = df.cache()，减少重复计算开销。

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