PySpark中如何按相邻行条件聚合分组

已有按user和time排序完成的DataFrame，原始数据如下：

activity          location  user
0 watch movie        house    A
1 sleep              house    A
2 cardio             gym      A
3 cardio             gym      B
4 buy biscuits       shop     B
5 cardio             gym      B
6 weight training    gym      B

需求说明 #

仅对同一user下相邻行location相同的记录，将activity字段按顺序拼接聚合，而非直接执行df.groupby(['user','location']).activity.collect(",")。最终期望输出如下：

duration                  location  user
watch movie,sleep          house     A
cardio                     gym       A
cardio                     gym       B
buy biscuits               shop      B
cardio, weight training    gym      B

PySpark 实现方案 #

利用Spark窗口函数识别连续相同的分组，再进行聚合操作，具体代码如下：

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

# 1. 定义窗口：按user分区，按time排序（确保顺序和原始数据一致）
window_spec = Window.partitionBy("user").orderBy("time")

# 2. 标记连续相同location的分组：当前行与上一行location不同时，生成1，否则0
df_with_flag = df.withColumn(
    "location_change",
    F.when(
        F.lag("location").over(window_spec) != F.col("location"),
        1
    ).otherwise(0)
)

# 3. 累加标记生成分组ID，同一连续组的ID相同
df_with_group = df_with_flag.withColumn(
    "group_id",
    F.sum("location_change").over(window_spec.rangeBetween(Window.unboundedPreceding, 0))
)

# 4. 按user、location、group_id分组，拼接activity字段
result_df = df_with_group.groupBy("user", "location", "group_id")\
    .agg(F.concat_ws(", ", F.collect_list("activity")).alias("duration"))\
    .select("duration", "location", "user")\
    .orderBy("user", "group_id")  # 保持原始顺序

result_df.show(truncate=False)

代码解释 #

• 窗口定义：partitionBy("user")确保只在同一用户内处理，orderBy("time")保证行顺序和原始数据一致；
• location_change标记：用lag函数获取上一行的location，判断是否变化，生成标记值；
• group_id生成：累加标记值，同一连续相同location的行会得到相同的group_id；
• 聚合拼接：按user、location、group_id分组，用concat_ws拼接activity，最后整理列顺序并排序。

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