我有一个这样的相框

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'])

字符串
我想基于group_col创建一个组ID列，但仅当组发生变化时。因此，当r1组再次出现时，它将获得与第一个出现的r1不同的组ID。

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|
+---+---------+---------+-----+

型
组ID不必是连续的，我只需要一种方法使每个组都是唯一的。
基本上，我希望在data.table R包中有类似rleid function的东西。等效的R代码是：

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

型
我已经尝试过rank()和dense_rank()超过group_col。

df = have.withColumn("rleid", F.dense_rank().over(Window.orderBy('group_col')))
df.show()
+---+---------+---------+-----+
| id|group_col|order_col|rleid|
+---+---------+---------+-----+
|  a|       r1|        1|    1|
|  b|       r1|        2|    1|
|  c|       r1|        3|    1|
|  g|       r1|        7|    1|
|  d|       s3|        4|    2|
|  e|       s3|        5|    2|
|  f|       s4|        6|    3|
+---+---------+---------+-----+

型
这并没有给予我想要的结果，因为id=g应该有rleid=4。
我也尝试过基于this answer的array_sort()，但不幸的是这也不起作用。

df = (have
    .withColumn("rank", F.array_sort(F.collect_set('group_col').over(Window.orderBy('order_col').rowsBetween(Window.unboundedPreceding, Window.currentRow))))
    .withColumn('rleid', F.expr("array_position(rank, group_col)")))
df.show()
+---+---------+---------+------------+-----+
| id|group_col|order_col|        rank|rleid|
+---+---------+---------+------------+-----+
|  a|       r1|        1|        [r1]|    1|
|  b|       r1|        2|        [r1]|    1|
|  c|       r1|        3|        [r1]|    1|
|  d|       s3|        4|    [r1, s3]|    2|
|  e|       s3|        5|    [r1, s3]|    2|
|  f|       s4|        6|[r1, s3, s4]|    3|
|  g|       r1|        7|[r1, s3, s4]|    1|
+---+---------+---------+------------+-----+

型