PySpark pivot as SQL query

user_4918765 0
Viewed 93

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()
pyspark pivot sql
1 Answers

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|
+-----+----+---+----+
user_91089120
Related