如何根据pyspark中数据框的值创建列

Question

I am having a df that contains below values, would like to create columns for each values present in the df. Looking for a solution in pyspark.Basically I could do this with a case when in pyspark, looking for a different approach. Any suggestions will be helpful.

DF:

|number|color|
|------|-----|
|123   |red  |
|234   |blue |
|555   |white|

Expected output:

number	red	blue	white
123	1	0	0
234	0	1	0
555	0	0	1

Answer 1

You can group by number , pivot by color and apply the lit(1) value. To removing null values, apply [dataframe].na.fill(0)

import pyspark.sql.functions as f

df = spark.createDataFrame([
  [123, 'red'],
  [234, 'blue'],
  [555, 'white']
], ['number', 'color'])

pivot_df = df.groupBy('number').pivot('color').agg(f.lit(1))
pivot_df = pivot_df.na.fill(0)

(pivot_df
 .select('number', 'red', 'blue', 'white')
 .sort('number')
 .show(truncate=False))
# +------+---+----+-----+
# |number|red|blue|white|
# +------+---+----+-----+
# |123   |1  |0   |0    |
# |234   |0  |1   |0    |
# |555   |0  |0   |1    |
# +------+---+----+-----+