Как сделать условную агрегацию после группового в фрейме данных pyspark?

Question

Я пытаюсь сгруппировать по столбцу идентификатора в pyspark кадре данных и суммировать столбец в зависимости от значения другого столбца.

Для иллюстрации рассмотрим следующий фиктивный кадр данных:

+-----+-------+---------+
|   ID|   type|   amount|
+-----+-------+---------+
|    1|      a|       55|
|    2|      b|     1455|
|    2|      a|       20|
|    2|      b|      100|
|    3|   null|      230|
+-----+-------+---------+

Мой желаемый результат:

+-----+--------+----------+----------+
|   ID|   sales|   sales_a|   sales_b|
+-----+--------+----------+----------+
|    1|      55|        55|         0| 
|    2|    1575|        20|      1555|        
|    3|     230|         0|         0|
+-----+--------+----------+----------+

Таким образом, sales будет суммой amount, а sales_a и sales_b - сумма amount, когда type равно a или b соответственно.

Для sales я знаю, что это можно сделать так:

from pyspark.sql import functions as F
df = df.groupBy("ID").agg(F.sum("amount").alias("sales"))

Для других, я думаю, F.when было бы полезно, но я не уверен, как это сделать.

Answer 1

Вы можете создать два столбца до агрегирования на основе значения типа.

df.withColumn("sales_a", F.when(col("type") == "a", col("amount"))) \
  .withColumn("sales_b", F.when(col("type") == "b", col("amount"))) \
  .groupBy("ID") \
  .agg(F.sum("amount").alias("sales"),
       F.sum("sales_a").alias("sales_a"),
       F.sum("sales_b").alias("sales_b"))

Answer 2

from pyspark.sql import functions as F
df = df.groupBy("ID").agg(F.sum("amount").alias("sales"))
dfPivot = df.filter("type is not null").groupBy("ID").pivot("type").agg(F.sum("amount").alias("sales"))

res = df.join(dfPivot, df.id== dfPivot.id,how='left')

Затем замените null на 0.Это универсальное решение будет работать независимо от значений в столбце type, поэтому, если type c добавлено в фрейм данных, будет создан столбец _c

.