Introduction

Spark pivot() function is used to pivot/rotate the data from one DataFrame/Dataset column into multiple columns (transform row to column) and unpivot is used to transform 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.

val data = Seq(
  ("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")
)

import spark.sqlContext.implicits._
val df = data.toDF("Product", "Amount", "Country")
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 Spark 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.

val pivotDF = df.groupBy("Product").pivot("Country").sum("Amount")
pivotDF.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

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

val countries = Seq("USA", "China", "Canada", "Mexico")
val pivotDF = df.groupBy("Product").pivot("Country", countries).sum("Amount")
pivotDF.show()

Another approach is to do two-phase aggregation. Spark 2.0 uses this implementation in order to improve the performance Spark-13749

val pivotDF = df.groupBy("Product", "Country")
  .sum("Amount")
  .groupBy("Product")
  .pivot("Country")
  .sum("sum(Amount)")
pivotDF.show()

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

3. Unpivot DataFrame

Unpivot is a reverse operation, we can achieve by rotating column values into row values. Spark SQL doesn’t have an unpivot function, hence will use the stack() function. The code below converts column countries to rows.

val unPivotDF = pivotDF.select($"Product",
  expr("stack(3, 'Canada', Canada, 'China', China, 'Mexico', Mexico) as (Country, Total)"))
  .where("Total is not null")
unPivotDF.show()

It converts the pivoted column “country” to rows.

output
+-------+-------+-----+
|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 do Spark DataFrame transpose or pivot without aggregation?

Of course, you can, but unfortunately, you can’t achieve this using the Pivot function. However, pivoting or transposing the DataFrame structure without aggregation from rows to columns and columns to rows can be easily done using a Spark and Scala hack. Please refer to this example.

The complete code can be downloaded from GitHub