scala - 如何动态地将列添加到DataFrame？

Question

我正在尝试将列从字符串序列动态添加到数据帧。

这是一个示例:源数据帧如下:

+-----+---+----+---+---+
|id | A | B  | C  | D  |
+-----+---+----+---+---+
|1 |toto|tata|titi|    |
|2 |bla |blo |   |     |
|3 |b   | c  |  a |  d |
+-----+---+----+---+---+

我还有一个字符串序列，其中包含我要添加的列的名称。如果源 DataFrame 中已存在列，则它必须执行某种差异，如下所示:

序列看起来像:

val columns = Seq("A", "B", "F", "G", "H")

期望是:

+-----+---+----+---+---+---+---+---+
|id | A | B  | C  | D  | F | G | H |
+-----+---+----+---+---+---+---+---+
|1 |toto|tata|titi|tutu|null|null|null
|2 |bla |blo |   |     |null|null|null|
|3 |b   | c  |  a |  d |null|null|null|
+-----+---+----+---+---+---+---+---+

到目前为止我所做的是这样的:

val difference = columns diff sourceDF.columns
val finalDF = difference.foldLeft(sourceDF)((df, field) => if (!sourceDF.columns.contains(field)) df.withColumn(field, lit(null))) else df)
  .select(columns.head, columns.tail:_*) 

但我不知道如何使用 Spark 以更简单、更容易阅读的方式有效地做到这一点......

提前致谢

Answer 1

这是使用 Seq.diff、单个 select 和 map 生成最终列列表的另一种方法:

import org.apache.spark.sql.functions.{lit, col}


val newCols = Seq("A", "B", "F", "G", "H")

val updatedCols = newCols.diff(df.columns).map{ c => lit(null).as(c)}

val selectExpr = df.columns.map(col) ++ updatedCols

df.select(selectExpr:_*).show

// +---+----+----+----+----+----+----+----+
// | id|   A|   B|   C|   D|   F|   G|   H|
// +---+----+----+----+----+----+----+----+
// |  1|toto|tata|titi|null|null|null|null|
// |  2| bla| blo|null|null|null|null|null|
// |  3|   b|   c|   a|   d|null|null|null|
// +---+----+----+----+----+----+----+----+

首先我们找到 newCols 和 df.columns 之间的差异，这给了我们:F, G, H。接下来我们通过 map 函数将列表中的每个元素转换为 lit(null).as(c) 。最后，我们将现有列表和新列表连接在一起以生成用于 select 的 selectExpr 。