如何查找 PySpark DataFrame 列中的連續值並替換值

Question

我有以下 Dataframe

ID    Name     Dept
1     John     ABC
2     Rio      BCD
3     Marry    BCD
4     Andy     BCD
5     Smith    PQR
6     Rich     XYZ
7     Lisa     LMN
8     Steve    LMN
9     Ali      STU

我們可以看到在 Dept 列中，BCD 重復了 3 次，而 LMN 重復了 2 次。

現在我需要創建新列 Dept_Updated 並檢查連續值，如果有連續值，只需在最后添加下划線並在下划線后添加數字，如果不是連續值，則保持原樣。

我需要以下格式的 output。

ID    Name     Dept   Dept_Updated
1     John     ABC        ABC
2     Rio      BCD        BCD_1
3     Marry    BCD        BCD_2
4     Andy     BCD        BCD_3
5     Smith    PQR        PQR
6     Rich     XYZ        XYZ
7     Lisa     LMN        LMN_1
8     Steve    LMN        LMN_2
9     Ali      STU        STU

我對 PySpark 很陌生，有什么辦法可以實現上面的 output 這真的很有幫助。

Answer 1

我們需要按 Dept 分區的 window，還需要檢查連續條目。 為此，我建議像下面這樣的東西，它將檢查下一行是否與當前行相同，並且只有 append 的 Rnk 列（計數列）用於具有重復項的條目：

from pyspark.sql import functions as F, Window as W
w = W.orderBy('ID')
w1 = W.partitionBy("Dept").orderBy("Dept")
condition = F.col("Check_Duplicate")| ((F.col("CheckLength")>1) & (F.col("Rnk")==1))

new_df = df.withColumn("Check_Duplicate",F.col("Dept")==F.lag("Dept").over(w))\
.withColumn("Rnk",F.row_number().over(w1))\
.withColumn("CheckLength",F.count("Dept").over(w1))\
.withColumn("Dept_Updated",F.when(condition,F.concat_ws("_",*["Dept","Rnk"]))
                                                   .otherwise(F.col("Dept")))

new_df.select(*df.columns,'Dept_Updated').orderBy("ID").show()

---

Output：

+---+-----+----+------------+
| ID| Name|Dept|Dept_Updated|
+---+-----+----+------------+
|  1| John| ABC|         ABC|
|  2|  Rio| BCD|       BCD_1|
|  3|Marry| BCD|       BCD_2|
|  4| Andy| BCD|       BCD_3|
|  5|Smith| PQR|         PQR|
|  6| Rich| XYZ|         XYZ|
|  7| Lisa| LMN|       LMN_1|
|  8|Steve| LMN|       LMN_2|
|  9|  Ali| STU|         STU|
+---+-----+----+------------+

---

測試表明，如果 Dept 沒有連續重復，則代碼不會 append 行號：

+---+-----+----+------------+
| ID| Name|Dept|Dept_Updated|
+---+-----+----+------------+
|  1| John| ABC|         ABC|
|  2|  Rio| BCD|       BCD_1|
|  3|Marry| BCD|       BCD_2|
|  4| Andy| BCD|       BCD_3|
|  5|Smith| PQR|         PQR|
|  6| Rich| BCD|         BCD| # <-- This entry is repeated but not consecutive
|  7| Lisa| LMN|       LMN_1|
|  8|Steve| LMN|       LMN_2|
|  9|  Ali| STU|         STU|