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

“已使用”数据帧的“ capacityId”列是“容量”数据帧的“ 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”列用于对“二手”数据框列进行排序。

火花版本：2.2

Answer 1

这可以通过在Spark中使用window functions来解决。请注意，要执行此操作，需要有一个列来跟踪每个capacityId的行顺序。

首先将两个数据帧结合在一起：

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

此处将capacity数据框中的id重命名为与used数据框中的id名称相匹配，以免保留重复的列。

现在创建一个窗口并计算所用列的总和。取capacity的值并减去总和以得到剩余的金额：

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删除任何不需要的列。

计算数据框列中的剩余量

1 个答案: