title: How to calculate rolling sum with varying window sizes in PySpark
body: <p>I have a spark dataframe that contains sales prediction data for some products in some stores over a time period. How do I calculate the rolling sum of Predictions for a window size of next N values?</p>
  
  <p><strong>Input Data</strong></p>
  
  <pre><code>+-----------+---------+------------+------------+---+
  | ProductId | StoreId |    Date    | Prediction | N |
  +-----------+---------+------------+------------+---+
  |         1 |     100 | 2019-07-01 | 0.92       | 2 |
  |         1 |     100 | 2019-07-02 | 0.62       | 2 |
  |         1 |     100 | 2019-07-03 | 0.89       | 2 |
  |         1 |     100 | 2019-07-04 | 0.57       | 2 |
  |         2 |     200 | 2019-07-01 | 1.39       | 3 |
  |         2 |     200 | 2019-07-02 | 1.22       | 3 |
  |         2 |     200 | 2019-07-03 | 1.33       | 3 |
  |         2 |     200 | 2019-07-04 | 1.61       | 3 |
  +-----------+---------+------------+------------+---+
  </code></pre>
  
  <p><strong>Expected Output Data</strong></p>
  
  <pre><code>+-----------+---------+------------+------------+---+------------------------+
  | ProductId | StoreId |    Date    | Prediction | N |       RollingSum       |
  +-----------+---------+------------+------------+---+------------------------+
  |         1 |     100 | 2019-07-01 | 0.92       | 2 | sum(0.92, 0.62)        |
  |         1 |     100 | 2019-07-02 | 0.62       | 2 | sum(0.62, 0.89)        |
  |         1 |     100 | 2019-07-03 | 0.89       | 2 | sum(0.89, 0.57)        |
  |         1 |     100 | 2019-07-04 | 0.57       | 2 | sum(0.57)              |
  |         2 |     200 | 2019-07-01 | 1.39       | 3 | sum(1.39, 1.22, 1.33)  |
  |         2 |     200 | 2019-07-02 | 1.22       | 3 | sum(1.22, 1.33, 1.61 ) |
  |         2 |     200 | 2019-07-03 | 1.33       | 3 | sum(1.33, 1.61)        |
  |         2 |     200 | 2019-07-04 | 1.61       | 3 | sum(1.61)              |
  +-----------+---------+------------+------------+---+------------------------+
  </code></pre>
  
  <p>There are lots of questions and answers to this problem in Python but I couldn't find any in PySpark. </p>
  
  <p><strong>Similar Question 1</strong> <br>
  There is a similar question <a href="https://stackoverflow.com/questions/46767807/how-to-calculate-rolling-median-in-pyspark-using-window">here</a> but in this one frame size is fixed to 3. In the provided answer rangeBetween function is used and it is only working with fixed sized frames so I cannot use it for varying sizes.</p>
  
  <p><strong>Similar Question 2</strong> <br>
  There is also a similar question <a href="https://stackoverflow.com/questions/39522074/spark-lag-function-with-parameter-as-dynamic">here</a>. In this one, writing cases for all possible sizes is suggested but it is not applicable for my case since I don't know how many distinct frame sizes I need to calculate.  </p>
  
  <p><strong>Solution attempt 1</strong> <br>
  I've tried to solve the problem using a pandas udf:</p>
  
  <pre><code>rolling_sum_predictions = predictions.groupBy('ProductId', 'StoreId').apply(calculate_rolling_sums)
  </code></pre>
  
  <p><em>calculate_rolling_sums</em> is a pandas udf where I solve the problem in python. This solution works with a small amount of test data. However, when the data gets bigger (in my case, the input df has around 1B rows), calculations take so long.</p>
  
  <p><strong>Solution attempt 2</strong> <br>
  I have used a workaround of the answer of <em>Similar Question 1</em> above. I've calculated the biggest possible N, created the list using it and then calculate the sum of predictions by slicing the list. </p>
  
  <pre><code>predictions = predictions.withColumn('DayIndex', F.rank().over(Window.partitionBy('ProductId', 'StoreId').orderBy('Date')))
  
  # find the biggest period
  biggest_period = predictions.agg({"N": "max"}).collect()[0][0]
  
  # calculate rolling predictions starting from the DayIndex
  w = (Window.partitionBy(F.col("ProductId"), F.col("StoreId")).orderBy(F.col('DayIndex')).rangeBetween(0, biggest_period - 1))
  rolling_prediction_lists = predictions.withColumn("next_preds", F.collect_list("Prediction").over(w))
  
  # calculate rolling forecast sums
  pred_sum_udf = udf(lambda preds, period: float(np.sum(preds[:period])), FloatType())
  rolling_pred_sums = rolling_prediction_lists \
      .withColumn("RollingSum", pred_sum_udf("next_preds", "N"))
  </code></pre>
  
  <p>This solution is also works with the test data. I couldn't have chance to test it with the original data yet but whether it works or not I do not like this solution. Is there any smarter way to solve this?</p>