pyspark column sum with transpose

2
votes

I have a data frame looks like -

+---+---+---+---+
| id| w1| w2| w3|
+---+---+---+---+
|  1|100|150|200|
|  2|200|400|500|
|  3|500|600|150|
+---+---+---+---+

I want output looks like - 

full   total_amt
 w1       800
 w2       1150
 w3       850

My code is - 

df = spark.createDataFrame(
    [(1, 100,150,200), (2, 200,400,500), (3, 500,600,150)], ("id", "w1","w2","w3"))

res = df.unionAll(
    df.select([
        F.lit('All').alias('id'), 
        F.sum(df.w1).alias('w1'),
        F.sum(df.w2).alias('w2'),
        F.sum(df.w3).alias('w3') 
    ]))
res.show()

But output gives me - 

+---+---+----+---+
| id| w1|  w2| w3|
+---+---+----+---+
|  1|100| 150|200|
|  2|200| 400|500|
|  3|500| 600|150|
|All|800|1150|850|
+---+---+----+---+

I think after addition need to create pivot. All the fields are numeric in nature.

2
pysparkpyspark-sqlpyspark-dataframes

2 Answers

0
votes

Try this approach -

First aggregate the data and then use stack function to convet columns to rows

import pyspark.sql.functions as psf

#perform aggregation
df_agg = df.agg(psf.sum('w1').alias('w1'), psf.sum('w2').alias('w2'), psf.sum('w3').alias('w3'))

#let's have a look at aggregated dataframe
df_agg.show()
#+---+----+---+
#| w1|  w2| w3|
#+---+----+---+
#|800|1150|850|
#+---+----+---+

#Use stack function to convert column to rows
df_agg.selectExpr("stack(3, 'w1', w1, 'w2', w2, 'w3', w3) as (full, total)").show()
#+----+-----+
#|full|total|
#+----+-----+
#|  w1|  800|
#|  w2| 1150|
#|  w3|  850|
#+----+-----+
0
votes

A quick solution could be

>>> df.createOrReplaceTempView('df')

>>> spark.sql('''
...    select 'w1' as full, sum(w1) as total  from df 
...    union
...    select 'w2' as full, sum(w2) as total  from df 
...    union
...    select 'w3' as full, sum(w3) as total  from df 
... ''').show()
+----+-----+                                                                    
|full|total|
+----+-----+
|  w2| 1150|
|  w3|  850|
|  w1|  800|
+----+-----+