带有字典参数的Spark UDF失败

时间:2019-07-15 10:04:58

标签: python apache-spark pyspark

我在Spark数据框中有一个列(myCol),其值为1,2,并且我想用该值的描述创建一个新列,例如1->'A',2->'B'等< / p>

我知道这可以通过联接来完成,但是我尝试了一下,因为它看起来更优雅:

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)

使用字典作为参数的用户定义函数是否不可能?

3 个答案:

答案 0 :(得分: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仅接受列作为参数,则需要有一个map列来替换字典。

答案 1 :(得分:1)

如果您使用的是Spark> = 2.4.0,则还可以使用内置map_from_arrays函数来动态创建地图,然后使用getItem获得所需的值,如下所示:

from pyspark.sql.functions import lit, col, map_from_arrays, array
df = spark.createDataFrame([[1],[2],[3]]).toDF("key")

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

map_keys = array([lit(k) for k in dict.keys()])
map_values = array([lit(v) for v in dict.values()])
map_func = map_from_arrays(map_keys, map_values) 

df = df.withColumn("description", map_func.getItem(df.key))

输出:

+---+-----------+
|key|description|
+---+-----------+
|  1|          A|
|  2|          B|
|  3|       null|
+---+-----------+

答案 2 :(得分:1)

以下是通过广播的字典解决此问题的方法(这是解决问题的最可靠方法,因为它也可以与大型词典一起使用):

def add_descriptions(dict_b):
    def f(x):
        return dict_b.value.get(x)
    return F.udf(f)

df = spark.createDataFrame([[1,], [2,], [3,]]).toDF("some_num")
dictionary= { 1:'A' , 2:'B' }
dict_b = spark.sparkContext.broadcast(dictionary)
df.withColumn(
    "res",
    add_descriptions(dict_b)(F.col("some_num"))
).show()
+--------+----+
|some_num| res|
+--------+----+
|       1|   A|
|       2|   B|
|       3|null|
+--------+----+

blog post解释了使用广播词典时的常见陷阱以及如何避免错误。很好的问题,这是PySpark程序员掌握的重要设计模式。