如何在Scala中对DataFrame及列应用UDF？新手求助

问题根源 #

你写的kill_4是普通Scala函数，接收String类型参数，但col("Letter")返回的是Spark的Column对象，两者类型不兼容，导致type mismatch错误。Spark的DataFrame API操作的是Column对象，不能直接传入普通Scala函数，必须将函数包装成Spark可识别的UDF。

另外你的函数逻辑和期望输出不符：当前kill_4在输入"A"时返回true，用它过滤会留下"A"这条数据，但你期望保留B、C，所以需要调整逻辑。

正确实现步骤 #

1. 修正函数逻辑并创建UDF #

首先导入Spark UDF相关依赖：

import org.apache.spark.sql.functions._

将普通函数包装为Spark UDF，同时修正逻辑（过滤掉"A"，所以当输入为"A"时返回false，其他返回true）：

// 修正逻辑的普通函数
def keepNonA(n: String): Boolean = n != "A"

// 包装为Spark UDF
val keepNonAUdf = udf(keepNonA _)

或者直接用匿名函数创建UDF，更简洁：

val keepNonAUdf = udf((n: String) => n != "A")

2. 用UDF实现withColumn #

withColumn需要传入Column类型的参数，使用UDF时直接传入col("Letter")即可：

df.withColumn("keep", keepNonAUdf(col("Letter"))).show()

输出结果：

+------+------+-----+
|Letter|Number| keep|
+------+------+-----+
|     A|     1|false|
|     B|     2| true|
|     C|     3| true|
+------+------+-----+

3. 用UDF实现filter过滤 #

直接将UDF传入filter方法，即可得到你期望的结果：

df.filter(keepNonAUdf(col("Letter"))).show()

输出结果：

+------+------+
|Letter|Number|
+------+------+
|     B|     2|
|     C|     3|
+------+------+

额外优化：无需UDF的简化写法 #

如果逻辑简单，其实不需要自定义UDF，直接用Spark内置的Column操作更高效（Spark会对内置表达式做优化）：

// withColumn示例
df.withColumn("keep", col("Letter") =!= "A").show()

// filter示例
df.filter(col("Letter") =!= "A").show()

同样能得到相同结果。

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