if-statement - if else 在 pyspark 中用于折叠列值

Question

我正在尝试使用简单的代码在索引后将数据框中的分类变量折叠为二进制类目前我的专栏有 3 个类别 - “A”、“B”、“C”我正在编写一个简单的 if else 语句来折叠类，例如

def condition(r):
if (r.wo_flag=="SLM" or r.wo_flag=="NON-SLM"):
    r.wo_flag="dispatch" 
else: 
    r.wo_flag="non_dispatch" 
return r.wo_flag 

df_final=df_new.map(lambda x: condition(x)) 

它不工作，它不理解 else 条件

|MData|Recode12|Status|DayOfWeekOfDispatch|MannerOfDispatch|Wo_flag|PlaceOfInjury|Race|
     M|      11|     M|                  4|               7|      C|           99| 1  |    
     M|       8|     D|                  3|               7|      A|           99| 1  |
     F|      10|     W|                  2|               7|      C|           99| 1  |
     M|       9|     D|                  1|               7|      B|           99| 1  |
     M|       8|     D|                  2|               7|      C|           99| 1  |

这是示例数据

Answer 1

由于使用了用户定义函数 (UDF)，所接受的答案不是很有效。

我认为大多数人都在寻找何时。

from pyspark.sql.functions import when

matches = df["wo_flag"].isin("SLM", "NON-SLM")
new_df = df.withColumn("wo_flag", when(matches, "dispatch").otherwise("non-dispatch"))