如何将Pandas自定义函数add_issue_state转换为PySpark实现？

原Pandas函数逻辑拆解 #

先明确你提供的Pandas函数核心逻辑：
针对指定列前缀tr_col，逐行计算{tr_col}_Issue列：

• 第1行：以False作为初始状态，flag_critical = False OR 当前行{tr_col}_critical，最终Issue等价于flag_critical AND NOT {tr_col}_healthy
• 后续行：以上一行的{tr_col}_Issue作为prev_issue，flag_critical = prev_issue OR 当前行{tr_col}_critical，最终Issue同样为flag_critical AND NOT {tr_col}_healthy

这是依赖前一行计算结果的递推逻辑，单纯用F.lag无法直接实现（因为lag只能获取原始数据的前一行值，而非衍生出的Issue值），需要用递归CTE（Common Table Expression）来完成。

PySpark实现方案 #

步骤说明 #

1. 给DataFrame添加有序行号：PySpark是分布式计算，没有天然的行顺序，必须显式指定排序规则（示例用monotonically_increasing_id()，实际业务中建议用时间戳、主键等有意义的排序字段）。
2. 用递归CTE实现递推计算：分为基础分支（处理第一行）和递归分支（关联前一行的计算结果生成当前行的Issue值）。

代码实现 #

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

def add_issue_state_spark(df_tr, tr_col):
    # 1. 添加行号，确保数据处理顺序与原Pandas逻辑一致
    # 替换orderBy中的字段为你的业务排序键（比如时间戳、ID）
    order_window = Window.orderBy(F.monotonically_increasing_id())
    df_with_row = df_tr.withColumn("row_num", F.row_number().over(order_window))
    
    # 2. 定义递归CTE
    # 基础分支：处理第一行
    base_df = df_with_row.filter(F.col("row_num") == 1).withColumn(
        f"{tr_col}_Issue",
        (F.lit(False) | F.col(f"{tr_col}_critical")) & ~F.col(f"{tr_col}_healthy")
    )
    
    # 递归分支：处理后续行，关联前一行的计算结果
    recursive_df = df_with_row.alias("curr").join(
        base_df.alias("prev"),
        F.col("curr.row_num") == F.col("prev.row_num") + 1,
        "rightouter"
    ).select(
        "curr.*",
        (F.col(f"prev.{tr_col}_Issue") | F.col(f"curr.{tr_col}_critical")) & ~F.col(f"curr.{tr_col}_healthy").alias(f"{tr_col}_Issue")
    )
    
    # 合并基础分支与递归分支，移除行号
    result_df = base_df.unionByName(recursive_df).drop("row_num")
    
    return result_df

补充说明 #

如果使用Spark 3.1+版本，也可以用SQL语法的递归CTE，逻辑更直观：

def add_issue_state_spark(df_tr, tr_col):
    order_window = Window.orderBy(F.monotonically_increasing_id())
    df_with_row = df_tr.withColumn("row_num", F.row_number().over(order_window))
    df_with_row.createOrReplaceTempView("temp_df")
    
    result_df = df_tr.sparkSession.sql(f"""
        WITH RECURSIVE issue_cte AS (
            SELECT *,
                   (FALSE OR {tr_col}_critical) AND NOT {tr_col}_healthy AS {tr_col}_Issue
            FROM temp_df
            WHERE row_num = 1
            UNION ALL
            SELECT curr.*,
                   (prev.{tr_col}_Issue OR curr.{tr_col}_critical) AND NOT curr.{tr_col}_healthy AS {tr_col}_Issue
            FROM temp_df curr
            JOIN issue_cte prev ON curr.row_num = prev.row_num + 1
        )
        SELECT * EXCEPT(row_num)
        FROM issue_cte
        ORDER BY row_num
    """)
    
    return result_df

内容的提问来源于stack exchange，提问作者Harsha S Gowda