Spark如何按规则为ID分组下的记录生成EventID列区分不同记录集

• 第一步：先为数据集添加全局行索引，保证后续处理时行顺序和原始数据一致（Spark DataFrame本身无内置行顺序，不做这一步会导致前后行对比结果出错）
• 第二步：按ID字段分区，按行索引排序，使用lag窗口函数取每个ID分组内上一行的Event值
• 第三步：用when条件函数生成递增标记：当前行Event为run且上一行Event为swim时标记为1，其余场景（包括每个ID的第一行）标记为0
• 第四步：按ID字段分区，按行索引排序，对递增标记做累计求和，求和结果加1就是最终的EventID

注意：如果你的原始数据集本身存在可确定行顺序的字段（例如事件发生时间戳），可以直接用该字段作为窗口排序依据，无需额外生成行索引。

• lag()：窗口函数，用于取同分区内前一行的字段值
• sum()：窗口聚合函数，用于累计求和标记值
• when()：条件判断函数，用于生成递增标记
• Window类：用于定义窗口的分区、排序规则
• zipWithIndex()：RDD方法，用于生成全局行序索引

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lag, sum, when
from pyspark.sql.types import StructType, StructField, IntegerType, StringType
from pyspark.sql.window import Window

# 生成示例数据集
spark = SparkSession.builder.appName("EventIDGen").getOrCreate()
data = [
       (1, "run"),
       (1, "run"),
       (1, "run"),
       (1, "swim"),
       (1, "run"),
       (1, "swim"),
       (1, "swim"),
       (1, "run"),
       (1, "run"),
       (1, "run"),
       (1, "swim"),
       (1, "run"),
       (1, "run"),
       (2, "run"),
       (2, "run"),
       (2, "run"),
       (2, "swim"),
       (2, "run"),
       (2, "swim"),
       (2, "swim"),
       (2, "swim"),
       (2, "swim"),
       (2, "run"),
       (3, "run"),
       (3, "run"),
       (3, "swim")
        ]
schema = StructType([
  StructField('ID', IntegerType(),True),
  StructField('Event', StringType(),True)
])
df = spark.createDataFrame(data=data, schema=schema)

# 步骤1：添加全局行索引保证顺序
df_with_index = df.rdd.zipWithIndex().toDF(["row", "idx"])\
    .select(col("idx"), col("row.ID"), col("row.Event"))

# 步骤2：定义窗口规则
window_part = Window.partitionBy("ID").orderBy("idx")
window_cum = Window.partitionBy("ID").orderBy("idx").rowsBetween(Window.unboundedPreceding, 0)

# 步骤3：取上一行Event，生成递增标记
df_with_lag = df_with_index.withColumn("prev_event", lag("Event", 1).over(window_part))
df_with_flag = df_with_lag.withColumn("inc_flag", 
    when((col("Event") == "run") & (col("prev_event") == "swim"), 1).otherwise(0)
)

# 步骤4：累计求和加1得到EventID
df_result = df_with_flag.withColumn("EventID", sum("inc_flag").over(window_cum) + 1)\
    .select("ID", "Event", "EventID")

df_result.show(30)

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