PySpark实现类似R rleid功能：生成连续变化的分组ID

需求说明 #

按order_col排序后，基于group_col生成连续组ID（rleid）：仅当group_col的值发生变化时生成新的组ID，即使同一group_col值非连续重复出现（如示例最后一行的r1），也需分配新的组ID。

原始DataFrame #

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

have = spark.createDataFrame(
    [('a', 'r1', '1'),
    ('b', 'r1', '2'),
    ('c', 'r1', '3'),
    ('d', 's3', '4'),
    ('e', 's3', '5'),
    ('f', 's4', '6'),
    ('g', 'r1', '7')],
    ['id', 'group_col', 'order_col'])

期望结果 #

want = spark.createDataFrame(
    [('a', 'r1', '1', '1'),
    ('b', 'r1', '2', '1'),
    ('c', 'r1', '3', '1'),
    ('d', 's3', '4', '2'),
    ('e', 's3', '5', '2'),
    ('f', 's4', '6', '3'),
    ('g', 'r1', '7', '4')],
    ['id', 'group_col', 'order_col', 'rleid'])

want.show()
+---+---------+---------+-----+
| id|group_col|order_col|rleid|
+---+---------+---------+-----+
|  a|       r1|        1|    1|
|  b|       r1|        2|    1|
|  c|       r1|        3|    1|
|  d|       s3|        4|    2|
|  e|       s3|        5|    2|
|  f|       s4|        6|    3|
|  g|       r1|        7|    4|
+---+---------+---------+-----+

R语言data.table参考实现 #

library(data.table)
df <- data.table(
    id = letters[1:7],
    group_col = c("r1", "r1", "r1", "s3", "s3", "s4", "r1"),
    order_col = c(1:7)
)
setorder(df, order_col)
df[, `:=` (rleid = rleid(group_col))]
df
   id group_col order_col rleid
1:  a        r1         1     1
2:  b        r1         2     1
3:  c        r1         3     1
4:  d        s3         4     2
5:  e        s3         5     2
6:  f        s4         6     3
7:  g        r1         7     4

PySpark实现方案 #

核心思路是通过窗口函数标记组的起始行，再对标记做累加求和得到连续组ID：

1. 先将order_col转为整数类型，确保排序逻辑正确；
2. 使用lag窗口函数获取前一行的group_col值；
3. 对比当前行与前一行的group_col，不同则标记为1（新组起始），否则标记为0；
4. 对标记列做累加求和，得到最终的rleid。

完整代码：

# 转换order_col为整数，避免字符串排序问题
have = have.withColumn("order_col", F.col("order_col").cast("int"))

# 定义窗口：按order_col排序，全表范围累加
window_spec = Window.orderBy("order_col")

# 生成组起始标记 + 累加得到rleid
result = have.withColumn(
    "is_new_group",
    F.when(F.lag("group_col").over(window_spec) != F.col("group_col"), 1)
     .otherwise(0)
     # 第一行没有前一行，标记为1作为第一个组的起始
     .na.fill(1)
).withColumn(
    "rleid",
    F.sum("is_new_group").over(window_spec.rangeBetween(Window.unboundedPreceding, 0))
).drop("is_new_group")

# 查看结果
result.show()

运行后输出与期望结果完全一致：

+---+---------+---------+-----+
| id|group_col|order_col|rleid|
+---+---------+---------+-----+
|  a|       r1|        1|    1|
|  b|       r1|        2|    1|
|  c|       r1|        3|    1|
|  d|       s3|        4|    2|
|  e|       s3|        5|    2|
|  f|       s4|        6|    3|
|  g|       r1|        7|    4|
+---+---------+---------+-----+

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