如何在PySpark中为连续相等值的行生成唯一标签？

这个需求在PySpark里完全可以实现，核心思路和你在Pandas里用的「判断与前一行是否不同+累计求和」是一致的，只不过要借助PySpark的窗口函数来完成，同时要注意Spark分布式特性带来的顺序问题。下面是具体的实现步骤：

首先，补充导入我们需要的函数和窗口模块：

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

关键前提：保证数据顺序 #

Spark的DataFrame本身没有天然的行顺序，而我们的需求完全依赖原始数据的顺序，所以第一步必须给DataFrame添加一个自增行号列，用来固定顺序：

# 添加自增行号，确保原始数据的顺序不会被打乱
df_sp = df_sp.withColumn("row_id", F.monotonically_increasing_id())

实现连续值分组标签 #

接下来用窗口函数完成核心逻辑：

1. 用lag函数获取当前行的前一行col1值
2. 判断当前行与前一行的值是否不同（第一行没有前一行，直接标记为新分组）
3. 对「是否为新分组」的标记做累计求和，得到唯一分组标签

完整代码如下：

# 定义按行号排序的窗口，保证计算前一行值时的顺序正确
window = Window.orderBy("row_id")

# 计算分组标签
result_df = df_sp.withColumn(
    # 标记是否为新分组：第一行或与前一行值不同则为1，否则0
    "is_new_group",
    F.when(
        F.coalesce(F.lag("col1").over(window) != F.col("col1"), F.lit(True)),
        1
    ).otherwise(0)
).withColumn(
    # 对标记列累计求和，得到和Pandas cumsum一致的分组标签
    "group_label",
    F.sum("is_new_group").over(window.rangeBetween(Window.unboundedPreceding, 0))
).select("col1", "group_label")  # 只保留需要的列

验证结果 #

执行result_df.show()会得到和你Pandas示例完全一致的结果：

+----+-----------+
|col1|group_label|
+----+-----------+
|   1|          1|
|   1|          1|
|   1|          1|
|   2|          2|
|   2|          2|
|   1|          3|
|   1|          3|
|   3|          4|
+----+-----------+

补充说明 #

• 关于行号：如果你的数据本身有可以用来排序的业务字段（比如时间戳、业务自增ID），可以直接用那个字段排序，不用额外添加row_id
• 简化写法：也可以把逻辑合并成一个withColumn调用，让代码更紧凑：

  result_df = df_sp.withColumn("row_id", F.monotonically_increasing_id()) \
      .withColumn(
          "group_label",
          F.sum(
              F.when(
                  F.coalesce(F.lag("col1").over(Window.orderBy("row_id")) != F.col("col1"), F.lit(True)),
                  1
              ).otherwise(0)
          ).over(Window.orderBy("row_id").rangeBetween(Window.unboundedPreceding, 0))
      ) \
      .select("col1", "group_label")
  

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