根据其他列中定义的间隔索引pyspark列

Question

我有一个pyspark数据框，我想在其中插入基于列pad_change的ID列我的数据框看起来像这样:

|TOOL_ID|pad_change|DATE      |Pad_ID        |
+-------+----------+----------+--------------+
|59628  |1         |2021-05-22|PAD_2021-05-22|
|59628  |0         |2021-05-23|null          |
|59628  |0         |2021-05-23|null          |
|59628  |0         |2021-05-24|null          |
|59628  |0         |2021-05-25|null          |
|59628  |0         |2021-05-26|null          |
|59628  |0         |2021-05-27|null          |
|59628  |0         |2021-05-27|null          |
|59628  |0         |2021-05-28|null          |
|59628  |0         |2021-05-29|null          |
|59628  |1         |2021-06-02|PAD_2021-06-02|
|59628  |0         |2021-06-02|null          |
|59628  |0         |2021-06-02|null          |
|59628  |0         |2021-06-02|null          |
|59628  |0         |2021-06-03|null          |
|59628  |0         |2021-06-04|null          |
|59628  |0         |2021-06-04|null          |
|59628  |0         |2021-06-04|null          |
|59628  |0         |2021-06-04|null          |
|59628  |1         |2021-06-05|PAD_2021-06-05|
|59628  |0         |2021-06-06|null          |
|59628  |0         |2021-06-07|null          |
|59628  |0         |2021-06-08|null          |
|59628  |0         |2021-06-08|null          |
|59628  |0         |2021-06-09|null          |
|59628  |0         |2021-06-09|null          |
|59628  |0         |2021-06-10|null          |

我希望Pad_ID列在pad_change=1时具有相同的值，直到下一次它从0变为1。

预期输出:

|TOOL_ID|pad_change|DATE      |Pad_ID        |
+-------+----------+----------+--------------+
|59628  |1         |2021-05-22|PAD_2021-05-22|
|59628  |0         |2021-05-23|PAD_2021-05-22|
|59628  |0         |2021-05-23|PAD_2021-05-22|
|59628  |0         |2021-05-24|PAD_2021-05-22|
|59628  |0         |2021-05-25|PAD_2021-05-22|
|59628  |0         |2021-05-26|PAD_2021-05-22|
|59628  |0         |2021-05-27|PAD_2021-05-22|
|59628  |0         |2021-05-27|PAD_2021-05-22|
|59628  |0         |2021-05-28|PAD_2021-05-22|
|59628  |0         |2021-05-29|PAD_2021-05-22|
|59628  |1         |2021-06-02|PAD_2021-06-02|
|59628  |0         |2021-06-02|PAD_2021-06-02|
|59628  |0         |2021-06-02|PAD_2021-06-02|
|59628  |0         |2021-06-02|PAD_2021-06-02|
|59628  |0         |2021-06-03|PAD_2021-06-02|
|59628  |0         |2021-06-04|PAD_2021-06-02|
|59628  |0         |2021-06-04|PAD_2021-06-02|
|59628  |0         |2021-06-04|PAD_2021-06-02|
|59628  |0         |2021-06-04|PAD_2021-06-02|
|59628  |1         |2021-06-05|PAD_2021-06-05|
|59628  |0         |2021-06-06|PAD_2021-06-05|
|59628  |0         |2021-06-07|PAD_2021-06-05|
|59628  |0         |2021-06-08|PAD_2021-06-05|
|59628  |0         |2021-06-08|PAD_2021-06-05|
|59628  |0         |2021-06-09|PAD_2021-06-05|
|59628  |0         |2021-06-09|PAD_2021-06-05|
|59628  |0         |2021-06-10|PAD_2021-06-05|

在Pyspark中有办法做到这一点吗?

Answer 1

您可以尝试下面的解决方案，创建2个窗口，一个用于为每个组创建pad_change的累积和。然后，您可以使用该助手列和Tool_ID列创建一个新窗口，将聚合作为first。最后删除helper列(Win_)

from pyspark.sql import functions as F, Window as W
w = W.partitionBy("TOOL_ID").orderBy("DATE")
.rangeBetween(W.unboundedPreceding,0)
w1 = W.partitionBy("TOOL_ID","Win_").orderBy("TOOL_ID")
out = (df.withColumn("Win_",F.sum("pad_change").over(w))
.withColumn("Pad_ID",F.first("Pad_ID").over(w1)).drop("Win_"))

---

out.show(30)
+---+-------+----------+----------+--------------+
|   |TOOL_ID|pad_change|      DATE|        Pad_ID|
+---+-------+----------+----------+--------------+
|   |  59628|         1|2021-05-22|PAD_2021-05-22|
|   |  59628|         0|2021-05-23|PAD_2021-05-22|
|   |  59628|         0|2021-05-23|PAD_2021-05-22|
|   |  59628|         0|2021-05-24|PAD_2021-05-22|
|   |  59628|         0|2021-05-25|PAD_2021-05-22|
|   |  59628|         0|2021-05-26|PAD_2021-05-22|
|   |  59628|         0|2021-05-27|PAD_2021-05-22|
|   |  59628|         0|2021-05-27|PAD_2021-05-22|
|   |  59628|         0|2021-05-28|PAD_2021-05-22|
|   |  59628|         0|2021-05-29|PAD_2021-05-22|
|   |  59628|         1|2021-06-02|PAD_2021-06-02|
|   |  59628|         0|2021-06-02|PAD_2021-06-02|
|   |  59628|         0|2021-06-02|PAD_2021-06-02|
|   |  59628|         0|2021-06-02|PAD_2021-06-02|
|   |  59628|         0|2021-06-03|PAD_2021-06-02|
|   |  59628|         0|2021-06-04|PAD_2021-06-02|
|   |  59628|         0|2021-06-04|PAD_2021-06-02|
|   |  59628|         0|2021-06-04|PAD_2021-06-02|
|   |  59628|         0|2021-06-04|PAD_2021-06-02|
|   |  59628|         1|2021-06-05|PAD_2021-06-05|
|   |  59628|         0|2021-06-06|PAD_2021-06-05|
|   |  59628|         0|2021-06-07|PAD_2021-06-05|
|   |  59628|         0|2021-06-08|PAD_2021-06-05|
|   |  59628|         0|2021-06-08|PAD_2021-06-05|
|   |  59628|         0|2021-06-09|PAD_2021-06-05|
|   |  59628|         0|2021-06-09|PAD_2021-06-05|
|   |  59628|         0|2021-06-10|PAD_2021-06-05|
+---+-------+----------+----------+--------------+