Scala Spark，如何为列添加值

Posted 2023-04-13

【中文标题】Scala Spark，如何为列添加值

【英文标题】：Scala Spark, how to add value to the column

【发布时间】：2017-08-25 19:46:22

【问题描述】：

我的目标是向 DataFrame 的给定列添加一个可配置的常量值。

val df = Seq(("A", 1), ("B", 2), ("C", 3)).toDF("col1", "col2")

+----+----+
|col1|col2|
+----+----+
|   A|   1|
|   B|   2|
|   C|   3|
+----+----+

为此，我可以定义一个带有硬编码数字的 UDF，如下所示：

val add100 = udf( (x: Int) => x + 100)
df.withColumn("col3", add100($"col2")).show()

+----+----+----+
|col1|col2|col3|
+----+----+----+
|   A|   1| 101|
|   B|   2| 102|
|   C|   3| 103|
+----+----+----+    

我的问题是，使数字（100 以上）可配置的最佳方法是什么？

我尝试了以下方法，它似乎有效。但我想知道是否有其他更好的方法来实现相同的操作结果？

val addP = udf( (x: Int, p: Int) => x + p )
df.withColumn("col4", addP($"col2", lit(100)))

+----+----+----+
|col1|col2|col4|
+----+----+----+
|   A|   1| 101|
|   B|   2| 102|
|   C|   3| 103|
+----+----+----+

【参考方案1】：

你可以定义一个curried函数，取出额外的参数并返回一个udf，它只接受列作为参数：

val addP = (p: Int) => udf( (x: Int) => x + p ) 
// addP: Int => org.apache.spark.sql.expressions.UserDefinedFunction = <function1>

df.withColumn("col3", addP(100)($"col2")).show
+----+----+----+
|col1|col2|col3|
+----+----+----+
|   A|   1| 101|
|   B|   2| 102|
|   C|   3| 103|
+----+----+----+

【讨论】：

虽然我一直在寻找一个像@mtoto 建议的简单表达式，但我真的很喜欢currying 的想法！谢谢@Psidom！

val addP = (p: Int) => udf( (x: Int) => x + p ) val addP = udf( (x: Int, p: Int)有什么区别=> x + p ) 为什么要求寻找原始 val addP = udf( (x: Int, p: Int) => x + p )

【参考方案2】：

这里不需要 udf：

df.withColumn("col3", df("col2") + 100).show
+----+----+----+
|col1|col2|col3|
+----+----+----+
|   A|   1| 101|
|   B|   2| 102|
|   C|   3| 103|
+----+----+----+

【讨论】：

谢谢！简单又好看！正是我正在寻找的表达方式！