在Spark中将多列追加到现有数据框

时间:2019-04-09 06:03:00

标签: scala apache-spark apache-spark-sql bigdata

我需要将多个列追加到现有的spark数据帧中,其中列表中给出了列名 假设新列的值是恒定的,例如给定的输入列和数据框为

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)))

            }
        }

有没有更好的方法和更实用的方法来做到这一点。

谢谢

2 个答案:

答案 0 :(得分:2)

是的,有一种更好,更简单的方法。基本上,您对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|
+-----+---+----+----+

答案 1 :(得分:2)

与递归相反,对于有限数量的列,使用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尽管上下文稍有不同,但其他答案都通过选择方法暗示了这一点。