Agfdhyk

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

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.