Pivot a spark dataframe and calculate sum of 2 rows above

user_4186589 0
Viewed 554
+--------------------+-------------------+------------+
|            name    |               1990|1991        |
+--------------------+-------------------+------------+
|abc                 |100                |        300 |
|bbc                 |200                |        400 |
|cbc                 |300                |        500 |
|xyz                 |700                |        500 |
|xzz                 |200                |        500 |

Above is a dataframe in pyspark. I would like to convert this to a pivoted table and calculate a rolling summation. Is there a way to convert the columns to row on the name? Belo sis the expected result which is basically summation of previous 2 years. How can we accomplish this?

+--------------------+-------------------+------------+
|            year    |               abc |bbc         |
+--------------------+-------------------+------------+
|1990                 |100               |        300 |
|1991                 |100+200           |    400+300 |
pyspark
1 Answers

You can split the task in two parts

  • first create the pivot table and then
  • use a window to calculate the sum
from pyspark.sql import types as T
from pyspark.sql import Window

df = ...

w=Window.orderBy("year").rowsBetween(-1,0)
df.withColumnRenamed("1990", "c1990").withColumnRenamed("1991", "c1991")\
    .selectExpr("name", "stack(2, '1990', c1990, '1991', c1991)") \
    .groupBy("col0") \
    .pivot("name", ["abc", "bbc"]) \
    .agg(F.first("col1")) \
    .withColumnRenamed("col0", "year") \
    .withColumn("abc_sum", F.sum("abc").over(w)) \
    .withColumn("bbc_sum", F.sum("bbc").over(w)) \
    .show()

Output:

+----+---+---+-------+-------+
|year|abc|bbc|abc_sum|bbc_sum|
+----+---+---+-------+-------+
|1990|100|200|    100|    200|
|1991|300|400|    400|    600|
+----+---+---+-------+-------+

Edit: using all values from the name column:

df2=df.withColumnRenamed("1990", "c1990").withColumnRenamed("1991", "c1991")\
    .selectExpr("name", "stack(2, '1990', c1990, '1991', c1991)") \
    .groupBy("col0") \
    .pivot("name") \
    .agg(F.first("col1")) \
    .withColumnRenamed("col0", "year")

cols =[F.col("year")] + [F.sum(col).over(w).alias(col) for col in df2.columns[1:]]
df2.select(cols).show()
user_21298010
Related