scala - 如何注册要在SQL和DataFrame中使用的UDF？

从我所看到的，要做到这一点，你必须

将udf设为普通函数

使用SQLContext注册用于SQL的函数

spark.sqlContext.udf.register("myUDF", myFunc)

将其转换为UserDefinedFunction的DataFrame

def myUDF = udf(myFunc)

是否没有办法将这一步骤组合成一个步骤并使udf可用于两者？另外，对于存在DataFrame而不是SQL的函数的情况，您如何进行注册而不重新复制代码？

最佳答案

带有UDFRegistration.register的scala.FunctionN变体，返回一个UserDefinedFunction，因此您可以在一个步骤中注册SQL函数并创建DSL友好的UDF:

val timesTwoUDF = spark.udf.register("timesTwo", (x: Int) => x * 2)
spark.sql("SELECT timesTwo(1)").show

+---------------+
|UDF:timesTwo(1)|
+---------------+
|              2|
+---------------+

spark.range(1, 2).toDF("x").select(timesTwoUDF($"x")).show

+------+
|UDF(x)|
+------+
|     2|
+------+