Spark：自定义 udf

Spark SQL 中提供的很多函数可以满足大部分的需求，但有时也需要根据需求自定义函数。项目中有一个需求是，判断某个字段的值是否在某个区间内（前闭后开）或者多个区间内。而 BETWEEN ... AND ... 表示的是前闭后闭，而且只能作用于一个区间。所以就需要实现一个自定义 udf。

自定义 udf 有以下步骤：

• 实现 UDF 接口，根据接收参数个数的不同，分为 UDF0、UDF1、UDF2…，除了接收的参数外，还有一个返回值类型。需要实现 call() 方法，在其中实现具体的逻辑。比如就目前的需求来说，判断一个值在不在区间内，只需返回是或者不是，所以返回值的类型为 Boolean。然后需要考虑接收哪几个参数。因为不同的数据类型比较的方法不同，所以需要把类型传递过来进行判断，然后还有字段的值和要比较的区间。所以有 3 个参数，需要实现 UDF3 接口。当在 SQL 中使用该 udf 时，对于每一行记录，会把字段所对应的值给传递进来。所以我们在 call() 方法中实现相应的比较逻辑即可。另外需要注意的一点是，虽然不同类型的数据比较方法不同，但是大致的逻辑是一样的，可以把比较的主逻辑抽取出来，然后对于不同的类型实现不同的 Comparator 接口，然后把比较对象传递到主逻辑中，这样可以避免使用 if...else... 写过多的重复代码。
• 在使用之前，需要在 SparkSession 中注册 udf，使用 spark.udf().register(name, new UDF(), DataTypes.xxxType); 进行注册。可以在 udf 中用一个静态常量来表示 name，也方便把 name 和定义相关的逻辑放在一起。如此，就可以在 Spark SQL 中使用 udf 了。