Filling missing values per group in pyspark?

Question

I have a DataFrame which looks like this:

| date      | week | quantity |
|-----------|------|----------|
| 1/1/2020  | 1    | 17       |
| 1/2/2020  | 1    | 15       |
| 1/3/2020  | 1    | 9        |
| 1/4/2020  | 1    |          |
| 1/5/2020  | 1    |          |
| 1/6/2020  | 1    | 3        |
| 1/7/2020  | 1    | 3        |
| 1/8/2020  | 2    | 10       |
| 1/9/2020  | 2    | 8        |
| 1/10/2020 | 2    | 8        |
| 1/11/2020 | 2    |          |
| 1/12/2020 | 2    |          |
| 1/13/2020 | 2    |          |
| 1/14/2020 | 2    | 8        |

I need to impute missing values by distributing the last known value between missing dates. So, for week 1 my output should look like:

| date     | week | quantity | quant_dist |
|----------|------|----------|------------|
| 1/1/2020 | 1    | 17       | 17         |
| 1/2/2020 | 1    | 15       | 15         |
| 1/3/2020 | 1    | 9        | 9          |
| 1/4/2020 | 1    |          | 1          |
| 1/5/2020 | 1    |          | 1          |
| 1/6/2020 | 1    | 3        | 1          |
| 1/7/2020 | 1    | 3        | 3          |

and for week 2 it should look like:

| date      | week | quantity | quant_dist |
|-----------|------|----------|------------|
| 1/8/2020  | 2    | 10       | 10         |
| 1/9/2020  | 2    | 8        | 8          |
| 1/10/2020 | 2    | 8        | 8          |
| 1/11/2020 | 2    |          | 2          |
| 1/12/2020 | 2    |          | 2          |
| 1/13/2020 | 2    |          | 2          |
| 1/14/2020 | 2    | 8        | 2          |

I have been stuck on this for over 2 days now. Any suggestions?

Does this answer your question? [Forward fill missing values in Spark/Python](https://stackoverflow.com/questions/38131982/forward-fill-missing-values-in-spark-python) — Equinox, Aug 30 '20 at 14:53

Lamanus · Accepted Answer · 2020-08-30T15:01:00.383

With many functions and Windows, I can achieve it.

from pyspark.sql.functions import *
from pyspark.sql import Window

w1 = Window.orderBy('date').rowsBetween(Window.currentRow, Window.unboundedFollowing)
w2 = Window.partitionBy('group')

df.withColumn('date', to_date('date', 'MM/dd/yyyy')) \
  .withColumn('group', sum(when(col('quantity').isNull(), 0).otherwise(1)).over(w1)) \
  .withColumn('count', count(when(col('quantity').isNull(), 1).otherwise(0)).over(w2)) \
  .withColumn('quantity', first('quantity', ignorenulls=True).over(w1) / col('count')).show()

+----------+----+--------+-----+-----+
|      date|week|quantity|group|count|
+----------+----+--------+-----+-----+
|2020-01-01|   1|    17.0|    9|    1|
|2020-01-02|   1|    15.0|    8|    1|
|2020-01-03|   1|     9.0|    7|    1|
|2020-01-04|   1|     1.0|    6|    3|
|2020-01-05|   1|     1.0|    6|    3|
|2020-01-06|   1|     1.0|    6|    3|
|2020-01-07|   1|     3.0|    5|    1|
|2020-01-08|   2|    10.0|    4|    1|
|2020-01-09|   2|     8.0|    3|    1|
|2020-01-10|   2|     8.0|    2|    1|
|2020-01-11|   2|     2.0|    1|    4|
|2020-01-12|   2|     2.0|    1|    4|
|2020-01-13|   2|     2.0|    1|    4|
|2020-01-14|   2|     2.0|    1|    4|
+----------+----+--------+-----+-----+

Thank you so much.. I don't have much reputation to upvote this answer — sisupalan, Aug 30 '20 at 15:36

Filling missing values per group in pyspark?

1 Answers1