枚举PySpark中连续相等值的组

Question

我试图在PySpark数据帧中用相等的值唯一地标记连续的行。在Pandas中，只需使用就可以做到这一点

s = pd.Series([1,1,1,2,2,1,1,3])
s.ne(s.shift()).cumsum()
0    1
1    1
2    1
3    2
4    2
5    3
6    3
7    4
dtype: int64

在PySpark中如何做到这一点？设置-

from pyspark.sql.types import IntegerType
from pyspark.sql.types import StructType
spark = SparkSession.builder.appName('pandasToSparkDF').getOrCreate()
mySchema = StructType([StructField("col1", IntegerType(), True)])
df_sp = spark.createDataFrame(s.to_frame(), schema=mySchema)

我发现了一些稍微相关的问题，比如这个问题，但没有一个是关于同一场景的。

我认为一个好的起点可以是找到这个答案中的第一个差异

Answer 1

我想出了一个解决方案。这个想法与熊猫身上的做法类似。我们首先添加一个唯一的标识符列over，我们将计算滞后列(这里使用over是必要的，因为它是一个窗口函数(。

然后，我们将感兴趣的列与滞后列进行比较，并将结果的累积和转换为int:

mySchema = StructType([StructField("col1", IntegerType(), True)])
df_sp = spark.createDataFrame(s.to_frame(), schema=mySchema)
win = Window.orderBy("id")
df_sp = (df_sp.withColumn("id", f.monotonically_increasing_id())
.withColumn("col1_shift", f.lag("col1", offset=1, default=0).over(win))
.withColumn("col1_shift_ne", (f.col("col1") != f.col("col1_shift")).cast("int"))
.withColumn("col1_shift_ne_cumsum", f.sum("col1_shift_ne").over(win))
.drop(*['id','col1_shift', 'col1_shift_ne']))
df_sp.show()
---+--------------------+
|col1|col1_shift_ne_cumsum|
+----+--------------------+
|   1|                   1|
|   1|                   1|
|   1|                   1|
|   2|                   2|
|   2|                   2|
|   1|                   3|
|   1|                   3|
|   3|                   4|
+----+--------------------+

Answer 2

解决此问题的另一种方法是使用介于和之间的范围，并在比较滞后后使用无界的前一个和：

from pyspark.sql import functions as F, Window as W
w1 = W.orderBy(F.monotonically_increasing_id())
w2 = W.orderBy(F.monotonically_increasing_id()).rangeBetween(W.unboundedPreceding,0)
cond = F.col("col1") != F.lag("col1").over(w1)
df_sp.withColumn("col1_shift_ne_cumsum",F.sum(F.when(cond,1).otherwise(0)).over(w2)+1).show()

---

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