Pyspark udf在条件定义上返回一列,接受多个列作为输入

Question

我正在使用spark 2.1,用法是pyscripting

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

ABC

SSS

S NS NS

S NS S.

SS NS

NS S NS

我的输出必须如下

A B C D

SSSS

S NS NS NS

S NS SS

SS 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())

任何人都可以帮我完成这个逻辑.

Answer 1

我试图使用内置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)

Answer 2

它应该是:

@udf
def return_string(a, b, c):
    if a == 's' and b == 'S' and c == 's':
        return 'S'
    if a == 's' and b == 'NS' and c == 's':
        return 'S'
    if a == 's' and b == 'S' and c == 'NS':
        return 'NS'

df = sc.parallelize([('s', 'S', 'NS'), ('?', '?', '?')]).toDF(['a', 'b', 'c'])

df.withColumn('result', return_string('a', 'b', 'c')).show()
## +---+---+---+------+
## |  a|  b|  c|result|
## +---+---+---+------+
## |  s|  S| NS|    NS|
## |  ?|  ?|  ?|  null|
## +---+---+---+------+

• 应列出所有参数(除非您将数据作为传递struct).
• 您应该使用and不&(您评估逻辑表达式不是SQL表达式).
• 条件应该是表达式而不是列表(非空列表总是真实的).

就个人而言,我会跳过所有ifs并使用简单的dict:

@udf
def return_string(a, b, c):
    mapping = {
        ('s', 'S', 's'): 'S',
        ('s', 'NS' 's'): 'S',
        ('s', 'S', 'NS'): 'NS',
    }
    return mapping.get((a, b, c))

根据您的要求调整条件.

总的来说,您应该更喜欢SQL表达式,如Steven Laan提供的优秀答案所示(您可以链接多个条件when(..., ...).when(..., ...)).