在 Pyspark 中将（转换）列转换为行

Question

我有以下格式的 spark 数据框，其中每个唯一 id 最多可以有 3 行，由 rank 列给出。

 id pred    prob      rank
485 9716    0.19205872  1
729 9767    0.19610429  1
729 9716    0.186840048 2
729 9748    0.173447074 3
818 9731    0.255104463 1
818 9748    0.215499913 2
818 9716    0.207307154 3

我想将（转换）成行数据，这样每个 id 只有一行，而 pred 和 prob 列有多个由等级变量（列后缀）区分的列。

id  pred_1  prob_1      pred_2  prob_2     pred_3   prob_3
485 9716    0.19205872              
729 9767    0.19610429  9716    0.186840048 9748    0.173447074
818 9731    0.255104463 9748    0.215499913 9716    0.207307154

我无法弄清楚如何在Pyspark 中使用它

输入数据创建示例代码：

# Loading the requisite packages 
from pyspark.sql.functions import col, explode, array, struct, expr, sum, lit        
# Creating the DataFrame
df = sqlContext.createDataFrame([(485,9716,19,1),(729,9767,19,1),(729,9716,18,2), (729,9748,17,3), (818,9731,25,1), (818,9748,21,2), (818,9716,20,3)],('id','pred','prob','rank'))
df.show()

Answer 1

这是多列问题的关键。尝试：

import pyspark.sql.functions as F

df_pivot = df.groupBy('id').pivot('rank').agg(F.first('pred').alias('pred'), F.first('prob').alias('prob')).orderBy('id')
df_pivot.show(truncate=False)