使用Spark Pivot獲得get_dummies替換

Question

以下代碼以完美的格式設置數據框，但我需要讓它「正確」命名列。

df = spark.createDataFrame([ 
    (0, "X", "a"), 
    (1, "Z", "b"), 
    (2, "X", "b"), 
    (3, "X", "c"), 
    (4, "Y", "c"), 
    (5, "Y", "a") 
], ["id","category", "other_thing"]) 

pivotDF = df.groupBy("id").pivot("category").count() 

pivotDF.show() 

+---+----+----+----+ 
| id| X| Y| Z| 
+---+----+----+----+ 
| 0| 1|null|null| 
| 5|null| 1|null| 
| 1|null|null| 1| 
| 3| 1|null|null| 
| 2| 1|null|null| 
| 4|null| 1|null| 
+---+----+----+----+ 

我需要這個作爲輸出：

+---+-------------+-------------+-------------+ 
| id| category_X| category_Y| category_Z| 
+---+-------------+-------------+-------------+ 
| 0| 1   |   null|   null| 
| 5|null   |   1|   null| 
| 1|null   |   null|   1| 
| 3| 1   |   null|   null| 
| 2| 1   |   null|   null| 
| 4|null   |   1|   null| 
+---+-------------+-------------+-------------+ 

如何添加列名編程方式（即我沒有在這種情況下手動鍵入「類別」

Answer 1

可以重命名：

>>> pivot_col = "category" 
>>> pivotDF = df.groupBy("id").pivot(pivot_col).count() 
>>> new_names = pivotDF.columns[:1] + \ 
... ["{0}_{1}".format(pivot_col, c) for c in pivotDF.columns[1:]] 
>>> pivotDF.toDF(*new_names)