如何让Spark的monotonically_increasing_id()遇FAIL时从0重新计数

你之前使用monotonically_increasing_id()无法实现需求的核心原因是：该函数会为全表所有行生成全局唯一、持续递增的ID值，不会在遇到FAIL行时自动重置计数，因此FAIL之后的PASS行会接着之前的全局ID继续递增，和预期效果不符。

要实现「遇到FAIL就重置PASS连续计数」的逻辑，不需要自定义UDF，用Spark原生窗口函数即可实现，性能更高，核心逻辑分两步：

• 先给每一段连续的同值块生成分组标记：每遇到一次FAIL，分组标记累加1，同一段连续的PASS会归属到同一个分组下
• 以分组标记做分区，按行顺序在组内生成连续计数，FAIL行直接赋值为0

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

# 注意：请将下方orderBy中的F.lit(1)替换为你表中能确定行顺序的字段（比如自增ID、时间戳字段）
# 分布式环境下无明确排序规则时，行顺序可能不稳定，会导致计数结果错误
sort_col = F.lit(1) # 替换为实际排序列，例：F.col("id")

# 步骤1：生成分组标记，每遇到1次FAIL分组号+1
cum_window = Window.orderBy(sort_col)
df = df.withColumn(
    "group_tag",
    F.sum(F.when(F.col("COLUMN_2") == "FAIL", 1).otherwise(0)).over(cum_window)
)

# 步骤2：组内生成连续计数
group_window = Window.partitionBy("group_tag").orderBy(sort_col)
df = df.withColumn(
    "COLUMN_3",
    F.when(
        F.col("COLUMN_2") == "FAIL", 0
    ).otherwise(
        # 若需要PASS计数从1开始，去掉末尾的-1即可
        F.row_number().over(group_window) - 1
    )
).drop("group_tag")

df.show()

运行上述代码后，FAIL行的COLUMN_3固定为0，每段连续PASS的计数值会从0开始逐行递增，遇到下一个FAIL时自动重置，和你给出的预期输出逻辑完全一致。

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