在 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)

我發現了一些稍微相關的問題，例如this one ，但沒有一個是關於同一場景的。

我認為一個好的起點可能是找到這個答案中的第一個差異

Answer 1

我想出了一個解決方案。 這個想法類似於在 Pandas 中所做的。 我們首先添加一個唯一標識符列，我們將over其上計算滯后列（ 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

解決此問題的另一種方法是在比較滯后后使用 range between 和使用無界先行總和：

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