PySpark：按EMP_ID分组后根据状态标记整体组状态的方法

需求说明 #

按EMP_ID分组：若分组内存在Status = 'Not Done'的记录，则该组所有行的overall_status标记为Not Done；若组内无该值，则标记为Done。

输入数据 #

预期输出 #

实现方法 #

方法一：PySpark DataFrame API（窗口函数） #

利用窗口函数直接在分组内计算状态，无需额外关联，性能更优：

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

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

# 构造输入DataFrame
data = [
    (1, "Done"),
    (1, "Not Done"),
    (1, "Done"),
    (2, "Done"),
    (2, "Done")
]
df = spark.createDataFrame(data, ["EMP_ID", "Status"])

# 定义窗口：按EMP_ID分区
emp_window = Window.partitionBy("EMP_ID")

# 计算overall_status
result_df = df.withColumn(
    # 标记组内是否存在Not Done
    "has_not_done",
    F.max(F.when(F.col("Status") == "Not Done", 1).otherwise(0)).over(emp_window)
).withColumn(
    "overall_status",
    F.when(F.col("has_not_done") == 1, "Not Done").otherwise("Done")
).drop("has_not_done")  # 移除中间辅助列

# 查看结果
result_df.show()

方法二：SparkSQL #

提供两种写法，推荐窗口函数版本（性能更优）：

写法1：窗口函数版

# 注册临时视图
df.createOrReplaceTempView("employee_status")

# 执行SQL
result_sql = spark.sql("""
SELECT
    EMP_ID,
    Status,
    CASE
        WHEN MAX(CASE WHEN Status = 'Not Done' THEN 1 ELSE 0 END) OVER (PARTITION BY EMP_ID) = 1
        THEN 'Not Done'
        ELSE 'Done'
    END AS overall_status
FROM employee_status
""")

result_sql.show()

写法2：关联子查询版

result_sql = spark.sql("""
SELECT
    es_main.EMP_ID,
    es_main.Status,
    CASE
        WHEN EXISTS (
            SELECT 1
            FROM employee_status es_sub
            WHERE es_sub.EMP_ID = es_main.EMP_ID
            AND es_sub.Status = 'Not Done'
        ) THEN 'Not Done'
        ELSE 'Done'
    END AS overall_status
FROM employee_status es_main
""")

result_sql.show()

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