如何在时序PySpark DataFrame中标记状态序列

针对你的需求，我们可以通过状态分组+序列匹配的方式，完全使用Spark内置函数实现，无需UDF，且支持任意长度的目标序列。以下是具体实现步骤和代码：

核心思路 #

1. 状态分组：将连续相同的目标列值划分为同一个状态组（比如所有连续的2为一组，连续的3为一组），无论每个状态有多少行，都只作为序列中的一个元素。
2. 序列匹配：将状态组的value序列与目标序列对比，找到完全匹配的连续状态组区间。
3. 打标关联：将匹配区间内的所有状态组对应的原始行标记为1，其余为0。

完整代码实现 #

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

def detect_sequence(df, col_name, target_seq):
    # 处理空序列的边界情况
    if not target_seq:
        return df.withColumn("Flag", F.lit(0))
    
    target_len = len(target_seq)
    
    # 步骤1：生成状态分组ID（按id分区，time排序，连续相同value为同一组）
    window_state = Window.partitionBy("id").orderBy("time")
    df_with_group = df.withColumn(
        "is_new_state",
        F.when(F.lag(col_name).over(window_state) != F.col(col_name), 1).otherwise(0)
    ).withColumn(
        "state_group_id",
        F.sum("is_new_state").over(window_state.rangeBetween(Window.unboundedPreceding, 0))
    )
    
    # 步骤2：提取每个状态组的value，并给状态组按顺序编号
    state_groups = df_with_group.select("id", "state_group_id", col_name).distinct()
    window_group_order = Window.partitionBy("id").orderBy("state_group_id")
    state_groups_ordered = state_groups.withColumn(
        "group_seq_pos",
        F.row_number().over(window_group_order) - 1  # 从0开始编号，方便后续匹配
    )
    
    # 步骤3：创建目标序列的映射表（位置+值）
    target_df = spark.createDataFrame(
        [(i, val) for i, val in enumerate(target_seq)],
        schema=["target_pos", "target_val"]
    )
    
    # 步骤4：匹配状态组与目标序列，计算可能的起始匹配位置
    state_groups_with_match = state_groups_ordered.join(
        target_df,
        F.col(col_name) == F.col("target_val"),
        "left"
    ).withColumn(
        "match_pos",
        F.when(F.col("target_val").isNotNull(), F.col("group_seq_pos") - F.col("target_pos")).otherwise(None)
    )
    
    # 步骤5：筛选出完全匹配的起始位置（匹配数量等于目标序列长度）
    valid_starts = state_groups_with_match.groupBy("id", "match_pos").agg(
        F.count("*").alias("match_count")
    ).filter(F.col("match_count") == target_len).select("id", F.col("match_pos").alias("start_group_seq_pos"))
    
    # 步骤6：计算匹配区间的结束位置
    valid_start_end = valid_starts.withColumn(
        "end_group_seq_pos",
        F.col("start_group_seq_pos") + target_len - 1
    )
    
    # 步骤7：筛选出所有属于匹配区间的状态组
    matched_groups = state_groups_ordered.join(
        valid_start_end,
        (F.col("group_seq_pos").between(F.col("start_group_seq_pos"), F.col("end_group_seq_pos"))) &
        (F.col("id") == F.col("id")),
        "inner"
    ).select("id", "state_group_id")
    
    # 步骤8：关联回原始数据，标记Flag
    result_df = df_with_group.join(
        matched_groups,
        on=["id", "state_group_id"],
        left=True
    ).withColumn(
        "Flag",
        F.when(F.col("state_group_id").isNotNull(), 1).otherwise(0)
    ).drop("is_new_state", "state_group_id")
    
    return result_df

测试示例 #

用你提供的DataFrame测试：

# 构建测试数据
df = spark.createDataFrame(
    [
        [1, 100, 1],
        [1, 200, 1],
        [1, 300, 2],
        [1, 400, 2],
        [1, 500, 2],
        [1, 600, 3],
        [1, 700, 3],
        [1, 800, 4],
        [1, 900, 4],
        [1, 1000, 5],
        [1, 1100, 5],
    ],
    schema=["id","time", "value"]
)

# 调用函数检测序列[2,3,4]
result = detect_sequence(df, "value", [2,3,4])
result.orderBy("time").show()

输出结果：

+---+-----+-----+----+
| id| time|value|Flag|
+---+-----+-----+----+
|  1|  100|    1|   0|
|  1|  200|    1|   0|
|  1|  300|    2|   1|
|  1|  400|    2|   1|
|  1|  500|    2|   1|
|  1|  600|    3|   1|
|  1|  700|    3|   1|
|  1|  800|    4|   1|
|  1|  900|    4|   1|
|  1| 1000|    5|   0|
|  1| 1100|    5|   0|
+---+-----+-----+----+

支持的序列类型 #

该函数可以处理任意长度和重复元素的目标序列：

• 比如[2,3,3,3,4,4]：会匹配状态组序列为[2,3,4]的区间（连续的3是一个状态组）
• 比如[2,2,2,2,3,4,4,4]：同样匹配状态组序列[2,3,4]的区间

边界情况处理 #

• 空目标序列：返回所有Flag为0
• 目标序列长度为1：所有对应value的行都标记为1
• 多id场景：每个id独立处理序列匹配
• 多次匹配：如果同一id内出现多次目标序列，所有匹配区间的行都会标记为1

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