How to pivot a table with dynamic columns in pyspark

user_9405212 0
Viewed 772

Is there a possibility to make a pivot for different columns at once in Pyspark? I have a dataframe like this:

sdf = spark.createDataFrame(
    pd.DataFrame([[1,2,6,1],[1,3,3,2],[1,6,0,3],[2,1,0,1],
        [2,1,7,2],[2,7,8,3]], columns = ['id','val1','val2','month'])
)
+----+------+------+-------+
| id | val1 | val2 | month |
+----+------+------+-------+
|  1 |   2  |   6  |   1   |
|  1 |   3  |   3  |   2   |
|  1 |   6  |   0  |   3   |
|  2 |   1  |   0  |   1   |
|  2 |   1  |   7  |   2   |
|  2 |   7  |   8  |   3   |
+----+------+------+-------+

This dataframe I want to pivot on multiple columns (val1, val2, ...) to have a dataframe that looks like this:

+----+-------------+-------------+-------------+-------------+-------------+-------------+
| id | val1_month1 | val1_month2 | val1_month3 | val2_month1 | val2_month2 | val2_month3 |
+----+-------------+-------------+-------------+-------------+-------------+-------------+
|  1 |           2 |           3 |           6 |           6 |           3 |           0 |
|  2 |           1 |           1 |           7 |           0 |           7 |           8 |
+----+-------------+-------------+-------------+-------------+-------------+-------------+

I've found a solution that works for hard-coded columns (see below), but I'm looking for a solution that can take val1, val2, etc dynamically.

sdf_pivot = (
    sdf
    .groupby('id')
    .pivot('month')
    .agg(sf.mean('val1'),sf.mean('val2'))
)

Something like this, but unfortunately this does not work...

col_to_pivot = ['val1','val2']
sdf_pivot = (
    sdf
    .groupby('id')
    .pivot('month')
    .agg(sf.mean(col_to_pivot))
)

Many thanks!

pyspark multiple-columns pivot python
1 Answers

IIUC, you can use a list comprehension:

newdf = sdf.groupby('id').pivot('month').agg(*[ sf.mean(c).alias(c) for c in col_to_pivot ])
#+---+------+------+------+------+------+------+                                 
#| id|1_val1|1_val2|2_val1|2_val2|3_val1|3_val2|
#+---+------+------+------+------+------+------+
#|  1|     2|     6|     3|     3|     6|     0|
#|  2|     1|     0|     1|     7|     7|     8|
#+---+------+------+------+------+------+------+

col_names = [ '{}_month{}'.format(x[1],x[0]) if len(x)>1 else x[0] for c in newdf.columns for x in [c.split('_')] ]
#['id',
# 'val1_month1',
# 'val2_month1',
# 'val1_month2',
# 'val2_month2',
# 'val1_month3',
# 'val2_month3']

newdf = newdf.toDF(*col_names)
user_95107290
Related