在pyspark datafram中循环两列时向新列添加值

Question

我有一个pyspark数据帧，其中包含列（除了其他一些列之外）： 每个月有多个ID。每个id的活动状态由amount列确定。如果金额大于0，则激活=1，否则为0。你知道吗

+-----------------------------+---
|id|amount|  dates   | active |
+-----------------------------+---
| X|     0|2019-05-01|    0   |
| X|   120|2019-06-01|    1   |      
| Y|    60|2019-06-01|    1   |
| X|     0|2019-07-01|    0   |
| Y|     0|2019-07-01|    0   |
| Z|    50|2019-06-01|    1   |
| Y|     0|2019-07-01|    0   |
+-----------------------------+---

我要计算并添加的新列是p3mactive。 根据过去三个月的活跃状态计算。 例如：对于id=x，date=2019-08-01，p3mactive=1，因为x在2019-06-01处于活动状态。 如果之前的几个月不存在，则p3m活动=0。如果只有1或2个月，那么p3m active可以简单地计算为max（active（month-1），active（month-2））。基本上是在现有柱子的基础上。你知道吗

+-----------------------------+-----------+
|id|amount|  dates   | active | p3mactive |
+-----------------------------+-----------+
| X|     0|2019-05-01|    0   |     0     |
| X|   120|2019-06-01|    1   |     0     |      
| Y|    60|2019-06-01|    1   |     0     |
| X|     0|2019-07-01|    0   |     1     |
| Y|     0|2019-07-01|    0   |     1     |
| Z|    50|2019-06-01|    1   |     0     |
| Y|     0|2019-07-01|    0   |     1     |
+-----------------------------+-----------+

所以基本上：

1. 05的X为0，在此之前没有月，因此p3mactive为0。你知道吗
2. Y在06中激活，因此p3mactive=1在07中，而p3mactive在06中仍然为0。你知道吗
3. Z只有06的数据，所以06中的p3mactive为0

等等。如果对流量有任何疑问，请告诉我。你知道吗

我想使用pyspark中更好的数据帧操作和函数来实现这一点。 一般来说，我可以很容易地想到如何使用pandas或python来实现这一点，但是我对spark还不熟悉，不能想到一种方法来循环通过id，对于每个给定的月份，然后将前三个月的活动状态选择到max（m1，m2，m3）函数中，如果前几个月不存在，则保留边缘条件。任何帮助都将不胜感激。你知道吗

Answer 1

可以使用^{}和^{}使用^{}函数来执行此操作：

from pyspark.sql.window import Window
from pyspark.sql.functions import when, col, lag

w = Window().partitionBy("id").orderBy("dates")
df = df.withColumn("p3mactive", when(
    (lag(df.active,1).over(w) == 1)| 
    (lag(df.active,2).over(w) == 1) | 
    (lag(df.active,3).over(w) == 1), 1).otherwise(0))

您不能在pyspark数据帧上循环，但可以使用Window跨越它们。可以使用when应用条件，也可以使用lag查看以前的行，使用lead查看将来的行。如果x之前的行不存在，条件的计算结果为false，您将得到一个0，正如您的用例所提到的那样。你知道吗

我希望这有帮助。你知道吗