Explode multiple pandas columns and unnest one column as column names

Viewed 136

I have the following pandas dataframe:

name score
[A, B, C] [1, 2, 0]
[A, B] [1, 0]
[B, D] [2, 0]
[A, B, C, D] [1, 2,3,4]

I would like to get the following pandas dataframe:

A B C D
1 2 0 NA
1 0 NA NA
NA 2 NA 0
1 2 3 4

So far, I have done the following:

l_df = []
for i in range(len(df)):
    df_ = pd.DataFrame(data  = [df.iloc[i]['score']], columns = df.iloc[i]['name'])
    l_df.append(df_) 
pdf_risk_all = pd.concat(l_df)

However, this takes a long time and it's not good for a dataset with >1MM rows. Any suggestions to do this in a more efficient way?

Thank you,

2 Answers

Let's try:

(pd.concat([df['name'].explode(), df['score'].explode()], axis=1)
   .set_index('name',append=True)
   ['score'].unstack()
)

Output:

name    A  B    C    D
0       1  2    0  NaN
1       1  0  NaN  NaN
2     NaN  2  NaN    0
3       1  2    3    4

In pyspark dataframe, you can zip columns score and name using arrays_zip function, then explode and pivot :

from pyspark.sql import functions as F


df1 = df.withColumn("idx", F.monotonically_increasing_id()) \
    .withColumn("name_score", F.explode(F.arrays_zip("name", "score"))) \
    .select("idx", "name_score.*") \
    .groupBy("idx").pivot("name").agg(first("score")).drop("idx")

df1.show(truncate=False)

#+----+---+----+----+
#|A   |B  |C   |D   |
#+----+---+----+----+
#|1   |2  |0   |null|
#|1   |2  |3   |4   |
#|null|2  |null|0   |
#|1   |0  |null|null|
#+----+---+----+----+
Related