Combining values across columns and pyspark and pivoting - pyspark

Question

I have a pypark df like so:
+-------+-------+-----------+-----------+-----------+-----------+-----------+-----------+
| SEQ_ID|TOOL_ID|kurtosis_1m|kurtosis_2m|kurtosis_3m|kurtosis_4m|kurtosis_5m|kurtosis_6m|
+-------+-------+-----------+-----------+-----------+-----------+-----------+-----------+
|3688539| 99725| 6.7484| 6.2753| 6.2055| 7.2076| 7.0501| 7.5099|
|3689076| 99705| 4.8109| 4.3774| 4.1131| 4.4084| 4.1568| 4.4445|
+-------+-------+-----------+-----------+-----------+-----------+-----------+-----------+
I need to pivot it in such a way that I end up with a dataframe like so:
+-------+-------+-----------+
| SEQ_ID|TOOL_ID|kurtosis |
+-------+-------+-----------+
|3688539| 99725| 6.7484|
|3688539| 99725| 6.2753|
|3688539| 99725| 6.2055|
|3688539| 99725| 7.2076|
|3688539| 99725| 7.0501|
|3688539| 99725| 7.5099|
|3689076| 99705| 4.8109|
|3689076| 99705| 4.3774|
|3689076| 99705| 4.1131|
|3689076| 99705| 4.4084|
|3689076| 99705| 4.1568|
|3689076| 99705| 4.4445|
+-------+-------+-----------+
I figured one way would be to create the kurtosis column as an array column and then exploding it. How do you combine the values of columns across a dataframe into a single column as an array.
I have other columns like mean_1m, mean_2m etc... that I need to pivot in the same manner.
Any insights?
Thank You

Answer 1

You can create an array of dataframes and union them.
First, identify the kurtosis columns:
sub_string = "kurtosis"
kurtosis_col = [x for x in df.schema.names if sub_string in x]
Now create an array of dataframes equal to the number of kurtosis columns:
from functools import reduce
from pyspark.sql import DataFrame
from pyspark.sql.functions import col
df_array = [df.withColumn('col', F.concat(F.col(x)))
.select('seq_id', 'tool_id', 'col') for x in kurtosis_col]
# Union them
reduce(DataFrame.unionAll, df_array).withColumnRenamed("col", "kurtosis").show()
Output:
+-------+-------+--------+
| seq_id|tool_id|kurtosis|
+-------+-------+--------+
|3688539| 99725| 6.7484|
|3688539| 99725| 7.2076|
|3688539| 99725| 6.2753|
|3688539| 99725| 6.2055|
|3688539| 99725| 7.5099|
|3688539| 99725| 7.0501|
|3689076| 99705| 4.4084|
|3689076| 99705| 4.1131|
|3689076| 99705| 4.8109|
|3689076| 99705| 4.4445|
|3689076| 99705| 4.3774|
|3689076| 99705| 4.1568|
+-------+-------+--------+
You can follow a similar approach for your other set columns like mean_1m, etc. One way to join them back and to avoid duplicates is by using row_number() and monotonically_increasing_id() before joining. Let me know if you need that piece of code.

Answer 2

You can use array to combine multiple columns into one array and then - like you have already described in your question - explode the array.
from pyspark.sql import functions as F
cols = [x for x in df.schema.names if "kurtosis" in x]
df.withColumn("kurtosis", F.explode(F.array(cols))) \
.drop(*cols) \
.show()
Output:
+-------+-------+--------+
| SEQ_ID|TOOL_ID|kurtosis|
+-------+-------+--------+
|3688539| 99725| 6.7484|
|3688539| 99725| 6.2753|
|3688539| 99725| 6.2055|
|3688539| 99725| 7.2076|
|3688539| 99725| 7.0501|
|3688539| 99725| 7.5099|
|3689076| 99705| 4.8109|
|3689076| 99705| 4.3774|
|3689076| 99705| 4.1131|
|3689076| 99705| 4.4084|
|3689076| 99705| 4.1568|
|3689076| 99705| 4.4445|
+-------+-------+--------+