Is there Spark equivalent for Pandas MultiIndex operation like set_index() or unstack()?

user_13788929 0
Viewed 775

I have a Spark DataFrame like:

+---+---------+------+
| id| timeSlot| ratio|
+---+---------+------+
|100|    lunch|   0.2|
|100|  morning|   0.3|
|100|   dinner|   0.5|
|101|  morning|   0.1|
|101|    lunch|   0.2|
|101|   dinner|   0.7|
|102|    lunch|   0.4|
|102|  morning|   0.3|
|102|   dinner|   0.3|
+---+---------+------+

I want to transform it to the following format: pivot by timeSlot in columns, id as row-index, and aggregate (sum) values of ratio:

+---+--------+------+-------+
| id| morning| lunch| dinner|
+---+--------+------+-------+
|100|     0.3|   0.2|    0.5|
|101|     0.1|   0.2|    0.7|
|102|     0.3|   0.4|    0.3|
+---+--------+------+-------+

For such an operation, in pandas I use MultiIndex methods such as set_index() and unstack():

import pandas as pd

ids = pd.Series([100, 100, 100, 101, 101, 101, 102 ,102, 102])
timeSlots = pd.Series(["lunch", "morning", "dinner", "morning", 
                       "lunch", "dinner", "lunch" ,"morning", "dinner"])
ratios = pd.Series([0.2,0.3,0.5,0.1,0.2,0.7,0.4,0.3,0.3])
df = pd.DataFrame({"id": ids, "timeSlot": timeSlots, "ratio": ratios}) # create initial data frame

df = df.set_index(['id', 'timeSlot']) # main goal1 : Set 'id' and 'timeSlot' as multi indexes.
df = df.unstack() # main goal2 : Switch 'timeSlot' to column index
df = df['ratio'] # remove one layer of column index
df = df[['morning', 'lunch', 'dinner']] # reorder columns

Is there any Spark equivalent method for Pandas MultiIndex operation like set_index() or unstack()?

If not, is there any other method to achieve the same goal?

pyspark pandas apache-spark-sql apache-spark python
1 Answers

You can use groupby and pivot in pyspark

import pandas as pd
import pyspark
from pyspark.sql import SparkSession
import pyspark.sql.functions as F

spark = SparkSession.builder.getOrCreate()
# your initial df
df = pd.DataFrame({"id": ids, "timeSlot": timeSlots, "ratio": ratios}) 
# convert pandas df to spark df
spark_df = spark.createDataFrame(df)


# combine groupby, pivot and agg methods to get your result
spark_df.groupby("id")\
        .pivot("timeSlot")\
        .agg(F.sum("ratio"))\ 
        .show()
#+---+------+-----+-------+
#| id|dinner|lunch|morning|
#+---+------+-----+-------+
#|100|   0.5|  0.2|    0.3|
#|101|   0.7|  0.2|    0.1|
#|102|   0.3|  0.4|    0.3|
#+---+------+-----+-------+

Please note that I have used sum as an aggregation function but as you only have 1 row by (id, timeSlot) combination you could also have used other functions such as mean, max...

user_33011130
Related