Как выбрать столбец на основе значения другого в Pyspark?

Question

У меня есть фрейм данных, где некоторый столбец special_column содержит такие значения, как one, two. В моем фрейме данных также есть столбцы one_processed и two_processed.

. Я хотел бы добавить новый столбец my_new_column, значения которого взяты из других столбцов из моего фрейма данных на основе обработанных значений из special_column. Например, если special_column == one я бы хотел, чтобы my_new_column был установлен на one_processed.

Я пытался .withColumn("my_new_column", F.col(F.concat(F.col("special_column"), F.lit("_processed")))), но Спарк жалуется, что я не могу параметризовать F.col с помощью столбца.

Как я могу получить строковое значение конкатенации, чтобы я мог выбрать нужный столбец?

Answer 1

Самый простой способ в вашем случае - это просто when/oterwise, например:

>>> df = spark.createDataFrame([(1, 2, "one"), (1,2,"two")], ["one_processed", "two_processed", "special_column"]) 
>>> df.withColumn("my_new_column", F.when(F.col("special_column") == "one", F.col("one_processed")).otherwise(F.col("two_processed"))).show()
+-------------+-------------+--------------+-------------+
|one_processed|two_processed|special_column|my_new_column|
+-------------+-------------+--------------+-------------+
|            1|            2|           one|            1|
|            1|            2|           two|            2|
+-------------+-------------+--------------+-------------+

Насколько я знаю, нет способа получить значение столбца по имени, поскольку план выполнения будет зависеть по данным.

Answer 2

from pyspark.sql.functions import when, col, lit, concat_ws

sdf.withColumn("my_new_column", when(col("special_column")=="one", col("one_processed"
).otherwise(concat_ws("_", col("special_column"), lit("processed"))