Рассчитать остаточное количество в столбце данных

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" кадра данных "Capacity".Я хочу вычислить столбец «acityLeft », который является остаточным количеством в тот момент.

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

В реальном сенарио столбец «madeDate» используется для упорядочивания «используемого» столбца данных.

Версия Spark: 2.2

Answer 1

Это можно решить с помощью оконных функций в Spark.Обратите внимание, что для этого должен существовать столбец, который отслеживает порядок строк для каждого capacityId.

. Начните с объединения двух фреймов данных:

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

Здесьid в capacity dataframe переименовывается, чтобы соответствовать имени id в used dataframe, чтобы не сохранять повторяющиеся столбцы.

Теперь создайте окно и вычислите сумму используемого столбца.Возьмите значение 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.