Question

我有一个包含A列和A列的数据框。 B的String类型。我们假设下面的数据框

+--------+
|A  | B  |
|1a | 1b |
|2a | 2b |

我想添加第三列，创建A＆amp; A的地图。 B栏

+-------------------------+
|A  | B  |  C             |
|1a | 1b | {A->1a, B->1b} |
|2a | 2b | {A->2a, B->2b} |

我试图按照以下方式进行。我有udf，它接收数据帧并返回一个地图

val test = udf((dataFrame: DataFrame) => {
val result = new mutable.HashMap[String, String]
dataFrame.columns.foreach(col => {
  result.put(col, dataFrame(col).asInstanceOf[String])
})
result
})

我正在以下面的方式调用这个udf，因为我试图将DataSet作为文字传递而抛出RunTimeException

df.withColumn("C", Helper.test(lit(df.select(df.columns.head, df.columns.tail: _*)))

我不想将df（'a'）df（'b'）传递给我的帮助器udf，因为我希望它们是我可以选择的列的通用列表。任何指针？

Answer 1

地图方式

您可以使用map 内置功能作为

import org.apache.spark.sql.functions._
val columns = df.columns
df.withColumn("C", map(columns.flatMap(x => Array(lit(x), col(x))): _*)).show(false)

应该给你

+---+---+---------------------+
|A  |B  |C                    |
+---+---+---------------------+
|1a |1b |Map(A -> 1a, B -> 1b)|
|2a |2b |Map(A -> 2a, B -> 2b)|
+---+---+---------------------+

Udf方式

或者您可以将udf定义为

//collecting column names to be used in the udf
val columns = df.columns
//definining udf function
import org.apache.spark.sql.functions._
def createMapUdf = udf((names: Seq[String], values: Seq[String])=> names.zip(values).toMap)
 //calling udf function 
df.withColumn("C", createMapUdf(array(columns.map(x => lit(x)): _*), array(col("A"), col("B")))).show(false)

我希望答案很有帮助

Answer 2

@Ramesh Maharjan - 你的答案已经很棒了，我的答案就是使用字符串插值以动态方式回答你的UDF答案。

专栏D以动态方式提供。

df.withColumn("C", createMapUdf(array(columns.map(x => lit(x)): _*), 
array(col("A"), col("B"))))
.withColumn("D", createMapUdf(array(columns.map(x => lit(x)): _*), 
array(columns.map(x => col(s"$x") ): _* ))).show()

将每行数据框转换为地图

2 个答案: