Spark DataFrame列赋值疑问：UDF入参作用与通用实现

疑问1：为什么传入col("b")或col("a")结果相同，却需要传入该列？ #

咱们先拆解你的UDF逻辑：

def myfun: Int => List[Int] = _ => uniqueVal
def myfun_udf = udf(myfun)

这里的myfun是个完全忽略输入参数的函数——不管你传什么Int类型的值，它都会直接返回提前计算好的uniqueVal列表。所以：

• 当你传入col("b")或col("a")时，这两列都是Int类型，UDF能正常接收参数（哪怕完全没用上），自然返回结果一致。
• 而你必须传入列的原因是：Spark的UDF调用语法要求必须绑定到列表达式上，哪怕你的逻辑根本不需要这个列的值。如果不传列，withColumn里的UDF调用会直接报错，因为UDF需要一个输入数据源（列）来触发执行流程。

疑问2：传入col("ID")时结果为null，如何实现对所有列通用的逻辑？ #

出现null的核心原因是col("ID")是字符串类型，而你的UDF定义的输入类型是Int，类型不匹配导致UDF执行失败，最终返回null。要实现通用逻辑，有两种更优雅的方案：

方法1：放弃UDF，用Spark内置函数直接生成数组（推荐）

既然uniqueVal是提前计算好的固定列表，完全不需要依赖DataFrame的任何列，直接用array函数把列表转成列即可——性能比UDF高很多（内置函数是Spark原生优化过的，没有Java序列化的额外开销）：

import spark.implicits._
import org.apache.spark.sql.functions._
val df = sc.parallelize(Seq( ("r1", 1, 1), ("r2", 6, 4), ("r3", 4, 1), ("r4", 1, 2) )).toDF("ID", "a", "b")
val uniqueVal = df.select("b").distinct().map(x => x.getAs[Int](0)).collect.toList
// 把uniqueVal转换为Spark可识别的数组列
val uniqueArrayCol = array(uniqueVal.map(lit(_)): _*)
df.withColumn("X", uniqueArrayCol).show

这样不管DataFrame有什么类型的列，都能正常生成X列，完全不依赖任何输入列。

方法2：修改UDF的输入类型为通用类型

如果你一定要用UDF，可以把UDF的输入类型改成Any，这样不管传入什么类型的列，都能正常接收并返回数组：

import spark.implicits._
import org.apache.spark.sql.functions._
val df = sc.parallelize(Seq( ("r1", 1, 1), ("r2", 6, 4), ("r3", 4, 1), ("r4", 1, 2) )).toDF("ID", "a", "b")
val uniqueVal = df.select("b").distinct().map(x => x.getAs[Int](0)).collect.toList
// 定义接收任意类型输入的UDF
def myfun: Any => List[Int] = _ => uniqueVal
def myfun_udf = udf(myfun)
// 现在传入任何列都能得到正确结果
df.withColumn("X", myfun_udf(col("ID"))).show

不过这种方法还是不如内置函数高效，优先推荐方法1。

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