Introduction

In Apache Spark, the pivot() function is used to rotate data from one DataFrame/Dataset column into multiple columns (transforming rows into columns), while unpivot transforms it back (transforming columns into rows).

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

Let’s create a DataFrame to work with.

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.catalyst.encoders.RowEncoder;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

import java.util.Arrays;
import java.util.List;

SparkSession spark = SparkSession.builder()
        .appName("Pivot Example")
        .master("local")
        .getOrCreate();

List<Row> data = Arrays.asList(
        RowFactory.create("Banana", 1000, "USA"),
        RowFactory.create("Carrots", 1500, "USA"),
        RowFactory.create("Beans", 1600, "USA"),
        RowFactory.create("Orange", 2000, "USA"),
        RowFactory.create("Orange", 2000, "USA"),
        RowFactory.create("Banana", 400, "China"),
        RowFactory.create("Carrots", 1200, "China"),
        RowFactory.create("Beans", 1500, "China"),
        RowFactory.create("Orange", 4000, "China"),
        RowFactory.create("Banana", 2000, "Canada"),
        RowFactory.create("Carrots", 2000, "Canada"),
        RowFactory.create("Beans", 2000, "Mexico")
);

StructType schema = new StructType(new StructField[]{
        DataTypes.createStructField("Product", DataTypes.StringType, false),
        DataTypes.createStructField("Amount", DataTypes.IntegerType, false),
        DataTypes.createStructField("Country", DataTypes.StringType, false)
});

Dataset<Row> df = spark.createDataFrame(data, schema);
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 column 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.

import org.apache.spark.sql.functions;

Dataset<Row> pivotDF = df.groupBy("Product")
        .pivot("Country")
        .sum("Amount");
pivotDF.show();

This will transpose the countries from DataFrame rows into columns, producing 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

Starting with Spark 2.0, performance has been improved for the Pivot, but if you are using a version below 2.0, note that Pivot is a very resource-intensive operation. Therefore, it is recommended to provide column data (if known) as an argument to the function as shown below.

List<String> countries = Arrays.asList("USA", "China", "Canada", "Mexico");
Dataset<Row> pivotDF = df.groupBy("Product")
        .pivot("Country", countries)
        .sum("Amount");
pivotDF.show();

Another approach is to do a two-phase aggregation. Spark 2.0 uses this implementation to enhance performance Spark-13749.

Dataset<Row> phaseAggDF = df.groupBy("Product", "Country")
        .agg(functions.sum("Amount").alias("TotalAmount"))
        .groupBy("Product")
        .pivot("Country")
        .sum("TotalAmount");
phaseAggDF.show();

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

3. Unpivot DataFrame

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

Dataset<Row> unPivotDF = pivotDF.selectExpr("Product", 
    "stack(3, 'Canada', Canada, 'China', China, 'Mexico', Mexico) as (Country, Total)")
    .filter("Total is not null");
unPivotDF.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 do Spark DataFrame transpose or pivot without aggregation?

Yes, it is possible; however, you can’t achieve this using the Pivot function. Pivoting or transposing the DataFrame structure without aggregation from rows to columns and columns to rows can be done using a method similar to the one described in this StackOverflow example.

The complete code can be downloaded from GitHub