python - pyspark when/otherwise 子句在使用 udf 时失败

Question

我有一个 udf 函数，它接受 key 并从 name_dict 返回相应的 value。

from pyspark.sql import *
from pyspark.sql.functions import udf, when, col

name_dict = {'James': "manager", 'Robert': 'director'}
func = udf(lambda name: name_dict[name])

原始数据框:James 和 Robert 在字典中，但 Michael 不在。

data = [("James","M"),("Michael","M"),("Robert",None)]
test = spark.createDataFrame(data = data, schema = ['name', 'gender'])
test.show()

+-------+------+
|   name|gender|
+-------+------+
|  James|     M|
|Michael|     M|
| Robert|  null|
+-------+------+

为了防止KeyError，我在任何操作之前使用when条件来过滤行，但它不起作用。

test.withColumn('senior', when(col('name').isin(['James', 'Robert']), func(col('name'))).otherwise(col('gender'))).show()

PythonException: An exception was thrown from a UDF: 'KeyError:'Michael'', from , line 8. Full tracebackbelow...

这是什么原因造成的，有什么可行的方法可以解决这个问题吗？假设并非所有名称都是字典的键，对于那些未包含的名称，我想从另一列复制值，在这里说 gender。

Answer 1

这实际上是 Spark 中用户定义函数的行为。您可以从 docs 中读取:

The user-defined functions do not support conditional expressions orshort circuiting in boolean expressions and it ends up with beingexecuted all internally. If the functions can fail on special rows,the workaround is to incorporate the condition into the functions.

因此在您的情况下，您需要将 UDF 重写为:

func = udf(lambda name: name_dict.get(name, "NA"))

然后调用它:

test.withColumn('senior', func(col('name'))).show()

#+-------+------+--------+
#|   name|gender|  senior|
#+-------+------+--------+
#|  James|     M| manager|
#|Michael|     M|      NA|
#| Robert|  null|director|
#+-------+------+--------+

但是，在您的情况下，您实际上可以通过使用映射列来执行此操作而不必使用 udf:

from itertools import chain
from pyspark.sql.functions import col, create_map, lit

map_col = create_map(*[lit(x) for x in chain(*name_dict.items())])
test.withColumn('senior', map_col[col('name')]).show()