Calculate residual amount in dataframe column - scala

Question

I have a "capacity" dataframe:
scala> sql("create table capacity (id String, capacity Int)");
scala> sql("insert into capacity values ('A', 50), ('B', 100)");
scala> sql("select * from capacity").show(false)
+---+--------+
|id |capacity|
+---+--------+
|A |50 |
|B |100 |
+---+--------+
I have another "used" dataframe with following information:
scala> sql ("create table used (id String, capacityId String, used Int)");
scala> sql ("insert into used values ('item1', 'A', 10), ('item2', 'A', 20), ('item3', 'A', 10), ('item4', 'B', 30), ('item5', 'B', 40), ('item6', 'B', 40)")
scala> sql("select * from used order by capacityId").show(false)
+-----+----------+----+
|id |capacityId|used|
+-----+----------+----+
|item1|A |10 |
|item3|A |10 |
|item2|A |20 |
|item6|B |40 |
|item4|B |30 |
|item5|B |40 |
+-----+----------+----+
Column "capacityId" of the "used" dataframe is foreign key to column "id" of the "capacity" dataframe.
I want to calculate the "capacityLeft" column which is residual amount at that point of time.
+-----+----------+----+--------------+
|id |capacityId|used| capacityLeft |
+-----+----------+----+--------------+
|item1|A |10 |40 | <- 50(capacity of 'A')-10
|item3|A |10 |30 | <- 40-10
|item2|A |20 |10 | <- 30-20
|item6|B |40 |60 | <- 100(capacity of 'B')-40
|item4|B |30 |30 | <- 60-30
|item5|B |40 |-10 | <- 30-40
+-----+----------+----+--------------+
In real senario, the "createdDate" column is used for ordering of "used" dataframe column.
Spark version: 2.2

Answer 1

This can be solved by using window functions in Spark. Note that for this to work there need to exist a column that keep track of the row order for each capacityId.
Start by joining the two dataframes together:
val df = used.join(capacity.withColumnRenamed("id", "capacityId"), Seq("capacityId"), "inner")
Here the id in the capacity dataframe is renamed to match the id name in the used dataframe as to not keep a duplicate columns.
Now create a window and calculate the cumsum of the used column. Take the value of the capacity and subtract the cumsum to get the remaining amount:
val w = Window.partitionBy("capacityId").orderBy("createdDate")
val df2 = df.withColumn("capacityLeft", $"capacity" - sum($"used").over(w))
Resulting dataframe with example createdDate column:
+----------+-----+----+-----------+--------+------------+
|capacityId| id|used|createdDate|capacity|capacityLeft|
+----------+-----+----+-----------+--------+------------+
| B|item6| 40| 1| 100| 60|
| B|item4| 30| 2| 100| 30|
| B|item5| 40| 3| 100| -10|
| A|item1| 10| 1| 50| 40|
| A|item3| 10| 2| 50| 30|
| A|item2| 20| 3| 50| 10|
+----------+-----+----+-----------+--------+------------+
Any unwanted columns can now be removed with drop.