如何在Scala编写的Spark UDF中返回自定义Schema实例或空值

你的代码中if表达式缺少else分支，条件不满足时默认返回Unit类型，导致map操作返回的集合元素类型被Scala推断为Any（Option[ReturnSchema]和Unit的共同父类），Spark无法识别Any类型对应的Schema，因此抛出报错。

方案1：保留所有元素，满足条件返回Some(实例)，不满足返回None #

补全else分支，统一返回类型为Option[ReturnSchema]：

import org.apache.spark.sql.functions.udf
import java.util.concurrent.TimeUnit

case class ReturnSchema(a: String, b: String,c: String, d: java.sql.Timestamp)

val myUDF = udf((abc: java.sql.Timestamp, xyz1: Seq[Row]) => {
    xyz1.map(xyz => {
        val a = xyz.getAs[String]("a")
        val b = xyz.getAs[String]("b")
        val c = xyz.getAs[String]("c")
        val d = xyz.getAs[java.sql.Timestamp]("d")
        val timeDifferenceInMinutes = TimeUnit.MINUTES.convert(d.getTime - abc.getTime, TimeUnit.MILLISECONDS)
        if(timeDifferenceInMinutes >= 15) {
            Some(ReturnSchema(a, b, c, d))
        } else {
            None
        }
    })
})

该方案返回值类型为Seq[Option[ReturnSchema]]，Spark会自动推断对应的Schema为ArrayType(StructType(四个字段定义), containsNull = true)。

方案2：仅保留满足条件的实例，过滤不满足的记录 #

如果不需要保留None占位，直接用flatMap过滤无效值，返回类型为Seq[ReturnSchema]：

val myUDF = udf((abc: java.sql.Timestamp, xyz1: Seq[Row]) => {
    xyz1.flatMap(xyz => {
        val a = xyz.getAs[String]("a")
        val b = xyz.getAs[String]("b")
        val c = xyz.getAs[String]("c")
        val d = xyz.getAs[java.sql.Timestamp]("d")
        val timeDifferenceInMinutes = TimeUnit.MINUTES.convert(d.getTime - abc.getTime, TimeUnit.MILLISECONDS)
        if(timeDifferenceInMinutes >= 15) {
            Some(ReturnSchema(a, b, c, d))
        } else {
            None
        }
    })
})

该方案返回的Schema为ArrayType(StructType(四个字段定义), containsNull = false)，性能更优，更符合大多数业务场景需求。

如果遇到Spark类型推断异常，可以显式给udf指定返回类型，进一步避免运行时错误。

内容的提问来源于stack exchange，提问作者Lumos