PySpark实现同用户连续False行分组并生成唯一标识

核心思路 #

通过窗口函数标记连续false块的起始点，再对起始点进行累加得到分组编号，同时将true行的分组统一设为0。

具体实现步骤 #

1. 创建测试数据（若已有DataFrame可跳过）

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

spark = SparkSession.builder.appName("FalseGrouping").getOrCreate()

# 构造示例数据
data = [
    (1, True), (1, True), (1, False), (1, False), (1, True), (1, True),
    (2, True), (2, True), (2, False), (2, False), (2, False), (2, True),
    (2, False), (2, False), (2, True), (3, True), (4, True), (5, True),
    (5, False), (5, False)
]
df = spark.createDataFrame(data, ["user_id", "keep"])

2. 标记false分组的起始行

定义按user_id分区的窗口，用lag函数获取上一行的keep值，判断当前行是否是新的false连续块的起点：

# 按用户分区，保证行顺序（若有时间戳等排序字段，替换monotonically_increasing_id）
user_window = Window.partitionBy("user_id").orderBy(F.monotonically_increasing_id())

df = df.withColumn(
    "is_start",
    F.when(
        (F.col("keep") == False) & (F.lag("keep", 1, True).over(user_window) == True),
        1
    ).otherwise(0)
)

3. 生成分组编号

对每个用户的起始标记进行累加，得到false块的分组号，同时将true行的分组设为0：

# 累加起始标记的窗口
group_window = Window.partitionBy("user_id").orderBy(F.monotonically_increasing_id()).rowsBetween(Window.unboundedPreceding, 0)

df = df.withColumn(
    "group",
    F.when(F.col("keep") == True, 0).otherwise(F.sum("is_start").over(group_window))
)

# 清理中间列
df = df.drop("is_start")

4. 查看结果

执行df.show()即可得到目标输出：

+-------+-----+-----+
|user_id| keep|group|
+-------+-----+-----+
|      1| true|    0|
|      1| true|    0|
|      1|false|    1|
|      1|false|    1|
|      1| true|    0|
|      1| true|    0|
|      2| true|    0|
|      2| true|    0|
|      2|false|    1|
|      2|false|    1|
|      2|false|    1|
|      2| true|    0|
|      2|false|    2|
|      2|false|    2|
|      2| true|    0|
|      3| true|    0|
|      4| true|    0|
|      5| true|    0|
|      5|false|    1|
|      5|false|    1|
+-------+-----+-----+

关键说明 #

• 若原数据有明确的排序字段（如时间戳），请将monotonically_increasing_id()替换为该字段，确保行顺序的准确性
• lag("keep", 1, True)中的第三个参数True是为了处理用户的第一行是false的情况，此时默认上一行是true，会被标记为起始行
• 累加窗口的rowsBetween(Window.unboundedPreceding, 0)保证从当前用户的第一行到当前行进行累加，得到正确的分组编号

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