Introduction

In PySpark, the pivot() function is used to pivot/rotate the data from one DataFrame column into multiple columns (transform row to column) and the inverse operation, unpivot, allows transforming it back (transform columns to rows).

Pivot() is an aggregation where one of the grouping columns' values is transposed into individual columns with distinct data.

Let’s create a DataFrame to work with.

from pyspark.sql import SparkSession
from pyspark.sql.functions import expr

spark = SparkSession.builder.master("local").appName("PivotExample").getOrCreate()

data = [
    ("Banana", 1000, "USA"),
    ("Carrots", 1500, "USA"),
    ("Beans", 1600, "USA"),
    ("Orange", 2000, "USA"),
    ("Orange", 2000, "USA"),
    ("Banana", 400, "China"),
    ("Carrots", 1200, "China"),
    ("Beans", 1500, "China"),
    ("Orange", 4000, "China"),
    ("Banana", 2000, "Canada"),
    ("Carrots", 2000, "Canada"),
    ("Beans", 2000, "Mexico")
]

columns = ["Product", "Amount", "Country"]
df = spark.createDataFrame(data, columns)
df.show()

DataFrame ‘df’ consists of 3 columns: Product, Amount, and Country, as shown below.

+-------+------+-------+
|Product|Amount|Country|
+-------+------+-------+
| Banana|  1000|    USA|
|Carrots|  1500|    USA|
|  Beans|  1600|    USA|
| Orange|  2000|    USA|
| Orange|  2000|    USA|
| Banana|   400|  China|
|Carrots|  1200|  China|
|  Beans|  1500|  China|
| Orange|  4000|  China|
| Banana|  2000| Canada|
|Carrots|  2000| Canada|
|  Beans|  2000| Mexico|
+-------+------+-------+

1. Pivot PySpark DataFrame

pivot() is an aggregation where one of the grouping columns' values is transposed into individual columns with distinct data. From the above DataFrame, to get the total amount exported to each country of each product, we can group by Product, pivot by Country, and sum the Amount.

pivot_df = df.groupBy("Product").pivot("Country").sum("Amount")
pivot_df.show()

This will transpose the countries from DataFrame rows into columns and produce the output below. Wherever data is not present, it is represented as null by default.

+-------+------+-----+------+----+
|Product|Canada|China|Mexico| USA|
+-------+------+-----+------+----+
| Orange|  null| 4000|  null|4000|
|  Beans|  null| 1500|  2000|1600|
| Banana|  2000|  400|  null|1000|
|Carrots|  2000| 1200|  null|1500|
+-------+------+-----+------+----+

2. Pivot Performance Improvement in Spark 2.0

From Spark 2.0 onwards, performance has been improved for Pivot; however, if you are using a lower version, note that Pivot is a very expensive operation. It is recommended to provide column data (if known) as an argument to the function as shown below.

countries = ["USA", "China", "Canada", "Mexico"]
pivot_df = df.groupBy("Product").pivot("Country", countries).sum("Amount")
pivot_df.show()

Another approach is to use two-phase aggregation. Spark 2.0 uses this implementation to improve performance Spark-13749.

pivot_df = df.groupBy("Product", "Country") \
    .sum("Amount") \
    .groupBy("Product") \
    .pivot("Country") \
    .sum("sum(Amount)")
pivot_df.show()

The above two examples return the same output but with better performance.

3. Unpivot DataFrame

Unpivot is a reverse operation, achieved by rotating column values into row values. PySpark doesn’t have an unpivot function, hence the stack() function is used. The code below converts column countries to rows.

unpivot_df = pivot_df.select("Product", 
    expr("stack(3, 'Canada', Canada, 'China', China, 'Mexico', Mexico) as (Country, Total)")
).where("Total is not null")
unpivot_df.show()

It converts the pivoted column “country” to rows.

+-------+-------+-----+
|Product|Country|Total|
+-------+-------+-----+
| Orange|  China| 4000|
|  Beans|  China| 1500|
|  Beans| Mexico| 2000|
| Banana| Canada| 2000|
| Banana|  China|  400|
|Carrots| Canada| 2000|
|Carrots|  China| 1200|
+-------+-------+-----+

4. Transpose or Pivot without Aggregation

Can we transpose or pivot a PySpark DataFrame without aggregation?

Indeed, it's possible, but unfortunately, you can’t achieve this using the Pivot function directly. However, pivoting or transposing the DataFrame structure without aggregation from rows to columns and columns to rows can be done using workaround methods in PySpark. Please refer to this discussion for examples.

The complete code can be downloaded from this GitHub repository.