spark使用udf给dataFrame新增列

在 spark 中给 dataframe 增加一列的方法一般使用 withColumn

// 新建一个dataFrame
val sparkconf = new SparkConf()
.setMaster("local")
.setAppName("test")
val spark = SparkSession.builder().config(sparkconf).getOrCreate()
val tempDataFrame = spark.createDataFrame(Seq(
(1, "asf"),
(2, "2143"),
(3, "rfds")
)).toDF("id", "content")
// 增加一列
val addColDataframe = tempDataFrame.withColumn("col", tempDataFrame("id")*0)
addColDataframe.show(10,false)


打印结果如下：

+---+-------+---+
|id |content|col|
+---+-------+---+
|1 |asf |0 |
|2 |2143 |0 |
|3 |rfds |0 |
+---+-------+---+


可以看到 withColumn 很依赖原来 dataFrame 的结构，但是假设没有 id 这一列，那么增加列的时候灵活度就降低了很多，假设原始 dataFrame 如下：

+---+-------+
| id|content|
+---+-------+
| a| asf|
| b| 2143|
| b| rfds|
+---+-------+


这样可以用 udf 写自定义函数进行增加列：

import org.apache.spark.sql.functions.udf
// 新建一个dataFrame
val sparkconf = new SparkConf()
.setMaster("local")
.setAppName("test")
val spark = SparkSession.builder().config(sparkconf).getOrCreate()
val tempDataFrame = spark.createDataFrame(Seq(
("a, "asf"),
("b, "2143"),
("c, "rfds")
)).toDF("id", "content")
// 自定义udf的函数
val code = (arg: String) => {
if (arg.getClass.getName == "java.lang.String") 1 else 0
}val addCol = udf(code)
// 增加一列
val addColDataframe = tempDataFrame.withColumn("col", addCol(tempDataFrame("id")))
addColDataframe.show(10, false)


得到结果：

+---+-------+---+
|id |content|col|
+---+-------+---+
|a |asf |1 |
|b |2143 |1 |
|c |rfds |1 |
+---+-------+---+


还可以写下更多的逻辑判断：

// 新建一个dataFrame
val sparkconf = new SparkConf()
.setMaster("local")
.setAppName("test")
val spark = SparkSession.builder().config(sparkconf).getOrCreate()
val tempDataFrame = spark.createDataFrame(Seq(
(1, "asf"),
(2, "2143"),
(3, "rfds")
)).toDF("id", "content")val code :(Int => String) = (arg: Int) => {if (arg < 2) "little" else "big"}
val addCol = udf(code)
val addColDataframe = tempDataFrame.withColumn("col", addCol(tempDataFrame("id")))
addColDataframe.show(10, false)


+---+-------+------+
|1 |asf |little|
|2 |2143 |big |
|3 |rfds |big |
+---+-------+------+


传入多个参数：

val sparkconf = new SparkConf()
.setMaster("local")
.setAppName("test")
val spark = SparkSession.builder().config(sparkconf).getOrCreate()
val tempDataFrame = spark.createDataFrame(Seq(
("1", "2"),
("2", "3"),
("3", "1")
)).toDF("content1", "content2")val code = (arg1: String, arg2: String) => {
Try(if (arg1.toInt > arg2.toInt) "arg1>arg2" else "arg1<=arg2").getOrElse("error")
}
val compareUdf = udf(code)val addColDataframe = tempDataFrame.withColumn("compare", compareUdf(tempDataFrame("content1"),tempDataFrame("content2")))
addColDataframe.show(10, false)


+--------+--------+----------+
|content1|content2|compare |
+--------+--------+----------+
|1 |2 |arg1<=arg2|
|2 |3 |arg1<=arg2|
|3 |1 |arg1>arg2 |
+--------+--------+----------+