如何用PySpark高效保留各ID最新UNSATISFACTORY记录后的行

核心思路 #

按ID分组，定位每个ID下TEST_COMPONENT为UNSATISFACTORY的最新记录日期，保留该ID下所有日期不早于这个临界日期的行。

实现步骤与代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import max, col
from pyspark.sql.types import TimestampType

# 初始化Spark会话
spark = SparkSession.builder.appName("FilterUnsatisfactoryHistory").getOrCreate()

# 模拟输入数据
data = [
    (5000, "ENGL", "SATISFACTORY", "2023-01-04T00:00:00.000+11:00"),
    (5000, "ENGL", "SATISFACTORY", "2022-09-07T00:00:00.000+10:00"),
    (5000, "OTHER", "NONE", "2022-09-07T00:00:00.000+10:00"),
    (5000, "ENGL", "UNSATISFACTORY", "2016-05-23T00:00:00.000+10:00"),
    (5000, "OTHER", "NONE", "2016-05-23T00:00:00.000+10:00"),
    (5000, "OTHER", "NONE", "2016-05-23T00:00:00.000+10:00"),
    (5000, "OTHER", "NONE", "2016-02-09T00:00:00.000+11:00"),
    (5000, "OTHER", "NONE", "2016-02-09T00:00:00.000+11:00"),
    (5000, "OTHER", "NONE", "2016-02-09T00:00:00.000+11:00"),
    (5000, "ENGL", "UNSATISFACTORY", "2014-05-29T00:00:00.000+10:00"),
    (5000, "OTHER", "NONE", "2013-09-27T00:00:00.000+10:00")
]

df = spark.createDataFrame(data, ["ID", "TEST_ID", "TEST_COMPONENT", "TEST_DT"])

# 1. 将日期字符串转为Timestamp类型（支持时区比较）
df = df.withColumn("TEST_DT", col("TEST_DT").cast(TimestampType()))

# 2. 计算每个ID的最新不合格记录日期
latest_unsat_dates = df.filter(col("TEST_COMPONENT") == "UNSATISFACTORY") \
                       .groupBy("ID") \
                       .agg(max("TEST_DT").alias("LATEST_UNSAT_DT"))

# 3. 关联原表并过滤数据
# 若ID无不合格记录，保留所有行；如需过滤掉这类ID，去掉`| col("LATEST_UNSAT_DT").isNull()`
result_df = df.join(latest_unsat_dates, on="ID", how="left") \
              .filter(col("TEST_DT") >= col("LATEST_UNSAT_DT") | col("LATEST_UNSAT_DT").isNull()) \
              .drop("LATEST_UNSAT_DT")

# 查看结果
result_df.show(truncate=False)

关键细节说明 #

1. 类型转换：必须将TEST_DT转为TimestampType，Spark会自动解析带时区的字符串，确保日期比较的准确性。
2. 高效聚合：使用groupBy+max计算临界日期，比窗口函数更高效——窗口函数需要对每个ID的全量数据排序，而聚合仅需计算最大值，减少shuffle开销。
3. 边界处理：左连接保证所有原始数据被覆盖，过滤条件中col("LATEST_UNSAT_DT").isNull()用于保留无不合格记录的ID的全部数据，可根据业务需求移除该条件。

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