PySpark pivot as SQL query - sql

Question

Looking to write the full-SQL equivalent of a pivot implemented in pyspark. Code below creates a pandas DataFrame.
import pandas as pd
df = pd.DataFrame({
'id': ['a','a','a','b','b','b','b','c','c'],
'name': ['up','down','left','up','down','left','right','up','down'],
'count': [6,7,5,3,4,2,9,12,4]})
# id name count
# 0 a up 6
# 1 a down 7
# 2 a left 5
# 3 b up 3
# 4 b down 4
# 5 b left 2
# 6 b right 9
# 7 c up 12
# 8 c down 4
Code below then converts to a pyspark DataFrame and implements a pivot on the name column.
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
ds = spark.createDataFrame(df)
dp = ds.groupBy('id').pivot('name').max().toPandas()
# id down left right up
# 0 c 4 NaN NaN 12
# 1 b 4 2.0 9.0 3
# 2 a 7 5.0 NaN 6
Trying to do the equivalent of ds.groupBy('id').pivot('name').max() in full-SQL, ie something like
ds.createOrReplaceTempView('ds')
dp = spark.sql(f"""
SELECT * FROM ds
PIVOT
(MAX(count)
FOR
...)""").toPandas()

Answer 1

Taking reference from SparkSQL Pivot -
Pivot
spark.sql(f"""
SELECT * FROM ds
PIVOT (
MAX(count)
FOR name in ('up','down','left','right')
)""").show()
+---+---+----+----+-----+
| id| up|down|left|right|
+---+---+----+----+-----+
| c| 12| 4|null| null|
| b| 3| 4| 2| 9|
| a| 6| 7| 5| null|
+---+---+----+----+-----+
Dynamic Approach
You can dynamically create the PIVOT Clause , I tried to create a general wrapper around the same
def pivot_by(inp_df,by):
distinct_by = inp_df[by].unique()
distinct_name_str = ''
for i,name in enumerate(distinct_by):
if i == 0:
distinct_name_str += f'\'{name}\''
else:
distinct_name_str += f',\'{name}\''
final_str = f'FOR {by} in ({distinct_name_str})'
return final_str
pivot_clause_str = pivot_by(df,'name')
### O/p - FOR name in ('up','down','left','right')
spark.sql(f"""
SELECT * FROM ds
PIVOT (
MAX(count)
{pivot_clause_str}
)""").show()
+---+---+----+----+-----+
| id| up|down|left|right|
+---+---+----+----+-----+
| c| 12| 4|null| null|
| b| 3| 4| 2| 9|
| a| 6| 7| 5| null|
+---+---+----+----+-----+
Dynamic Approach Usage
pivot_clause_str = pivot_by(df,'id')
##O/p - FOR id in ('a','b','c')
ds.createOrReplaceTempView('ds')
sql.sql(f"""
SELECT * FROM ds
PIVOT (
MAX(count)
{pivot_clause_str}
)""").show()
+-----+----+---+----+
| name| a| b| c|
+-----+----+---+----+
| down| 7| 4| 4|
| left| 5| 2|null|
| up| 6| 3| 12|
|right|null| 9|null|
+-----+----+---+----+