计算数据框列中的剩余量

Question

我有一个“容量”数据框：

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     |
+---+--------+

我有另一个“使用过的”数据框，其中包含以下信息：

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  |
+-----+----------+----+

“已使用”数据帧的“容量 ID”列是“容量”数据帧的“id”列的外键。 我想计算“capacityLeft”列，它是那个时间点的剩余量。

+-----+----------+----+--------------+
|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
+-----+----------+----+--------------+

在实际场景中，“createdDate”列用于对“used”数据框列进行排序。

Spark 版本：2.2

Answer 1

这可以通过在 Spark 中使用 window functions 来解决。请注意，要使其正常工作，需要存在一个列来跟踪每个 capacityId 的行顺序。

首先将两个数据框连接在一起：

val df = used.join(capacity.withColumnRenamed("id", "capacityId"), Seq("capacityId"), "inner")

这里 capacity 数据框中的 id 被重命名以匹配 used 数据框中的 id 名称，以免保留重复的列。

现在创建一个窗口并计算所用列的累积和。取capacity的值，减去cumsum得到余数：

val w = Window.partitionBy("capacityId").orderBy("createdDate")
val df2 = df.withColumn("capacityLeft", $"capacity" - sum($"used").over(w))

带有示例createdDate 列的结果数据框：

+----------+-----+----+-----------+--------+------------+
|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|
+----------+-----+----+-----------+--------+------------+

现在可以使用drop 删除任何不需要的列。