Question

我使用的是spark 2.1，用法是pyscripting

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

a b c

S S S

NS NS

S NS S

S S NS

NS S NS

我的输出必须如下

a b c d

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

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

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提供的the excellent answer所示（您可以使用when(..., ...).when(..., ...)链接多个条件）。

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

2 个答案: