python之Pyspark udf 在接受多列作为输入的条件定义上返回一列

我正在使用 spark 2.1 并且使用的是 pyscripting

问题陈述:有一个场景需要传递多列作为输入并返回一列作为输出，下面是我的 3 列输入数据框

a b c

S S S

S NS NS

S NS S

S S NS

NS S NS

我的输出必须如下所示

a b c d

S S S S S

S NS NS NS

S NS S S

S S NS NS

NS S NS NS

我正在尝试注册一个 UDF 以将这 3 列 [a,b,c] 作为输入传递并返回 d 列作为输出，这里 a,b,c,d 是列名

我发现很难得到下面的输出是使用的语法

def return_string(x): 
      if [x.a=='s' & x.b=='S' & x.c=='s'] 
          return 'S' 
      else if[x.a=='s' & x.b=='NS' & x.c=='s'] 
          return 'S' 
      else if[x.a=='s' & x.b=='S' & x.c=='NS'] 
          return 'NS; 
 
func= udf(returnstring,types.StringType()) 

谁能帮我完成这个逻辑。

请您参考如下方法:

我试图使用内置的 withColumn 来做到这一点和 when功能:

from pyspark.sql.functions import col, when, lit 
 
df.withColumn('d', when( 
     ((col('A') == 'S') & (col('B') == 'S') & (col('C')=='S')) 
   | ((col('A') == 'S') & (col('B') == 'NS') & (col('C')=='S')) 
 , lit('S') 
 ).otherwise(lit('NS')) 
).show() 

这也假设这两个值是互斥的(因此 otherwise)