Pyspark 窗口函数与其他列上的过滤器答案

【问题标题】：Pyspark window function with filter on other columnPyspark 窗口函数与其他列上的过滤器
【发布时间】：2019-05-04 17:57:15
【问题描述】：

我有一个包含以下数据的 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)，但没有返回正确的结果。

【问题讨论】：

标签： apache-spark pyspark apache-spark-sql window-functions

【解决方案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 行，此解决方案可能不起作用，请考虑一下

【讨论】：