Создайте столбцы в pyspark df из списка, если столбец еще не существует

Question

Вот список значений, я хотел бы, чтобы мой фрейм данных имел:

 cols=['USA','CAN','UK','DEN']

Мой текущий df:

| ID | USA | DEN | VEN | NOR|
| 98 |  1  |  0  | 1   |  1 |
| 99 |  0  |  1  | 0   |  0 |

Я хочу проверить, все ли в моем существующем df Значения в списке в виде столбцов, если нет, то создайте эти столбцы и заполните их нулями, например:

| ID | USA | DEN | VEN | NOR| CAN | UK|
| 98 |  1  |  0  | 1   |  1 |  0  | 0 |
| 99 |  0  |  1  | 0   |  0 |  0  | 0 |

Answer 1

Попробуйте с помощью for + if l oop проверить, существует ли столбец в df.columns или добавить столбец с 0.

from pyspark.sql.functions import *

df=spark.createDataFrame([(98,1,0,1,1,)],['ID','USA','DEN','VEN','NOR'])
cols=['USA','CAN','UK','DEN']

for i in cols:
     if not i in df.columns:
        df=df.withColumn(i,lit("0"))

df.show()

#+---+---+---+---+---+---+---+
#| ID|USA|DEN|VEN|NOR|CAN| UK|
#+---+---+---+---+---+---+---+
#| 98|  1|  0|  1|  1|  0|  0|
#+---+---+---+---+---+---+---+

Answer 2

Вы можете использовать простое выражение выбора:

from pyspark.sql.functions import lit

select_cols = df.columns + [lit(0).alias(c) for c in cols if c not in df.columns]

df.select(*select_cols).show()