PySpark中提取连续重复状态行的首个记录

问题描述 #

需要按created_when日期排序后，提取连续重复状态行的首个记录（状态可重复但不能连续多行重复），输入输出如下：

输入表：

期望输出表：

解决方案 #

核心思路是用**窗口函数lag()**对比当前行和前一行的status，筛选出状态发生变化（或为分组首行）的记录，具体实现步骤如下：

1. 导入所需模块：

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

2. 定义窗口规则：按GP分区（支持多分组场景），按created_when升序排序：

window_spec = Window.partitionBy("GP").orderBy("created_when")

3. 添加辅助列，获取前一行的status值：

df_with_prev_status = df.withColumn(
    "prev_status",
    lag(col("status")).over(window_spec)
)

4. 筛选目标记录：保留分组第一行（prev_status为null），或者当前行status与前一行不同的记录：

result_df = df_with_prev_status.filter(
    (col("prev_status").isNull()) | (col("status") != col("prev_status"))
).drop("prev_status")

5. 查看结果：

result_df.show()

代码解释 #

• lag(col("status")).over(window_spec)：在每个GP分组内，按created_when排序后，获取当前行的上一行status值，分组的第一行没有上一行，所以值为null。
• 过滤条件：要么是分组的第一行，要么当前行和上一行状态不同，这样既保留了连续相同状态的首个记录，也允许非连续的重复状态存在（比如输出里的B出现了两次，因为不是连续的）。

完整可运行代码 #

假设你还未构建输入DataFrame，可直接运行以下完整代码：

from pyspark.sql import SparkSession
from pyspark.sql import Window
from pyspark.sql.functions import lag, col

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

# 构建输入数据
data = [
    ("A", "2022-10-10", "A1"),
    ("B", "2022-10-12", "A1"),
    ("B", "2022-10-13", "A1"),
    ("C", "2022-10-13", "A1"),
    ("C", "2022-10-14", "A1"),
    ("B", "2022-12-15", "A1"),
    ("C", "2022-12-16", "A1"),
    ("D", "2022-12-17", "A1"),
    ("A", "2022-12-18", "A1")
]

df = spark.createDataFrame(data, ["status", "created_when", "GP"])

# 定义窗口并处理数据
result_df = df.withColumn("prev_status", lag(col("status")).over(window_spec)) \
              .filter((col("prev_status").isNull()) | (col("status") != col("prev_status"))) \
              .drop("prev_status")

# 展示结果
result_df.show()

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