时间戳

首页 > 解决方案 > 如何在pyspark数据框中的groupby之后进行条件聚合？

apache-spark - 如何在pyspark数据框中的groupby之后进行条件聚合？

问题描述

我正在尝试按数据框中的 ID 列进行分组，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，whilesales_a和sales_b的总和，分别是amountwhentype是a或的总和b。

对于sales，我知道这可以这样做：

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

对于其他人，我猜F.when会很有用，但我不知道该怎么做。

标签： apache-sparkpyspark

解决方案

您可以根据类型的值在聚合之前创建两列。

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"))

推荐阅读