Question

我正在尝试按pyspark数据帧中的ID列进行分组，并根据另一列的值对一列求和。

为说明起见，请考虑以下虚拟数据帧：

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

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

Answer 2

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

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

在pyspark数据帧中进行分组后如何进行条件聚合？

2 个答案: