Spark技术问询：如何选取每条QUIT故障前的X条记录

Hey there! Since you're new to Spark, let's walk through this problem step by step—no jargon overload, just clear, actionable steps to get the result you need.

需求回顾 #

你有包含serial_number（设备序列号）、timestamp（时间戳）、code（故障码）、fault_type（故障类型）的日志数据，需要找到每条fault_type = 'QUIT'的记录，并提取该记录之前的X条对应设备的日志记录，最终生成一个每行包含对应QUIT故障前X条code列表的DataFrame。

核心思路 #

要实现这个需求，我们需要利用Spark的窗口函数来按设备分组、按时间排序，然后关联每个QUIT记录与其之前的历史数据。具体分为以下几个关键步骤：

步骤1：数据预处理（确保时间戳格式正确） #

首先要确保timestamp字段是Spark可识别的时间类型，否则排序会出错。如果你的原始数据中timestamp是字符串，先转换为TimestampType：

from pyspark.sql.types import TimestampType
from pyspark.sql.functions import col

# 假设原始DataFrame名为raw_df
processed_df = raw_df.withColumn("timestamp", col("timestamp").cast(TimestampType()))

步骤2：标记QUIT记录并分组 #

我们需要为每个设备的日志按时间排序，然后为每个QUIT记录创建一个"分组ID"，这样就能把每个QUIT和它之前的记录归为一组：

from pyspark.sql.window import Window
from pyspark.sql.functions import sum

# 定义窗口：按设备分区，按时间正序排序
window_partition = Window.partitionBy("serial_number").orderBy("timestamp")

# 标记是否为QUIT记录，并生成分组ID
grouped_df = processed_df.withColumn(
    "is_quit", 
    (col("fault_type") == "QUIT").cast("int")  # QUIT标记为1，其他为0
).withColumn(
    "quit_group", 
    sum("is_quit").over(window_partition)  # 累加标记值，每个QUIT会开启新的分组
)

这里的quit_group会为同一个设备的每一段"从上次QUIT（或起始）到当前QUIT"的日志分配唯一ID，方便后续筛选。

步骤3：为每组内的记录按时间倒序编号 #

为了快速定位QUIT记录之前的X条数据，我们在每个分组内按时间倒序排序，给每条记录编号——这样QUIT记录会是编号1，它之前的记录就是编号2、3...X+1：

from pyspark.sql.functions import row_number

# 定义分组内的窗口：按设备和分组ID分区，时间倒序排序
window_group = Window.partitionBy("serial_number", "quit_group").orderBy(col("timestamp").desc())

numbered_df = grouped_df.withColumn(
    "row_num", 
    row_number().over(window_group)
)

步骤4：筛选并收集目标code列表 #

现在我们只需要筛选出每个分组内编号在2到X+1之间的记录（跳过编号1的QUIT本身），然后按设备和分组ID收集code列表：

X = 5  # 替换为你需要的X值

result_df = numbered_df.filter(col("row_num").between(2, X+1)) \
    .groupBy("serial_number", "quit_group") \
    .agg(collect_list("code").alias("previous_codes")) \
    .drop("quit_group")  # 如果不需要分组ID，可以删除这一行

如果某个QUIT记录之前不足X条数据，collect_list会自动收集所有存在的历史记录，不会报错。

验证结果 #

你可以用result_df.show(truncate=False)查看最终结果，每条数据对应一个设备的某条QUIT故障，previous_codes就是该故障前X条记录的code列表。

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