Question

我有一个pyspark数据框，其中包含以下数据：

| y | date       | amount| id |
 ----------------------------- 
| 1 | 2017-01-01 | 10    | 1  |
| 0 | 2017-01-01 | 2     | 1  |
| 1 | 2017-01-02 | 20    | 1  |
| 0 | 2017-01-02 | 3     | 1  |
| 1 | 2017-01-03 | 2     | 1  |
| 0 | 2017-01-03 | 5     | 1  |

我想应用窗口函数，但仅将y == 1的列应用sum聚合函数，但仍保留其他列。我要应用的窗口是：

w = Window \
        .partitionBy(df.id) \
        .orderBy(df.date.asc()) \
        .rowsBetween(Window.unboundedPreceding, -1)

结果数据帧如下：

| y | date       | amount| id | sum |
 ----------------------------------- 
| 1 | 2017-01-01 | 10    | 1  | 0   |
| 0 | 2017-01-01 | 2     | 1  | 0   |
| 1 | 2017-01-02 | 20    | 1  | 10  | // =10 (considering only the row with y==1)
| 0 | 2017-01-02 | 3     | 1  | 10  | // same as above
| 1 | 2017-01-03 | 2     | 1  | 30  | // =10+20
| 0 | 2017-01-03 | 5     | 1  | 30  | // same as above

这是否可行？

我尝试使用sum(when(df.y==1, df.amount)).over(w)，但未返回正确的结果。

Answer 1

实际上，使用一个窗口功能很难处理它。我认为您应该首先创建一些虚拟列以计算总和列。您可以在下面找到我的解决方案。

>>> from pyspark.sql.window import Window
>>> import pyspark.sql.functions as F
>>> 
>>> df.show()
+---+----------+------+---+
|  y|      date|amount| id|
+---+----------+------+---+
|  1|2017-01-01|    10|  1|
|  0|2017-01-01|     2|  1|
|  1|2017-01-02|    20|  1|
|  0|2017-01-02|     3|  1|
|  1|2017-01-03|     2|  1|
|  0|2017-01-03|     5|  1|
+---+----------+------+---+

>>> 
>>> df = df.withColumn('c1', F.when(F.col('y')==1,F.col('amount')).otherwise(0))
>>> 
>>> window1 = Window.partitionBy(df.id).orderBy(df.date.asc()).rowsBetween(Window.unboundedPreceding, -1)
>>> df = df.withColumn('c2', F.sum(df.c1).over(window1)).fillna(0)
>>> 
>>> window2 = Window.partitionBy(df.id).orderBy(df.date.asc())
>>> df = df.withColumn('c3', F.lag(df.c2).over(window2)).fillna(0)
>>> 
>>> df = df.withColumn('sum', F.when(df.y==0,df.c3).otherwise(df.c2))
>>> 
>>> df = df.select('y','date','amount','id','sum')
>>> 
>>> df.show()
+---+----------+------+---+---+                                                 
|  y|      date|amount| id|sum|
+---+----------+------+---+---+
|  1|2017-01-01|    10|  1|  0|
|  0|2017-01-01|     2|  1|  0|
|  1|2017-01-02|    20|  1| 10|
|  0|2017-01-02|     3|  1| 10|
|  1|2017-01-03|     2|  1| 30|
|  0|2017-01-03|     5|  1| 30|
+---+----------+------+---+---+

如果每天有y = 1或y = 0行，请考虑使用此解决方案

Pyspark窗口功能，另一列带有过滤器

1 个答案: