Spark rolling window with multiple occurrence in same date

XBJ 0
Viewed 31

Input:

enter image description here

Request: Would like to calculate 3 months rolling sum, avg across time. But, there are two rows for "2022-07-01". Would like to get result for both row.

Expected Output:

enter image description here

pyspark
1 Answers

If you can use rowsBetween instead of rangeBetween, you can assign a row number to each of the dates in an order and then use that row number in the sum and avg window.

Below is an example.

data_sdf. \
    withColumn('rn', func.row_number().over(wd.partitionBy().orderBy('ym'))). \
    withColumn('sum_val_roll3m', 
               func.sum('val').over(wd.partitionBy().orderBy('rn').rowsBetween(-2, 0))
               ). \
    withColumn('mean_val_roll3m', 
               func.avg('val').over(wd.partitionBy().orderBy('rn').rowsBetween(-2, 0))
               ). \
    show()

# +----------+---+---+--------------+-----------------+
# |        ym|val| rn|sum_val_roll3m|  mean_val_roll3m|
# +----------+---+---+--------------+-----------------+
# |2022-03-01|  8|  1|             8|              8.0|
# |2022-04-01|  7|  2|            15|              7.5|
# |2022-05-01|  7|  3|            22|7.333333333333333|
# |2022-06-01| 10|  4|            24|              8.0|
# |2022-07-01|  4|  5|            21|              7.0|
# |2022-07-01|  1|  6|            15|              5.0|
# +----------+---+---+--------------+-----------------+
samkart0
Related