Pyspark: я хочу вручную отобразить значения одного из столбцов в моем фрейме данных

Question

У меня есть датафрейм в искре, и я хочу вручную отобразить значения одного из столбцов:

 Col1
  Y
  N
  N
  Y
  N
  Y

Я хочу, чтобы "Y" был равен 1, а "N" был равен 0, например:

 Col1
  1
  0
  0
  1
  0
  1

Я пробовал StringIndexer, но я думаю, что он случайным образом кодирует категориальные данные. (Я не уверен)

Эквивалент Python для этого:

df["Col1"] = df["Col1"].map({"Y": 1, "N": 0})

Подскажите, пожалуйста, как мне этого добиться в Pyspark?

Answer 1

Поскольку вы хотите отобразить значения в 1 и 0, проще всего указать логическое условие и привести результат к int

from pyspark.sql.functions import col
df.withColumn("Col1", (col("Col1")=="Y").cast("int"))

Для более общего случаяВы можете использовать pyspark.sql.functions.when для реализации логики if-then-else :

from pyspark.sql.functions import when
df.withColumn("Col1", when(col("Col1").isin(["Y"]), 1).otherwise(0))