我需要在现有的spark数据框中附加多个列,其中列名在List中给出,假设新列的值是常量,例如给定的输入列和数据帧是
val columnsNames=List("col1","col2")
val data = Seq(("one", 1), ("two", 2), ("three", 3), ("four", 4))
在追加两列之后,假设col1的常量值为“val1”,col2的值为“val2”,输出数据帧应为
+-----+---+-------+------+
| _1| _2|col1 |col2|
+-----+---+-------+------+
| one| 1|val1 |val2|
| two| 2|val1 |val2|
|three| 3|val1 |val2|
| four| 4|val1 |val2|
+-----+---+-------+------+
我写了一个附加列的函数
def appendColumns (cols: List[String], ds: DataFrame): DataFrame = {
cols match {
case Nil => ds
case h :: Nil => appendColumns(Nil, ds.withColumn(h, lit(h)))
case h :: tail => appendColumns(tail, ds.withColumn(h, lit(h)))
}
}
有没有更好的方法和更多的功能方式来做到这一点。
谢谢
是的,有一种更好,更简单的方法。基本上,你可以使用withColumn
进行尽可能多的调用。有了很多列,催化剂,优化火花查询的引擎可能会感到有些不知所措(我曾经有过类似用例的经验)。我甚至看到它在尝试数千列时会在驱动程序上产生OOM。为了避免压力催化剂(并编写更少的代码;-)),你可以简单地使用如下所示的select
来完成一个spark命令:
val data = Seq(("one", 1), ("two", 2), ("three", 3), ("four", 4)).toDF
// let's assume that we have a map that associates column names to their values
val columnMap = Map("col1" -> "val1", "col2" -> "val2")
// Let's create the new columns from the map
val newCols = columnMap.keys.map(k => lit(columnMap(k)) as k)
// selecting the old columns + the new ones
data.select(data.columns.map(col) ++ newCols : _*).show
+-----+---+----+----+
| _1| _2|col1|col2|
+-----+---+----+----+
| one| 1|val1|val2|
| two| 2|val1|val2|
|three| 3|val1|val2|
| four| 4|val1|val2|
+-----+---+----+----+
与递归相反,对于有限数量的列,使用foldLeft的更一般方法我认为更通用。使用Databricks Notebook:
import org.apache.spark.sql._
import org.apache.spark.sql.functions._
import spark.implicits._
val columnNames = Seq("c3","c4")
val df = Seq(("one", 1), ("two", 2), ("three", 3), ("four", 4)).toDF("c1", "c2")
def addCols(df: DataFrame, columns: Seq[String]): DataFrame = {
columns.foldLeft(df)((acc, col) => {
acc.withColumn(col, lit(col)) })
}
val df2 = addCols(df, columnNames)
df2.show(false)
收益:
+-----+---+---+---+
|c1 |c2 |c3 |c4 |
+-----+---+---+---+
|one |1 |c3 |c4 |
|two |2 |c3 |c4 |
|three|3 |c3 |c4 |
|four |4 |c3 |c4 |
+-----+---+---+---+
请注意以下内容:https://medium.com/@manuzhang/the-hidden-cost-of-spark-withcolumn-8ffea517c015虽然情况略有不同,但另一个答案通过选择方法提到了这一点。