How to pivot Spark DataFrame?

I overcame this by writing a for loop to dynamically create a SQL query. Say I have:

id tag value 1 US 50 1 UK 100 1 Can 125 2 US 75 2 UK 150 2 Can 175 

and I want:

id US UK Can 1 50 100 125 2 75 150 175 

I can create a list with the value I want to pivot and then create a string containing the SQL query I need.

val countries = List("US", "UK", "Can") val numCountries = countries.length - 1 var query = "select *, " for (i <- 0 to numCountries-1) { query += """case when tag = """" + countries(i) + """" then value else 0 end as """ + countries(i) + ", " } query += """case when tag = """" + countries.last + """" then value else 0 end as """ + countries.last + " from myTable" myDataFrame.registerTempTable("myTable") val myDF1 = sqlContext.sql(query) 

I can create similar query to then do the aggregation. Not a very elegant solution but it works and is flexible for any list of values, which can also be passed in as an argument when your code is called.

A pivot operator has been added to the Spark dataframe API, and is part of Spark 1.6.

See https://github.com/apache/spark/pull/7841 for details.

I have solved a similar problem using dataframes with the following steps:

Create columns for all your countries, with 'value' as the value:

import org.apache.spark.sql.functions._ val countries = List("US", "UK", "Can") val countryValue = udf{(countryToCheck: String, countryInRow: String, value: Long) => if(countryToCheck == countryInRow) value else 0 } val countryFuncs = countries.map{country => (dataFrame: DataFrame) => dataFrame.withColumn(country, countryValue(lit(country), df("tag"), df("value"))) } val dfWithCountries = Function.chain(countryFuncs)(df).drop("tag").drop("value") 

Your dataframe 'dfWithCountries' will look like this:

+--+--+---+---+ |id|US| UK|Can| +--+--+---+---+ | 1|50| 0| 0| | 1| 0|100| 0| | 1| 0| 0|125| | 2|75| 0| 0| | 2| 0|150| 0| | 2| 0| 0|175| +--+--+---+---+ 

Now you can sum together all the values for your desired result:

dfWithCountries.groupBy("id").sum(countries: _*).show 

Result:

+--+-------+-------+--------+ |id|SUM(US)|SUM(UK)|SUM(Can)| +--+-------+-------+--------+ | 1| 50| 100| 125| | 2| 75| 150| 175| +--+-------+-------+--------+ 

It's not a very elegant solution though. I had to create a chain of functions to add in all the columns. Also if I have lots of countries, I will expand my temporary data set to a very wide set with lots of zeroes.

There is a SIMPLE method for pivoting :

 id tag value 1 US 50 1 UK 100 1 Can 125 2 US 75 2 UK 150 2 Can 175 import sparkSession.implicits._ val data = Seq( (1,"US",50), (1,"UK",100), (1,"Can",125), (2,"US",75), (2,"UK",150), (2,"Can",175), ) val dataFrame = data.toDF("id","tag","value") val df2 = dataFrame .groupBy("id") .pivot("tag") .max("value") df2.show() +---+---+---+---+ | id|Can| UK| US| +---+---+---+---+ | 1|125|100| 50| | 2|175|150| 75| +---+---+---+---+ 

There is simple and elegant solution.

scala> spark.sql("select * from k_tags limit 10").show() +---------------+-------------+------+ | imsi| name| value| +---------------+-------------+------+ |246021000000000| age| 37| |246021000000000| gender|Female| |246021000000000| arpu| 22| |246021000000000| DeviceType| Phone| |246021000000000|DataAllowance| 6GB| +---------------+-------------+------+ scala> spark.sql("select * from k_tags limit 10").groupBy($"imsi").pivot("name").agg(min($"value")).show() +---------------+-------------+----------+---+----+------+ | imsi|DataAllowance|DeviceType|age|arpu|gender| +---------------+-------------+----------+---+----+------+ |246021000000000| 6GB| Phone| 37| 22|Female| |246021000000001| 1GB| Phone| 72| 10| Male| +---------------+-------------+----------+---+----+------+ 

There are plenty of examples of pivot operation on dataset/dataframe, but I could not find many using SQL. Here is an example that worked for me.

create or replace temporary view faang as SELECT stock.date AS `Date`, stock.adj_close AS `Price`, stock.symbol as `Symbol` FROM stock WHERE (stock.symbol rlike '^(FB|AAPL|GOOG|AMZN)$') and year(date) > 2010; SELECT * from faang PIVOT (max(price) for symbol in ('AAPL', 'FB', 'GOOG', 'AMZN')) order by date; 

Initially i adopted Al M's solution. Later took the same thought and rewrote this function as a transpose function.

This method transposes any df rows to columns of any data-format with using key and value column

for input csv

id,tag,value 1,US,50a 1,UK,100 1,Can,125 2,US,75 2,UK,150 2,Can,175 

ouput

+--+---+---+---+ |id| UK| US|Can| +--+---+---+---+ | 2|150| 75|175| | 1|100|50a|125| +--+---+---+---+ 

transpose method :

def transpose(hc : HiveContext , df: DataFrame,compositeId: List[String], key: String, value: String) = { val distinctCols = df.select(key).distinct.map { r => r(0) }.collect().toList val rdd = df.map { row => (compositeId.collect { case id => row.getAs(id).asInstanceOf[Any] }, scala.collection.mutable.Map(row.getAs(key).asInstanceOf[Any] -> row.getAs(value).asInstanceOf[Any])) } val pairRdd = rdd.reduceByKey(_ ++ _) val rowRdd = pairRdd.map(r => dynamicRow(r, distinctCols)) hc.createDataFrame(rowRdd, getSchema(df.schema, compositeId, (key, distinctCols))) } private def dynamicRow(r: (List[Any], scala.collection.mutable.Map[Any, Any]), colNames: List[Any]) = { val cols = colNames.collect { case col => r._2.getOrElse(col.toString(), null) } val array = r._1 ++ cols Row(array: _*) } private def getSchema(srcSchema: StructType, idCols: List[String], distinctCols: (String, List[Any])): StructType = { val idSchema = idCols.map { idCol => srcSchema.apply(idCol) } val colSchema = srcSchema.apply(distinctCols._1) val colsSchema = distinctCols._2.map { col => StructField(col.asInstanceOf[String], colSchema.dataType, colSchema.nullable) } StructType(idSchema ++ colsSchema) } 

main snippet

import java.util.Date import org.apache.spark.SparkConf import org.apache.spark.SparkContext import org.apache.spark.sql.Row import org.apache.spark.sql.DataFrame import org.apache.spark.sql.types.StructType import org.apache.spark.sql.hive.HiveContext import org.apache.spark.sql.types.StructField ... ... def main(args: Array[String]): Unit = { val sc = new SparkContext(conf) val sqlContext = new org.apache.spark.sql.SQLContext(sc) val dfdata1 = sqlContext.read.format("com.databricks.spark.csv").option("header", "true").option("inferSchema", "true") .load("data.csv") dfdata1.show() val dfOutput = transpose(new HiveContext(sc), dfdata1, List("id"), "tag", "value") dfOutput.show } 

The built-in spark pivot function is inefficient. The bellow implementation works on spark 2.4+ - the idea is to aggregate a map and extract the values as columns. The only limitation is it does not handle aggregate function in the pivoted columns, only column(s).

On a 8M table, those functions applies on 3 secondes, versus 40 minutes in the built-in spark version:

# pass an optional list of string to avoid computation of columns def pivot(df, group_by, key, aggFunction, levels=[]): if not levels: levels = [row[key] for row in df.filter(col(key).isNotNull()).groupBy(col(key)).agg(count(key)).select(key).collect()] return df.filter(col(key).isin(*levels) == True).groupBy(group_by).agg(map_from_entries(collect_list(struct(key, expr(aggFunction)))).alias("group_map")).select([group_by] + ["group_map." + l for l in levels]) # Usage pivot(df, "id", "key", "value") pivot(df, "id", "key", "array(value)") 

// pass an optional list of string to avoid computation of columns def pivot(df: DataFrame, groupBy: Column, key: Column, aggFunct: String, _levels: List[String] = Nil): DataFrame = { val levels = if (_levels.isEmpty) df.filter(key.isNotNull).select(key).distinct().collect().map(row => row.getString(0)).toList else _levels df .filter(key.isInCollection(levels)) .groupBy(groupBy) .agg(map_from_entries(collect_list(struct(key, expr(aggFunct)))).alias("group_map")) .select(groupBy.toString, levels.map(f => "group_map." + f): _*) } // Usage: pivot(df, col("id"), col("key"), "value") pivot(df, col("id"), col("key"), "array(value)") 

Spark has been providing improvements to Pivoting the Spark DataFrame. A pivot function has been added to the Spark DataFrame API to Spark 1.6 version and it has a performance issue and that has been corrected in Spark 2.0

however, if you are using lower version; note that pivot is a very expensive operation hence, it is recommended to provide column data (if known) as an argument to function as shown below.

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

This has been explained detailed at Pivoting and Unpivoting Spark DataFrame

Happy Learning !!