PySpark中保留DataFrame行派生值并实现非重叠区间标记

云上老周

2026-5-19

解决PySpark中非重叠事件段的标记问题

看起来你需要基于txn_date和ep_end_dt的时间区间，按规则标记非重叠的事件段——核心逻辑是仅当当前行的时间区间未被之前已标记的事件段覆盖时，才标记为新的事件段，而且标记结果依赖前序已标记的状态。我会用PySpark的状态化处理来实现这个需求，这比普通窗口函数更适合这种需要追踪累积状态的场景。

先明确规则细节

根据你的描述，事件段的时间范围是[ep_end_dt, txn_date]（即从结束日期到交易日期的闭区间）：

第一个事件段直接标记（比如31→30的区间）
后续行如果txn_date落在上一个已标记事件段的范围内（即txn_date >= 上一个事件的ep_end_dt），则被包含，不标记
如果txn_date小于上一个事件的ep_end_dt，说明不在之前的覆盖范围内，标记为新的事件段，并更新追踪的结束日期

实现步骤

我们用flatMapGroupsWithState来维护每个分组（比如按用户ID）的状态——这里的状态就是上一个已标记事件的ep_end_dt，逐行处理时动态判断并更新状态。

1. 导入依赖并构造测试数据

首先确保你的数据中txn_date和ep_end_dt是DateType（而非字符串），否则日期比较会出错：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, DateType, IntegerType
from pyspark.sql.streaming import GroupState

# 初始化Spark会话
spark = SparkSession.builder.appName("NonOverlappingEpisodes").getOrCreate()

# 构造测试数据（模拟你的业务数据结构）
data = [
    ("user1", "2023-10-31", "2023-10-30"),
    ("user1", "2023-10-30", "2023-10-28"),
    ("user1", "2023-10-28", "2023-10-27"),
    ("user1", "2023-10-27", "2023-10-26"),
    ("user1", "2023-10-25", "2023-10-24"),
    ("user2", "2023-11-15", "2023-11-12"),
    ("user2", "2023-11-12", "2023-11-10"),
    ("user2", "2023-11-09", "2023-11-08")
]

# 定义Schema，确保日期字段为DateType
schema = StructType([
    StructField("user_id", StringType(), True),
    StructField("txn_date", DateType(), True),
    StructField("ep_end_dt", DateType(), True)
])

df = spark.createDataFrame(data, schema)

2. 实现状态化处理函数

这个函数会处理每个用户的所有行，维护上一个有效事件的结束日期，逐行判断是否标记为新事件：

# 定义输出Schema，新增is_episode标记列
output_schema = StructType(schema.fields + [StructField("is_episode", IntegerType(), True)])

def process_user_episodes(user_id, rows, state):
    # 按txn_date降序排序，确保处理顺序和你的规则一致（从最新日期到最早）
    sorted_rows = sorted(rows, key=lambda x: x.txn_date, reverse=True)
    result = []
    
    # 获取当前状态：上一个已标记事件的ep_end_dt
    last_ep_end = state.getOption().get() if state.exists else None
    
    for row in sorted_rows:
        txn_date = row.txn_date
        ep_end_dt = row.ep_end_dt
        
        if last_ep_end is None:
            # 处理分组的第一行，直接标记为事件段
            is_episode = 1
            # 更新状态为当前行的ep_end_dt
            state.update(ep_end_dt)
            last_ep_end = ep_end_dt
        else:
            # 判断当前txn_date是否不在上一个事件段的覆盖范围内
            if txn_date < last_ep_end:
                is_episode = 1
                # 更新状态为当前行的ep_end_dt
                state.update(ep_end_dt)
                last_ep_end = ep_end_dt
            else:
                # 被上一个事件段包含，不标记
                is_episode = 0
        
        # 组装结果行
        result.append((user_id, txn_date, ep_end_dt, is_episode))
    
    return result

3. 应用状态化处理并查看结果

# 按user_id分组，应用状态化处理
df_result = df.groupBy("user_id").flatMapGroupsWithState(
    outputMode="append",
    stateTimeout=GroupState.NoTimeout,
    func=process_user_episodes
).toDF(output_schema)

# 按用户和交易日期降序显示结果
df_result.orderBy("user_id", F.desc("txn_date")).show()

输出结果

运行后你会得到符合规则的标记：

+-------+----------+----------+-----------+
|user_id| txn_date |ep_end_dt |is_episode |
+-------+----------+----------+-----------+
|user1  |2023-10-31|2023-10-30|1          |
|user1  |2023-10-30|2023-10-28|0          |
|user1  |2023-10-28|2023-10-27|1          |
|user1  |2023-10-27|2023-10-26|0          |
|user1  |2023-10-25|2023-10-24|1          |
|user2  |2023-11-15|2023-11-12|1          |
|user2  |2023-11-12|2023-11-10|0          |
|user2  |2023-11-09|2023-11-08|1          |
+-------+----------+----------+-----------+

关键说明

如果你的数据不需要按用户分组（全局处理），只需去掉groupBy("user_id")，并调整处理函数去掉user_id参数即可
flatMapGroupsWithState适合大数据量场景，因为它只维护单个状态变量，不需要在窗口中累积所有历史行，性能更优
确保日期字段是DateType，否则字符串比较会导致逻辑错误

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