Copy missed data from top/bottom row col values - scala

Question

I have a dataframe, with index, category and few other columns. index and category never be empty/null. but other columns data comes null, When all other columns data is null then we have to copy from top/bottom row values based on cateogry.
val df = Seq(
(1,1, null, null, null ),
(2,1, null, null, null ),
(3,1, null, null, null ),
(4,1,"123.12", "124.52", "95.98" ),
(5,1, "452.12", "478.65", "1865.12" ),
(1,2,"2014.21", "147", "265"),
(2,2, "1457", "12483.00", "215.21"),
(3,2, null, null, null),
(4,2, null, null, null) ).toDF("index", "category", "col1", "col2", "col3")
scala> df.show
+-----+--------+-------+--------+-------+
|index|category| col1| col2| col3|
+-----+--------+-------+--------+-------+
| 1| 1| null| null| null|
| 2| 1| null| null| null|
| 3| 1| null| null| null|
| 4| 1| 123.12| 124.52| 95.98|
| 5| 1| 452.12| 478.65|1865.12|
| 1| 2|2014.21| 147| 265|
| 2| 2| 1457|12483.00| 215.21|
| 3| 2| null| null| null|
| 4| 2| null| null| null|
+-----+--------+-------+--------+-------+
Expecting dataframe as below
+-----+--------+-------+--------+-------+
|index|category| col1| col2| col3|
+-----+--------+-------+--------+-------+
| 1| 1| 123.12| 124.52| 95.98| // Copied from below for same category
| 2| 1| 123.12| 124.52| 95.98| // Copied from below for same category
| 3| 1| 123.12| 124.52| 95.98|
| 4| 1| 123.12| 124.52| 95.98|
| 5| 1| 452.12| 478.65|1865.12|
| 1| 2|2014.21| 147| 265|
| 2| 2| 1457|12483.00| 215.21|
| 3| 2| 1457|12483.00| 215.21| // Copied from above for same category
| 4| 2| 1457|12483.00| 215.21| // Copied from above for same category
+-----+--------+-------+--------+-------+

Answer 1

Update When several rows with nulls possible, advanced Windows have to be used:
val cols = Seq("col1", "col2", "col3")
val beforeWindow = Window
.partitionBy("category")
.orderBy("index")
.rangeBetween(Window.unboundedPreceding, Window.currentRow)
val afterWindow = Window
.partitionBy("category")
.orderBy("index")
.rangeBetween(Window.currentRow, Window.unboundedFollowing)
val result = cols.foldLeft(df)((updated, columnName) =>
updated.withColumn(columnName,
coalesce(col(columnName),
last(columnName, ignoreNulls = true).over(beforeWindow),
first(columnName, ignoreNulls = true).over(afterWindow)
))
)
In one null case can be resolved with Window functions "lead" and "lag", and "coalesce":
val cols = Seq("col1", "col2", "col3")
val categoryWindow = Window.partitionBy("category").orderBy("index")
val result = cols.foldLeft(df)((updated, columnName) =>
updated.withColumn(columnName,
coalesce(col(columnName),
lag(col(columnName), 1).over(categoryWindow),
lead(col(columnName), 1).over(categoryWindow)
))
)
result.show(false)
Output:
+-----+--------+-------+--------+-------+
|index|category|col1 |col2 |col3 |
+-----+--------+-------+--------+-------+
|1 |1 |123.12 |124.52 |95.98 |
|2 |1 |123.12 |124.52 |95.98 |
|3 |1 |452.12 |478.65 |1865.12|
|1 |2 |2014.21|147 |265 |
|2 |2 |1457 |12483.00|215.21 |
|3 |2 |1.25 |3.45 |26.3 |
|4 |2 |1.25 |3.45 |26.3 |
+-----+--------+-------+--------+-------+