如何在pyspark中向sqlContext添加udf

时间:2018-04-13 16:56:32

标签: python apache-spark user-defined-functions

我知道我可以将Python函数注册为UDF并在SQL查询中使用它:

udf

或者我可以用from pyspark.sql.functions import udf from pyspark.sql.types import IntegerType example_udf = udf(example) data.select(example_udf('col')) 包装Python函数,因此它可以应用于数据帧:

from pyspark.sql.types import BooleanType
from pyspark.sql.functions import col
def my_udf(other_par)
    def example(s):
        return len(s) == other_par
    return udf(example, BooleanType())

dataframe.select(...).where(my_udf(5)(col('col')))

在我的情况下,因为我需要将一些其他参数传递给UDF,所以我为UDF构建了一个嵌套函数:

sqlContext.udf.register

现在我已经有了一个UDF,我可以在数据帧上应用它。但我也想在spark.sql中使用它,就像第一个块中的SQL查询一样,而不是数据帧的select或where方法。所以我想知道我该怎么做。看起来=SUMIFS($J:$J,$G:$G,"<"&N3,$H:$H,">="&M3,$I:$I,">="&NETWORKDAYS($G:$G,N3)/30)只能接受Python函数而不是UDF。

1 个答案:

答案 0 :(得分:2)

如果您使用最新且最好的(2.3),请不要直接使用udf

def my_udf(other_par, spark):
    def _(s):
        return len(s) == other_par
    return spark.udf.register("my_udf_{}".format(other_par), _, BooleanType())

my_udf_42 = my_udf(42, spark)

spark.sql("SELECT my_udf_42(array(1, 2))").show()
# +----------------------+
# |my_udf_42(array(1, 2))|
# +----------------------+
# |                 false|
# +----------------------+

spark.createDataFrame([([1] * 42, )], ("id", )).select(my_udf_42("id")).show()
# +-------------+
# |my_udf_42(id)|
# +-------------+
# |         true|
# +-------------+

否则直接调用注册副作用:

def my_udf(other_par, spark):
    def _(s):
        return len(s) == other_par
    name = "my_udf_{}".format(other_par)
    spark.udf.register(name, _, BooleanType())
    return udf(_, BooleanType())

my_udf_0 = my_udf(0, spark)

spark.sql("SELECT my_udf_0(array())").show()
# +-----------------+
# |my_udf_0(array())|
# +-----------------+
# |             true|
# +-----------------+

当然,这样的简单操作不应该用udf完成,但我认为这只是一个玩具的例子。如果没有,

from pyspark.sql.functions import size, length

size("some_col") == 42
length("some_col") == 42

是更好的选择。