如何在 Pyspark 的一个窗口中对每 N 行求和？答案

【问题标题】：How to sum every N rows over a Window in Pyspark?如何在 Pyspark 的一个窗口中对每 N 行求和？
【发布时间】：2021-07-07 21:48:25
【问题描述】：

我尝试了不同的窗口函数来做这个练习，但没有成功。谁能想到不同的方法？考虑添加和索引列或 r_number。

year	month	week	item	department	state	sales	sum(sales)_2wks	sum(sales)_4wks	sum(sales)_6wks
2020	1	1	1	1	TX	$100	$250	$680	$1380
2020	1	2	1	1	TX	$150	$250	$680	$1380
2020	1	3	1	1	TX	$200	$430	$680	$1380
2020	1	4	1	1	TX	$230	$430	$680	$1380
2020	1	5	1	1	TX	$400	$700	$1050	$1380
2020	1	6	1	1	TX	$300	$700	$1050	$1380
2020	1	7	1	1	TX	$250	$350	$1050	$1200
2020	1	8	1	1	TX	$100	$350	$1050	$1200
2020	1	9	1	1	TX	$200	$400	$850	$1200
2020	1	10	1	1	TX	$200	$400	$850	$1200
2020	1	11	1	1	TX	$300	$450	$850	$1200
2020	1	11	1	1	TX	$150	$450	$850	$1200

【问题讨论】：

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

【解决方案1】：

您可以分配行号，将它们四舍五入到最接近的 2/4/6 并将其用作分区列以对窗口求和：

from pyspark.sql import functions as F, Window

result = df.withColumn(
    'rn', 
    F.row_number().over(Window.partitionBy('item', 'department', 'state').orderBy('year', 'month', 'week')) - 1
).withColumn(
    'sum_2wks', 
    F.sum('sales').over(Window.partitionBy('item', 'department', 'state', (F.col('rn') / 2).cast('int')))
).withColumn(
    'sum_4wks', 
    F.sum('sales').over(Window.partitionBy('item', 'department', 'state', (F.col('rn') / 4).cast('int')))
).withColumn(
    'sum_6wks', 
    F.sum('sales').over(Window.partitionBy('item', 'department', 'state', (F.col('rn') / 6).cast('int')))
)

result.show()
+----+-----+----+----+----------+-----+-----+---+--------+--------+--------+
|year|month|week|item|department|state|sales| rn|sum_2wks|sum_4wks|sum_6wks|
+----+-----+----+----+----------+-----+-----+---+--------+--------+--------+
|2020|    1|   1|   1|         1|   TX|  100|  0|     250|     680|    1380|
|2020|    1|   2|   1|         1|   TX|  150|  1|     250|     680|    1380|
|2020|    1|   3|   1|         1|   TX|  200|  2|     430|     680|    1380|
|2020|    1|   4|   1|         1|   TX|  230|  3|     430|     680|    1380|
|2020|    1|   5|   1|         1|   TX|  400|  4|     700|    1050|    1380|
|2020|    1|   6|   1|         1|   TX|  300|  5|     700|    1050|    1380|
|2020|    1|   7|   1|         1|   TX|  250|  6|     350|    1050|    1200|
|2020|    1|   8|   1|         1|   TX|  100|  7|     350|    1050|    1200|
|2020|    1|   9|   1|         1|   TX|  200|  8|     400|     850|    1200|
|2020|    1|  10|   1|         1|   TX|  200|  9|     400|     850|    1200|
|2020|    1|  11|   1|         1|   TX|  300| 10|     450|     850|    1200|
|2020|    1|  12|   1|         1|   TX|  150| 11|     450|     850|    1200|
+----+-----+----+----+----------+-----+-----+---+--------+--------+--------+

【讨论】：

【解决方案2】：

上面的解决方案很好，只是如果我们有同一周的多行，row_number 会给人一种错误的印象，因为模数（row_number/2）对于同一周的行应该是相同的。相反，出于显而易见的原因，更喜欢使用 dense_rank() 而不是 row_number() 和 rank() 函数。

val sales_data = Seq((2020,1,1,"1","1","TX",100),
                       (2020,1,1,"1","1","TX",150),
                       (2020,1,2,"1","1","TX",150),
                       (2020,1,3,"1","1","TX",200),
                       (2020,1,4,"1","1","TX",230),
                       (2020,1,5,"1","1","TX",400),
                       (2020,1,6,"1","1","TX",300),
                       (2020,1,7,"1","1","TX",250),
                       (2020,1,8,"1","1","TX",100),
                       (2020,1,9,"1","1","TX",200),
                       (2020,1,10,"1","1","TX",200),
                       (2020,1,11,"1","1","TX",300),
                       (2020,1,11,"1","1","TX",150))
                       
  //Calculate moving sales for 2 weeks, 4 weeks, 6 weeks
                               
  val sales_df = sales_data.toDF("year", "month", "week", "item", "dept", "state", "sale")
//  sales_df.show
  
  sales_df.withColumn("row_no", dense_rank().over(Window.partitionBy("item", "state","dept").orderBy("year", "month", "week"))-1)
          .withColumn("sum(sales)_2wks", sum($"sale").over(Window.partitionBy($"item", $"state",$"dept", ($"row_no"/2).cast("int"))))
          .withColumn("sum(sales)_3wks", sum($"sale").over(Window.partitionBy($"item", $"state",$"dept", ($"row_no"/3).cast("int"))))
          .withColumn("sum(sales)_4wks", sum($"sale").over(Window.partitionBy($"item", $"state",$"dept", ($"row_no"/4).cast("int"))))
          .withColumn("sum(sales)_6wks", sum($"sale").over(Window.partitionBy($"item", $"state",$"dept", ($"row_no"/6).cast("int"))))
          .show

【讨论】：