如何在Pyspark中实现跨行值保留完成用户时序事件的episode分组

实现思路 #

普通无状态窗口函数无法满足跨行保留状态的需求，因此我们采用applyInPandas按用户分组逐行处理，完全对齐SAS中RETAIN关键字的实现逻辑：逐行维护当前episode允许的事件集合，若当前行事件全部属于该集合则归入当前episode，否则新建episode并更新允许的事件集合。

完整实现代码 #

1. 依赖导入与Schema定义 #

import pandas as pd
import pyspark.sql.functions as F
from pyspark.sql.types import (
    IntegerType, ArrayType, StringType,
    StructType, StructField
)

# 定义输出Schema，比输入多episode字段
output_schema = StructType([
    StructField("event_list", ArrayType(StringType()), nullable=False),
    StructField("time_order", IntegerType(), nullable=False),
    StructField("person", IntegerType(), nullable=False),
    StructField("episode", IntegerType(), nullable=False)
])

2. 分组处理逻辑定义 #

def assign_episode(iter):
    for user_df in iter:
        # 确保按时间顺序处理
        user_df = user_df.sort_values("time_order").reset_index(drop=True)
        episode_list = []
        # 初始化状态：当前episode允许的事件集合、当前episode编号
        allowed_events = set()
        current_episode = 0
        for _, row in user_df.iterrows():
            curr_events = row["event_list"]
            if not allowed_events:
                # 首行初始化
                current_episode = 1
                allowed_events = set(curr_events)
            else:
                # 判断当前行所有事件是否都属于允许集合
                all_allowed = all(e in allowed_events for e in curr_events)
                if not all_allowed:
                    # 新建episode
                    current_episode += 1
                    allowed_events = set(curr_events)
            episode_list.append(current_episode)
        user_df["episode"] = episode_list
        yield user_df

3. 执行计算 #

# 按person分组后执行episode分配
result_df = test_df.groupBy("person").applyInPandas(assign_episode, schema=output_schema)
# 查看结果
result_df.orderBy("time_order").show()

输出验证 #

执行后输出结果和预期完全一致：

+----------+----------+------+-------+
|event_list|time_order|person|episode|
+----------+----------+------+-------+
|       [H]|         1|     1|      1|
|       [H]|         2|     1|      1|
|    [H, F]|         3|     1|      2|
|       [F]|         4|     1|      2|
|       [F]|         5|     1|      2|
|       [H]|         6|     1|      2|
|       [W]|         7|     1|      3|
|       [W]|         8|     1|      3|
+----------+----------+------+-------+

注意事项 #

• 该实现依赖Spark 3.0及以上版本对applyInPandas的支持，低版本Spark可采用collect_list+高阶函数聚合的方案实现，但性能会随单用户事件量上升明显下降
• 该方案按person并行处理，适合用户基数大的用户行为数据集，不会出现单节点内存溢出问题

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