How to do data transpose using pyspark, no aggregation, but just reshape - pyspark

Question

Again, sorry for asking stupid question, I'm new to Pyspark.
I'm trying to transpose my dataset from:
id dif BA AB
1 1 100 30
1 2 200 40
1 3 300 20
1 4 100 15
2 1 99 18
2 2 89 9
3 1 29 20
3 2 70 32
3 3 39 13
4 1 24 14
4 2 56 23
4 3 24 26
5 1 27 31
5 2 30 17
6 1 100 19
to become
id BA1 BA2 BA3 BA4 AB1 AB2 AB3 AB4
1 100 200 300 100 30 40 20 15
2 99 89 NaN NaN 18 9 NaN NaN
3 29 70 39 NaN 20 32 13 NaN
4 24 56 24 NaN 14 23 26 NaN
5 27 30 NaN NaN 31 17 NaN NaN
6 100 NaN NaN NaN 19 NaN NaN NaN
In python, I used the following code and it works.
df['AB_idx'] = 'AB' + df.dif.astype(str)
AB = df.pivot(index='id',columns='AB_idx',values='AB').reset_index()
df['BA_idx'] = 'BA' + df.dif.astype(str)
BA = df.pivot(index='id',columns='BA_idx',values='BA').reset_index()
df1=pd.merge(BA, AB, on='id', how='left')
Now the problem is I have to translate into pyspark again since I have more than 1 billion records.
any thought? thanks
I do see some example like this post:
Pyspark: reshape data without aggregation
But it seems do not work for me, since I have more than 100 different dif values.

Answer 1

You can try pivot with agg as first and then rename the columns:
output = (df.groupBy("id").pivot("dif")
.agg(*[F.first(i).alias(i) for i in ["BA","AB"]]).orderBy("id"))
d = dict(zip(output.columns,[''.join(i.split('_')[::-1]) for i in output.columns]))
output.select(*[F.col(k).alias(v) for k,v in d.items()]).show()
+---+---+---+----+----+----+----+----+----+
| id|BA1|AB1| BA2| AB2| BA3| AB3| BA4| AB4|
+---+---+---+----+----+----+----+----+----+
| 1|100| 30| 200| 40| 300| 20| 100| 15|
| 2| 99| 18| 89| 9|null|null|null|null|
| 3| 29| 20| 70| 32| 39| 13|null|null|
| 4| 24| 14| 56| 23| 24| 26|null|null|
| 5| 27| 31| 30| 17|null|null|null|null|
| 6|100| 19|null|null|null|null|null|null|
+---+---+---+----+----+----+----+----+----+