Pivot dataframe in pyspark using column for suffix - pyspark

Question

This question is similar to one I've asked before (Pandas pivot ussing column as suffix) but this time I need to do it using Pyspark instead of Pandas. The problem is as follows.
I have a dataframe like the following example:
Id
Type
Value_1
Value_2
1234
A
1
2
1234
B
1
2
789
A
1
2
789
B
1
2
567
A
1
2
And I want to transform to get the following:
Id
Value_1_A
Value_1_B
Value_2_A
Value_2_B
1234
1
1
2
2
789
1
1
2
2
567
1
1
In summary: replicating the value columns using the 'Type' column as a suffix and convert the dataframe to a wide format.
One solution I can think of is creating the columns with the suffix manually and then aggregating.
Other solutions I've tried are using pyspark GroupedData pivot function as follows:
import pandas as pd
df = spark.createDataFrame(pd.DataFrame({'Id': {0: 1234, 1: 1234, 2: 789, 3: 789, 4: 567},
'Type': {0: 'A', 1: 'B', 2: 'A', 3: 'B', 4: 'A'},
'Value_1': {0: 1, 1: 1, 2: 1, 3: 1, 4: 1},
'Value_2': {0: 2, 1: 2, 2: 2, 3: 2, 4: 2}}))
df.groupBy("Id").pivot("Type").avg().show()
The issue of this solution is that the resulting dataframe would include the Id column repeated 3 times, and the inability to name the columns adding Type as the suffix, since they would be named liked this:
['Id',
'A_avg(Id)',
'A_avg(Value_1)',
'A_avg(Value_2)',
'B_avg(Id)',
'B_avg(Value_1)',
'B_avg(Value_2)']
I also tried specifying the value columns to the pivot functions as follows
df.groupBy("Id").pivot("Type", values=["Value_1", "Value_2"]).avg().show()
This removes the extra Id columns, but the rest of the columns only have null values.
Is there any elegant way to do the transformation I'm attempting on pyspark?

Answer 1

Option 1:
If you don't mind having your Type values as column prefixes rather than suffixes, you can use a combination of agg, avg, and alias:
import pyspark.sql.functions as F
df_pivot = df \
.groupBy("Id") \
.pivot("Type") \
.agg(F.avg("Value_1").alias("Value_1"), F.avg("Value_2").alias("Value_2"))
+----+---------+---------+---------+---------+
|Id |A_Value_1|A_Value_2|B_Value_1|B_Value_2|
+----+---------+---------+---------+---------+
|789 |1.0 |2.0 |1.0 |2.0 |
|567 |1.0 |2.0 |null |null |
|1234|1.0 |2.0 |1.0 |2.0 |
+----+---------+---------+---------+---------+
Separately, it's worth noting here that the values argument in the pivot method is used to limit which values you want to retain from your pivot (i.e., Type) column. For example, if you only wanted A and not B in your output, you would specify pivot("Type", values=["A"]).
Option 2:
If you do still want them as suffixes, you'll likely have to use some regex and withColumnRenamed, which could look something like this:
import pyspark.sql.functions as F
import re
df_pivot = df \
.groupBy("Id") \
.pivot("Type") \
.agg(F.avg("Value_1"), F.avg("Value_2"))
for col in df_pivot.columns:
if "avg(" in col:
suffix = re.findall("^.*(?=_avg\()|$", col)[0]
base_name = re.findall("(?<=\().*(?=\)$)|$", col)[0]
df_pivot = df_pivot.withColumnRenamed(col, "_".join([base_name, suffix]))
+----+---------+---------+---------+---------+
|Id |Value_1_A|Value_2_A|Value_1_B|Value_2_B|
+----+---------+---------+---------+---------+
|789 |1.0 |2.0 |1.0 |2.0 |
|567 |1.0 |2.0 |null |null |
|1234|1.0 |2.0 |1.0 |2.0 |
+----+---------+---------+---------+---------+