Использование Spark Pivot для get_dummies Substitute

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)