每组按行求和并将总数添加为 Pyspark 数据框中的新行 [英] Rowwise sum per group and add total as a new row in dataframe in Pyspark

问题描述

我有一个像这个示例的数据框

I have a dataframe like this sample

df = spark.createDataFrame(
    [(2, "A" , "A2" , 2500),
    (2, "A" , "A11" , 3500),
    (2, "A" , "A12" , 5500),
    (4, "B" , "B25" , 7600),
    (4, "B", "B26" ,5600),
    (5, "C" , "c25" ,2658),
    (5, "C" , "c27" , 1100),
    (5, "C" , "c28" , 1200)],
    ['parent', 'group' , "brand" , "usage"])


output :
+------+-----+-----+-----+
|parent|group|brand|usage|
+------+-----+-----+-----+
|     2|    A|   A2| 2500|
|     2|    A|  A11| 3500|
|     4|    B|  B25| 7600|
|     4|    B|  B26| 5600|
|     5|    C|  c25| 2658|
|     5|    C|  c27| 1100|
|     5|    C|  c28| 1200|
+------+-----+-----+-----+

我想要做的是计算每个组的使用总量，并将其添加为带有品牌总价值的新行.如何在 PySpark 中执行此操作?:

What I would like to do is to compute, for each group total of usage and add it as a new row with Total value for brand. How can I do this in PySpark?:

Expected result:

+------+-----+-----+-----+
|parent|group|brand|usage|
+------+-----+-----+-----+
|     2|    A|   A2| 2500|
|     2|    A|  A11| 3500|
|     2|    A|Total| 6000|
|     4|    B|  B25| 7600|
|     4|    B|  B26| 5600|
|     4|    B|Total|18700|
|     5|    C|  c25| 2658|
|     5|    C|  c27| 1100|
|     5|    C|  c28| 1200|
|     5|    C|Total| 4958|
+------+-----+-----+-----+

推荐答案

import pyspark.sql.functions as F

df = spark.createDataFrame(
[(2, "A" , "A2" , 2500),
(2, "A" , "A11" , 3500),
(2, "A" , "A12" , 5500),
(4, "B" , "B25" , 7600),
(4, "B", "B26" ,5600),
(5, "C" , "c25" ,2658),
(5, "C" , "c27" , 1100),
(5, "C" , "c28" , 1200)],
['parent', 'group' , "brand" , "usage"])

df.show()
+------+-----+-----+-----+
|parent|group|brand|usage|
+------+-----+-----+-----+
|     2|    A|   A2| 2500|
|     2|    A|  A11| 3500|
|     2|    A|  A12| 5500|
|     4|    B|  B25| 7600|
|     4|    B|  B26| 5600|
|     5|    C|  c25| 2658|
|     5|    C|  c27| 1100|
|     5|    C|  c28| 1200|
+------+-----+-----+-----+

#Group by and sum to get the totals
totals = df.groupBy(['group','parent']).agg(F.sum('usage').alias('usage')).withColumn('brand', F.lit('Total'))

# create a temp variable to sort
totals = totals.withColumn('sort_id', F.lit(2))
df = df.withColumn('sort_id', F.lit(1))

#Union dataframes, drop temp variable and show
df.unionByName(totals).sort(['group','sort_id']).drop('sort_id').show()

+------+-----+-----+-----+
|parent|group|brand|usage|
+------+-----+-----+-----+
|     2|    A|  A12| 5500|
|     2|    A|  A11| 3500|
|     2|    A|   A2| 2500|
|     2|    A|Total|11500|
|     4|    B|  B25| 7600|
|     4|    B|  B26| 5600|
|     4|    B|Total|13200|
|     5|    C|  c25| 2658|
|     5|    C|  c28| 1200|
|     5|    C|  c27| 1100|
|     5|    C|Total| 4958|
+------+-----+-----+-----+

这篇关于每组按行求和并将总数添加为 Pyspark 数据框中的新行的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！