使用 PySpark 的窗口函数

Question

我有一个 PySpark Dataframe，我的目标是创建一个 Flag 列，其值取决于 Amount 列的值。 基本上，对于每个组，我想知道在前三个月中是否存在金额大于 0，以及是否存在在这种情况下，所有组的 Flag 列的值将为 1，否则该值将为 0。

我将提供一个示例来更好地说明。

初始 PySpark 数据框：

组	月	金额
A	1	0
A	2	0
A	3	35
A	4	0
A	5	0
B	1	0
B	2	0
C	1	0
C	2	0
C	3	0
C	4	13
D	1	0
D	2	24
D	3	0

Final PySpark 数据框：

组	月份	金额	标志
A	1	0	1
A	2	0	1
A	3	35	1
A	4	0	1
A	5	0	1
B	1	0	0
B	2	0	0
C	1	0	0
C	2	0	0
C	3	0	0
C	4	13	0
D	1	0	1
D	2	24	1
D	3	0	1

基本上，我想要什么是针对每组，将前 3 个月的金额相加。如果该总和大于 0，则该组的所有元素的标志为 1，否则为 0。

Answer 1

您可以通过应用Window 函数来创建flag 列。创建一个伪列，如果满足条件，则该伪列变为 1，然后最后对伪列求和，如果它大于 0，则至少有一次满足条件的行并设置 flag至 1。

from pyspark.sql import functions as F
from pyspark.sql import Window as W

data = [("A", 1, 0, ), 
("A", 2, 0, ), 
("A", 3, 35, ), 
("A", 4, 0, ), 
("A", 5, 0, ), 
("B", 1, 0, ), 
("B", 2, 0, ), 
("C", 1, 0, ), 
("C", 2, 0, ), 
("C", 3, 0, ), 
("C", 4, 13, ), 
("D", 1, 0, ), 
("D", 2, 24, ), 
("D", 3, 0, ), ]

df = spark.createDataFrame(data, ("Group", "Month", "Amount", ))

ws = W.partitionBy("Group").orderBy("Month").rowsBetween(W.unboundedPreceding, W.unboundedFollowing)

criteria = F.when((F.col("Month") < 4) & (F.col("Amount") > 0), F.lit(1)).otherwise(F.lit(0))

(df.withColumn("flag", F.when(F.sum(criteria).over(ws) > 0, F.lit(1)).otherwise(F.lit(0)))
).show()

"""
+-----+-----+------+----+
|Group|Month|Amount|flag|
+-----+-----+------+----+
|    A|    1|     0|   1|
|    A|    2|     0|   1|
|    A|    3|    35|   1|
|    A|    4|     0|   1|
|    A|    5|     0|   1|
|    B|    1|     0|   0|
|    B|    2|     0|   0|
|    C|    1|     0|   0|
|    C|    2|     0|   0|
|    C|    3|     0|   0|
|    C|    4|    13|   0|
|    D|    1|     0|   1|
|    D|    2|    24|   1|
|    D|    3|     0|   1|
+-----+-----+------+----+
"""

Answer 2

您可以将窗口函数与 count 和 when 结合使用。

w = Window.partitionBy('Group')
df = df.withColumn('Flag', F.count(
        F.when((F.col('Month') < 4) & (F.col('Amount') > 0), True)).over(w))
     .withColumn('Flag', F.when(F.col('Flag') > 0, 1).otherwise(0))