How to do reshape operation on a single column in pyspark dataframe?

user_15458124 0
Viewed 207

I have a long pyspark dataframe as shown below:

+------+
|number|
+------+
|12.4  |
|13.4  |
|42.3  |
|33.4  |
|42.3  |
|32.4  |
|44.2  |
|12.3  |
|45.4  |
+------+

Ideally, I want this to be reshaped to an nxn matrix where n is sqrt(length of pyspark dataframe).

While there is a solution by converting it into a numpy array and then reshaping it to nxn matrix but I want that to be done in pyspark. Because my data is super long (about 100 million rows).

So the expected output I am looking for is something along these lines:

+------+------+------+
|12.4  | 13.4 | 42.3 |
|33.4  | 42.3 | 32.4 |
|44.2  | 12.3 | 45.4 |
+------+------+------+

While I was able to do it properly by converting it to pandas then to numpy and then doing the reshape operation. But I want to do this transformation in Pyspark itself. Because the below code works fine for only a few thousand rows.

covarianceMatrix_pd = covarianceMatrix_df.toPandas()
nrows = np.sqrt(len(covarianceMatrix_pd))
covarianceMatrix_pd = covarianceMatrix_pd.to_numpy().reshape((int(nrows),int(nrows)))
covarianceMatrix_pd
pyspark pandas apache-spark-sql apache-spark numpy
1 Answers

One way to do this is using row_number with pivot after we have a count of the dataframe:

from pyspark.sql import functions as F, Window
from math import sqrt

c = int(sqrt(df.count())) #this gives 3
rnum = F.row_number().over(Window.orderBy(F.lit(1)))

out = (df.withColumn("Rnum",((rnum-1)/c).cast("Integer"))
 .withColumn("idx",F.row_number().over(Window.partitionBy("Rnum").orderBy("Rnum")))
.groupby("Rnum").pivot("idx").agg(F.first("number")))

out.show()

+----+----+----+----+
|Rnum|   1|   2|   3|
+----+----+----+----+
|   0|12.4|13.4|42.3|
|   1|33.4|42.3|32.4|
|   2|44.2|12.3|45.4|
+----+----+----+----+
user_98406370
Related