PySpark DataFrame如何提取按时间排序后ID与location连续同组的首行数据

解法思路 #

核心使用PySpark窗口函数lag()获取同ID分组下前一行的location值，和当前行location对比后筛选符合要求的记录即可，该方案不会因为同一个ID后续重新访问之前去过的location就遗漏记录，完全满足不能直接用dropDuplicates去重的要求。

完整实现步骤 #

1. 导入依赖工具

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

2. 定义窗口规则：按ID分组，组内按已排好序的loggedTime升序排列

window_spec = Window.partitionBy("ID").orderBy("loggedTime")

3. 新增辅助列记录同ID下前一行的location值

df_with_prev_loc = df.withColumn("prev_location", F.lag("location").over(window_spec))

4. 筛选符合要求的记录：保留每个ID的第一条记录（前一行location为空）、当前location和前一行不同的记录，最后删除辅助列

result_df = df_with_prev_loc.filter(
    (F.col("prev_location").isNull()) | (F.col("prev_location") != F.col("location"))
).drop("prev_location")

5. 查看输出结果

result_df.show()

输出结果 #

+---+--------+----------+
| ID|location|loggedTime|
+---+--------+----------+
| 67|     312|  12:09:00|
| 67|     375|  12:23:00|
| 67|     650|  12:26:00|
| 75|     650|  12:27:00|
| 75|     800|  12:30:00|
+---+--------+----------+

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