python - 带有字典参数的 Spark UDF 失败

Question

我在 Spark 数据框中有一列 (myCol)，它的值为 1,2，我想创建一个新列来描述这些值，例如 1-> 'A'、2->'B' 等

我知道这可以通过连接来完成，但我试过了，因为它看起来更优雅:

dictionary= { 1:'A' , 2:'B' }

add_descriptions = udf(lambda x , dictionary: dictionary[x] if x in dictionary.keys() else None)

df.withColumn("description",add_descriptions(df.myCol,dictionary))

失败并报错

lib/py4j-0.10.4-src.zip/py4j/protocol.py", line 323, in get_return_value
py4j.protocol.Py4JError: An error occurred while calling z:org.apache.spark.sql.function  s.col. Trace:
py4j.Py4JException: Method col([class java.util.HashMap]) does not exist
        at py4j.reflection.ReflectionEngine.getMethod(ReflectionEngine.java:318)
        at py4j.reflection.ReflectionEngine.getMethod(ReflectionEngine.java:339)
        at py4j.Gateway.invoke(Gateway.java:274)
        at py4j.commands.AbstractCommand.invokeMethod(AbstractCommand.java:132)
        at py4j.commands.CallCommand.execute(CallCommand.java:79)
        at py4j.GatewayConnection.run(GatewayConnection.java:214)
        at java.lang.Thread.run(Thread.java:745)

用户定义的函数不能以字典作为参数吗？

Answer 1

这是可能的，你只需要做一些不同的事情。

dictionary= { 1:'A' , 2:'B' }

def add_descriptions(in_dict):
    def f(x):
        return in_dict.get(x)
    return udf(f)

df.withColumn(
    "description",
    add_descriptions(dictionary)(df.myCol)
)

如果你想在 UDF 中直接添加你的字典，因为 UDF 只接受列作为参数，你需要有一个映射列来替换你的字典。